Batch Size and Gradient Accumulation

~10 min · batch, gradient-accumulation

Level 0Token

0 XP0/94 lessons0/10 achievements

0/120 XP to next level120 XP to go0% complete

LLM pretraining uses massive effective batch sizes — millions of tokens per gradient update. Big batches stabilize gradient estimates (more samples = lower noise in the average gradient) and improve hardware utilization. The challenge: GPU memory.

Gradient accumulation

If you cannot fit your desired batch in memory, run several smaller "micro-batches," accumulate gradients across them, and update once. The effective batch size is micro-batch × accumulation steps. Memory cost is determined by the micro-batch; compute cost is identical to a single big batch.

This pattern is so universal that all modern training frameworks — PyTorch, DeepSpeed, Megatron-LM — make it the default. Production runs often use micro-batch ~1-4 sequences per GPU, accumulation ~16-128, distributed across hundreds or thousands of GPUs.

Code

Gradient accumulation in PyTorch·python

optimizer.zero_grad()
for step, micro_batch in enumerate(loader):
    logits = model(micro_batch)
    loss = causal_lm_loss(logits, micro_batch)
    # Scale to keep the effective gradient magnitude correct
    (loss / accumulation_steps).backward()
    if (step + 1) % accumulation_steps == 0:
        # Gradient clipping is standard for LLM training
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

External links

Exercise

Train a small model with effective batch size 256, achieved via three configurations: (a) micro-batch 256, accum 1 (assuming you have memory); (b) micro-batch 32, accum 8; (c) micro-batch 8, accum 32. Compare wall-clock time and final loss. Are the curves identical? They should be, modulo GPU optimization differences.

Progress

Progress is local-only — sign in to sync across devices.

← PreviousLearning Rate: Warmup + Cosine Decay Next →Mixed Precision: BF16, FP16, FP8

Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)

No comments yet — be the first.