Micro-batching & Gradient Accumulation

Gradient accumulation lets you train with effective batch sizes far larger than what fits in GPU memory. Process small micro-batches, accumulate gradients, then update once — same mathematical result as a large batch.

1. The Problem

Large batch sizes improve training throughput and stability, but activation memory scales linearly with batch size:

Batch size	Activation memory (70B, s=4K, ckpt)	Fits on 80 GB with TP=8?
1	~1.3 GB	✅
4	~5.2 GB	✅
16	~20.8 GB	⚠️ Tight
64	~83.2 GB	❌ OOM

But you want effective batch size of 1024+ tokens × millions of sequences for good convergence.

2. How Gradient Accumulation Works

accumulation_steps = 16  # effective batch = micro_batch × accumulation_steps
optimizer.zero_grad()

for step, micro_batch in enumerate(dataloader):
    # Forward + backward with micro-batch
    with torch.autocast('cuda', dtype=torch.bfloat16):
        loss = model(micro_batch) / accumulation_steps  # normalize loss
    loss.backward()  # gradients accumulate in .grad buffers
    
    if (step + 1) % accumulation_steps == 0:
        # Only update after K micro-batches
        optimizer.step()
        optimizer.zero_grad()

Key: loss / accumulation_steps normalizes the gradients so the accumulated result equals the gradient from one large batch.

Mathematical Equivalence

Standard large-batch gradient:

g = \frac{1}{B} \sum_{i=1}^{B} \nabla L(x_i)

Accumulated gradient over $K$ micro-batches of size $b_\mu$ (where $B = K \times b_\mu$ ):

g = \frac{1}{K} \sum_{k=1}^{K} \frac{1}{b_\mu} \sum_{i \in \text{batch}_k} \nabla L(x_i) = \frac{1}{B} \sum_{i=1}^{B} \nabla L(x_i)

Identical. The optimizer sees the same gradient either way.

3. Interaction with Distributed Training

With DP/FSDP

DP AllReduce should happen only after all accumulation steps, not after each micro-batch:

# With FSDP: use no_sync context manager
for step, micro_batch in enumerate(dataloader):
    # Skip gradient sync for accumulation steps
    context = model.no_sync() if (step + 1) % accum_steps != 0 else nullcontext()
    
    with context:
        loss = model(micro_batch) / accum_steps
        loss.backward()
    
    if (step + 1) % accum_steps == 0:
        # AllReduce happens here (on the non-no_sync step)
        optimizer.step()
        optimizer.zero_grad()

Without no_sync(), gradients are AllReduced every micro-batch — wasting $K\times$ communication.

With Pipeline Parallelism

PP naturally uses micro-batches — gradient accumulation is built into the schedule:

PP with 16 micro-batches:
  Forward μ1 → Forward μ2 → ... → Forward μ16
  Backward μ16 → ... → Backward μ1
  [Gradients accumulated across all 16 micro-batches]
  Optimizer step

The number of micro-batches in PP determines both the pipeline efficiency (less bubble) and the effective batch size.

4. Practical Considerations

Factor	Impact
Peak memory	Same as single micro-batch (gradients accumulate in-place)
Communication	AllReduce once per $K$ steps instead of every step
Speed	$K$ forward+backward passes per optimizer step (slower per step, but larger effective batch)
Numerical precision	BF16 gradient accumulation may lose precision at large $K$ — use FP32 accumulation if $K > 32$

For pipeline schedules that use micro-batching, see Pipeline Schedules.

Self-Check

Explain why loss / accumulation_steps is needed for mathematical equivalence
Implement gradient accumulation with model.no_sync() for FSDP
Calculate effective batch size from micro-batch size, accumulation steps, and DP degree
Explain when FP32 gradient accumulation is needed instead of BF16

References

PyTorch Gradient Accumulation