Python / PyTorch Fundamentals Interview Questions
How does gradient accumulation work in PyTorch and when would you use it?
Gradient accumulation simulates a larger effective batch size than fits in GPU memory by summing gradients over several smaller forward/backward passes before calling optimizer.step(). This is useful when training large models on limited GPU memory.
import torch import torch.nn as nn model = nn.Linear(100, 10) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) criterion = nn.CrossEntropyLoss() # Simulate effective batch_size=128 using micro_batch=32 (4 accumulation steps) accumulation_steps = 4 optimizer.zero_grad() for step, (X_micro, y_micro) in enumerate(loader): # loader yields micro-batches logits = model(X_micro) loss = criterion(logits, y_micro) # CRITICAL: scale loss by 1/accumulation_steps before backward # so the accumulated gradient matches what a single large-batch # backward pass would have produced loss = loss / accumulation_steps loss.backward() # gradients ACCUMULATE (not cleared) if (step + 1) % accumulation_steps == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() # update only every N micro-batches optimizer.zero_grad(set_to_none=True) # clear for next accumulation cycle # Effective batch size = micro_batch_size * accumulation_steps # This trades extra forward/backward compute for lower peak memory usage
| Aspect | Effect |
|---|---|
| GPU memory | Stays at micro-batch level — much lower peak usage |
| Wall-clock time | Slightly slower than one large batch (more Python overhead) |
| Effective batch size | micro_batch_size × accumulation_steps |
| BatchNorm caveat | Statistics computed per micro-batch, not the full effective batch — can behave differently than true large-batch training |
Why must the loss be divided by accumulation_steps before calling backward() in gradient accumulation?
What is the main motivation for using gradient accumulation in PyTorch?
More Related questions...
