gradient accumulation

**Gradient Accumulation and Micro-Batching** is **a training technique that simulates large effective batch sizes by accumulating gradients across multiple small forward/backward passes before optimizer step — enabling training with batch sizes beyond GPU memory through gradient summation while maintaining the convergence properties of large-batch training**. **Core Mechanism:** - **Accumulation Process**: computing loss and gradients on small batch (e.g., 32 examples), accumulating gradients without optimizer step; repeating N times; then stepping optimizer on accumulated gradients - **Effective Batch Size**: accumulation_steps × per_gpu_batch_size = effective batch size (e.g., 4 × 32 = 128 effective) - **Gradient Summation**: ∇L_total = Σᵢ₌₁^N ∇L_i where each ∇L_i from small batch — equivalent to single large batch update - **Memory Savings**: enabling same model with micro_batch_size=32 instead of batch_size=128 — 4x memory reduction (KV cache + activations) **Gradient Accumulation Workflow:** - **Step 1 - Forward**: compute output for first micro-batch (32 examples) with gradient computation enabled - **Step 2 - Backward**: compute gradients for first micro-batch, accumulate in optimizer buffer (don't zero or step) - **Step 3 - Repeat**: repeat forward/backward for N-1 remaining micro-batches (gradient buffer grows) - **Step 4 - Optimizer Step**: single optimizer step using accumulated gradients; zero gradient buffer for next accumulation cycle - **Time Cost**: N forward/backward passes (same compute as single large batch) plus 1 optimizer step (negligible vs forward/backward) **Memory Efficiency Analysis:** - **Activation Memory**: forward pass stores activations for backward; micro-batching reduces peak activation storage by 1/N - **KV Cache**: autoregressive generation stores cache for all tokens; gradient accumulation doesn't reduce this (cache still computed N times) - **Optimizer State**: Adam maintains velocity/second moment buffers; same size as model weights, independent of batch size - **Peak Memory**: reduced from batch_size×feature_dim to (batch_size/N)×feature_dim enabling 4-8x larger models **Practical Training Configurations:** - **Standard Setup**: per_gpu_batch=32, accumulation_steps=4, effective_batch=128 with 1-GPU VRAM (80GB A100) - **Large Model Training**: 70B parameter model requires 140GB memory for weights; effective batch 32 achievable through 8×4 accumulation - **Distributed Setup**: gradient accumulation combined with data parallelism: N_GPUs × per_gpu_batch × accumulation_steps = effective batch - **FSDP/DDP**: fully sharded data parallel stores model partitions; gradient accumulation reduces per-partition batch size requirement **Convergence and Optimization Properties:** - **Noise Scaling**: gradient variance scales as 1/effective_batch_size — larger effective batches produce smoother gradient updates - **Convergence Behavior**: with large effective batch, convergence curve smoother, fewer oscillations — matches large-batch training - **Noise Schedule**: early training (high noise) benefits from larger batches; late training (fine-tuning) uses smaller batches effectively - **Learning Rate Scaling**: with larger effective batch size, enabling proportionally larger learning rates (linear scaling hypothesis) **Practical Trade-offs:** - **Correctness**: mathematically equivalent to single large batch (same gradient computation, same optimizer step) - **Temporal Coupling**: gradients from step i and step j are temporally coupled (computed at different times) — potential issue for some optimizers - **Staleness**: if using momentum, older micro-batch gradients mixed with newer ones — typically negligible impact (<0.5% performance) - **Synchronization**: distributed accumulation requires careful synchronization across GPUs/nodes — synchronous training required **Implementation Details:** - **PyTorch Training Loop**: ``` for step, (input, target) in enumerate(dataloader): output = model(input) loss = criterion(output, target) / accumulation_steps loss.backward() if (step + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad() ``` - **Loss Scaling**: dividing loss by accumulation_steps enables consistent learning rates across different accumulation configurations - **Gradient Clipping**: applied after accumulation (before optimizer step) to cumulative gradients — critical for stability **Distributed Training Considerations:** - **Synchronous AllGather**: in distributed setting, gradients from all devices must be accumulated before stepping — requires synchronization barrier - **Communication Overhead**: gradient communication happens once per accumulation cycle (not per micro-batch) — reduces communication 4-8x - **Load Balancing**: micro-batches should be evenly distributed across GPUs; skewed distribution causes waiting idle time - **Checkpointing**: checkpointing every N optimizer steps (not micro-batch steps); critical for resuming large-scale training **Interaction with Other Techniques:** - **Mixed Precision Training**: gradient scaling and accumulation work together; loss scaling enables FP16 gradient computation - **Learning Rate Schedules**: warmup and cosine decay applied to optimizer steps (not micro-batch steps) — unchanged semantics - **Gradient Clipping**: clipping applied to accumulated gradients (sum from all micro-batches) — clipping threshold may need adjustment - **Weight Decay**: applied per optimizer step; accumulated with weight updates — equivalent to single large batch **Batch Size and Learning Rate Relationships:** - **Linear Scaling Rule**: learning_rate ∝ effective_batch_size enables stable training across batch configurations - **Gradient Noise Scale**: noise variance ∝ 1/effective_batch — important for generalization; larger batches may overfit more - **Batch Size Sweet Spot**: optimal batch size 32-512 for LLM training; beyond 512 marginal returns diminish - **Fine-tuning**: smaller effective batches (32-64) often better for downstream tasks; larger batches (256-512) better for pre-training **Real-World Examples:** - **BERT Training**: effective batch size 256-512 achieved with per-GPU batch 32-64 and accumulation on single GPU - **GPT-3 Training**: batch size 3.2M tokens simulated through gradient accumulation across 1000+ GPUs; enables optimal convergence - **Llama 2 Training**: effective batch 4M tokens using per-GPU batch 16M words with accumulation and pipeline parallelism - **Fine-tuning on Limited VRAM**: 24GB GPU with model-parallel batch 4, accumulation 8 achieves effective batch 32 **Limitations and When Not to Use:** - **Numerical Issues**: extremely small per-batch sizes (batch=1-2) with accumulation can accumulate numerical errors - **Batch Norm Incompatibility**: batch normalization statistics computed per micro-batch (not effective batch) — accuracy degradation possible - **Communication Overhead**: in communication-bound settings, accumulation reduces benefits (bandwidth not the bottleneck) - **Debugging Difficulty**: gradients from multiple steps mixed; harder to debug gradient flow issues **Gradient Accumulation and Micro-Batching are essential training techniques — enabling simulation of large batch sizes on limited hardware through careful gradient accumulation while maintaining convergence properties of large-batch optimization.**

Go deeper with CFSGPT

Get AI-powered deep-dives, save terms, and run advanced simulations — free account.

Create Free Account