- [Define Your Configuration](../basics/define_your_config.md)
- [Training Booster](../basics/booster_api.md)
## Introduction
Gradient accumulation is a common way to enlarge your batch size for training. When training large-scale models, memory can easily become the bottleneck and the batch size can be very small, (e.g. 2), leading to unsatisfactory convergence. Gradient accumulation works by adding up the gradients calculated in multiple iterations, and only update the parameters in the preset iteration.
## Usage
It is simple to use gradient accumulation in Colossal-AI. Just call `booster.no_sync()` which returns a context manager. It accumulate gradients without synchronization, meanwhile you should not update the weights.
## Hands-on Practice
We now demonstrate gradient accumulation. In this example, we let the gradient accumulation size to be 4.
### Step 1. Import libraries in train.py
Create a `train.py` and import the necessary dependencies. The version of `torch` should not be lower than 1.8.1.
```python
import os
from pathlib import Path
import torch
from torchvision import transforms
from torchvision.datasets import CIFAR10
from torchvision.models import resnet18
from torch.utils.data import DataLoader
import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import TorchDDPPlugin
from colossalai.logging import get_dist_logger
from colossalai.cluster.dist_coordinator import priority_execution
```
### Step 2. Initialize Distributed Environment
We then need to initialize distributed environment. For demo purpose, we uses `launch_from_torch`. You can refer to [Launch Colossal-AI](../basics/launch_colossalai.md) for other initialization methods.
```python
# initialize distributed setting
parser = colossalai.get_default_parser()
args = parser.parse_args()
# launch from torch
colossalai.launch_from_torch(config=dict())
```
### Step 3. Create training components
Build your model, optimizer, loss function, lr scheduler and dataloaders. Note that the root path of the dataset is obtained from the environment variable `DATA`. You may `export DATA=/path/to/data` or change `Path(os.environ['DATA'])` to a path on your machine. Data will be automatically downloaded to the root path.
Use booster in a normal training loops, and verify gradient accumulation. `param_by_iter` is to record the distributed training information.
```python
optimizer.zero_grad()
for idx, (img, label) in enumerate(train_dataloader):
sync_context = booster.no_sync(model)
img = img.cuda()
label = label.cuda()
if idx % (GRADIENT_ACCUMULATION - 1) != 0:
with sync_context:
output = model(img)
train_loss = criterion(output, label)
booster.backward(train_loss, optimizer)
else:
output = model(img)
train_loss = criterion(output, label)
booster.backward(train_loss, optimizer)
optimizer.step()
optimizer.zero_grad()
ele_1st = next(model.parameters()).flatten()[0]
param_by_iter.append(str(ele_1st.item()))
if idx != 0 and idx % (GRADIENT_ACCUMULATION - 1) == 0:
break
for iteration, val in enumerate(param_by_iter):
print(f'iteration {iteration} - value: {val}')
if param_by_iter[-1] != param_by_iter[0]:
print('The parameter is only updated in the last iteration')
```
### Step 6. Invoke Training Scripts
To verify gradient accumulation, we can just check the change of parameter values. When gradient accumulation is set, parameters are only updated in the last step. You can run the script using this command:
