gradient accumulation
**Gradient Accumulation**
**What is Gradient Accumulation?**
Accumulate gradients over multiple mini-batches before updating weights, simulating a larger batch size without requiring more memory.
**Why Use It?**
| Constraint | Solution |
|------------|----------|
| GPU memory limits batch size | Accumulate smaller batches |
| Need larger effective batch | More stable gradients |
| Single GPU training | Match multi-GPU batch sizes |
**How It Works**
**Standard Training**
```python
# Each step: forward → backward → update
for batch in dataloader:
loss = model(batch)
loss.backward()
optimizer.step() # Update every batch
optimizer.zero_grad()
```
**With Gradient Accumulation**
```python
accumulation_steps = 4
for i, batch in enumerate(dataloader):
loss = model(batch)
loss = loss / accumulation_steps # Scale loss
loss.backward() # Accumulate gradients
if (i + 1) % accumulation_steps == 0:
optimizer.step() # Update every N batches
optimizer.zero_grad()
```
**Effective Batch Size**
```
effective_batch_size = batch_size × accumulation_steps × num_gpus
Example:
batch_size = 4
accumulation_steps = 8
num_gpus = 1
effective_batch_size = 4 × 8 × 1 = 32
```
**Important Considerations**
**Loss Scaling**
Divide loss by accumulation steps to maintain correct gradient magnitude:
```python
loss = loss / accumulation_steps
```
**Learning Rate**
May need to adjust LR for larger effective batch:
- Linear scaling rule: `lr = base_lr × effective_batch_size / base_batch_size`
- Or use warmup to find optimal LR
**Memory Usage**
| Component | With Accumulation |
|-----------|-------------------|
| Model weights | Same |
| Activations | Per micro-batch |
| Gradients | Accumulate (same size) |
| Optimizer states | Same |
**Batch Normalization**
If using BatchNorm (rare in LLMs), statistics may differ with smaller micro-batches.
**Hugging Face Implementation**
```python
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4, # Micro-batch
gradient_accumulation_steps=8, # Accumulate 8 steps
# Effective: 4 × 8 = 32 per GPU
)
```
**Complete Example**
```python
model.train()
optimizer.zero_grad()
for step, batch in enumerate(dataloader):
outputs = model(**batch)
loss = outputs.loss / gradient_accumulation_steps
loss.backward()
if (step + 1) % gradient_accumulation_steps == 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
print(f"Step {step + 1}: loss = {loss.item() * gradient_accumulation_steps:.4f}")
```