mbp-pytorch 0.2.6

Enabling large batch size training for DNN models beyond the memory limit

Micro-Batch Processing (MBP)

MBP is a simple method to enable deep learning (DL) models to train in large mini-batches, even when the devices (i.e., GPUs) have limited memory size. This GitHub repository provides an implementation of PyTorch-based MBP. You can easily train your DL models in large-batch training without the need to apply complex techniques, increase the number of GPUs, or GPU memory size.

The MBP implementation was designed with reference to the research presented in: "Enabling large batch size training for DNN models beyond the memory limit while maintaining performance" by XinYu Piao, DoangJoo Synn, Jooyoung Park, and Jong-Kook Kim (IEEE Access 2023).

Install

Requirements: Python 3.10 or higher and a CUDA-enabled version of PyTorch are required.

You can easily install the MBP implementation using pip:

pip install mbp-pytorch

If you want to install MBP from source code, please follow the steps below:

git clone https://github.com/HPIC/MBP.git
cd MBP/
pip install .

Usage

import torch
import mbp

device = torch.device("cuda:0")

@mbp.apply(["x", "label"], 16, device) # Condition-1
def train_fn(model, criterion, x, label, *args, **kwargs):
    output = model(x)
    loss = criterion(output, label)
    return loss, output, ... # Condition-2

for image, label in dataloader:
    optimizer.zero_grad()
    loss, *others = train_fn(model, criterion, x=image, label=label) # Condition-3
    optimizer.step()

You can easily apply MBP to your DL training by following three conditions [see exmaple code]:

1. Specify Batch Names, Micro-Batch Size, and Target Device: Use the @mbp.apply decorator to specify the batch names to be split, the desired micro-batch size, and the target device.
2. Return Loss Without Backpropagation: Verify that the decorated function returns the loss without calling .backward(). There is no need to call .backward() because the MBP decorator handles backpropagation automatically. Also, make sure that the loss value is the first item returned by the function decorated with the MBP decorator.
3. Define Batches Explicitly: Explicitly define the corresponding batches in key=value format when using/calling the function decorated with MBP. If not specified, MBP will not be applied to those batches.

⚠️ Caution: The function decorated with MBP decorator returns the loss value for verification purposes only. Do not call .backward() on this loss value, as the MBP decorator will handle the backpropagation automatically. Do not follow the example below:

loss, *others = train_fn(model, criterion, x=image, label=label)
loss.backward() # Don't try this!

MBP automatically detects whether the model is allocated to a CPU or GPU. It then splits a large mini-batch into smaller micro-batches and streams them sequentially to the detected device.

⚠️ Caution: Ensure that the specified device has enough memory to hold the micro-batches and intermediate computations. To apply MBP, ensure that after uploading the model to the GPU memory, there is enough remaining GPU memory capacity to allocate at least one micro-batch size and store intermediate computations computed by each layer of the model.

Arguments

• batch_names (type: List[str] or Tuple[str]): List of batch names to be applied MBP.
• ub_size (type: int): The maximum size of a micro-batch. This value represents the maximum size of a micro-batch that can be executed on each device, and the batch size split by MBP cannot exceed this value. (Default: 1)
• device (type: str | int | Device, optional): Device to which the micro-batch will be sent. (Default: "cpu")
• device_ids (type: List[int] or Tuple[int]): List of GPU IDs for multi-GPU training. When using multiple GPUs with torch.nn.DataParallel, you need to specify the IDs of the GPUs to be used. (Defualt: None)

Multi-GPU

MBP can be applied to both a single GPU and multi-GPUs. One of the key scalability features of MBP is its compatibility with PyTorch's torch.nn.DataParallel() without requiring any code modifications. Simply specify the device_ids as shown in the example below [see exmaple code]:

import torch
import torch.nn as nn
import mbp

device = torch.device("cuda")
device_ids = [0, 1]
model = nn.DataParallel(model, device_ids=device_ids).to(device)

@mbp.apply(["x", "label"], 16, device=device, device_ids=device_ids)
def train_fn(model, criterion, x, label, *args, **kwargs):
    output = model(x)
    loss = criterion(output, label)
    return loss

for image, label in dataloader:
    optimizer.zero_grad()
    loss, *others = train_fn(model, criterion, x=image, label=label)
    optimizer.step()

Citation

If you use this MBP implementation in your research, please cite the following paper:

@article{piao2023enabling,
    title={Enabling large batch size training for DNN models beyond the memory limit while maintaining performance},
    author={Piao, XinYu and Synn, DoangJoo and Park, Jooyoung and Kim, Jong-Kook},
    journal={IEEE Access},
    year={2023},
    publisher={IEEE}
}