Gradient Accumulation › Code Implementation
Language AI HandbookAll chapters Gradient Accumulation › Code ImplementationOne section · Sections in this chapter9 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We will build a complete gradient accumulation training loop in PyTorch that demonstrates the correct implementation, including the loss scaling, optimizer synchronization, and gradient clipping interaction.
Setting Up the Training Environment
Let us start by defining the model and data:
import random
import numpy as np
import torch
import torch.nn as nn
# Reproducibility
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)
# Simple transformer-like language model for demonstration
class SmallTransformer(nn.Module):
def __init__(
self, vocab_size=1000, d_model=128, nhead=4, num_layers=2, seq_len=32
):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=256,
dropout=0.0,
batch_first=True,
)
self.transformer = nn.TransformerEncoder(
encoder_layer, num_layers=num_layers
)
self.head = nn.Linear(d_model, vocab_size)
self.d_model = d_model
def forward(self, x):
emb = self.embedding(x) * (self.d_model**0.5)
out = self.transformer(emb)
return self.head(out)
# Synthetic token dataset
VOCAB_SIZE = 1000
SEQ_LEN = 32
DATASET_SIZE = 512
tokens = torch.randint(0, VOCAB_SIZE, (DATASET_SIZE, SEQ_LEN + 1))
inputs = tokens[:, :-1] # shape: (512, 32)
targets = tokens[:, 1:] # shape: (512, 32)
dataset = torch.utils.data.TensorDataset(inputs, targets)Baseline: Standard Training Loop
First, let us implement the standard training loop with a batch size of 32, to establish the baseline gradient behavior. The demonstration model sets dropout to zero so that both paths compute deterministic gradients; the earlier dropout discussion still applies to real training runs.
from torch.utils.data import DataLoader
BATCH_SIZE = 32
MICRO_BATCH = 8
ACCUM_STEPS = BATCH_SIZE // MICRO_BATCH # = 4
model_baseline = SmallTransformer(vocab_size=VOCAB_SIZE)
optimizer_baseline = torch.optim.Adam(model_baseline.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
loader_baseline = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False)
# Run one step to capture gradient norms
model_baseline.train()
optimizer_baseline.zero_grad()
x_full, y_full = next(iter(loader_baseline))
logits = model_baseline(x_full) # (32, 32, 1000)
loss = loss_fn(logits.reshape(-1, VOCAB_SIZE), y_full.reshape(-1))
loss.backward()
# Record gradient norms before optimizer step
grad_norms_baseline = {}
for name, param in model_baseline.named_parameters():
if param.grad is not None:
grad_norms_baseline[name] = param.grad.norm().item()
total_grad_norm_baseline = (
sum(v**2 for v in grad_norms_baseline.values()) ** 0.5
)Baseline batch size: 32 Baseline loss: 7.0769 Total gradient norm (baseline): 1.075834 Number of parameter tensors: 27
The baseline run processes 32 samples simultaneously and computes the mean loss across all of them. The gradient norm reflects the average gradient direction over the full batch.
Gradient Accumulation Loop
Now let us implement the equivalent training step using gradient accumulation with micro-batches of size 8 and accumulation steps:
# Create a fresh model with the same initial weights as baseline
model_accum = SmallTransformer(vocab_size=VOCAB_SIZE)
model_accum.load_state_dict(model_baseline.state_dict())
optimizer_accum = torch.optim.Adam(model_accum.parameters(), lr=1e-3)
loader_micro = DataLoader(dataset, batch_size=MICRO_BATCH, shuffle=False)
micro_iter = iter(loader_micro)
model_accum.train()
optimizer_accum.zero_grad() # Zero ONCE per logical batch
accum_loss = 0.0
for step in range(ACCUM_STEPS):
x_micro, y_micro = next(micro_iter)
logits_micro = model_accum(x_micro) # (8, 32, 1000)
# Compute mean loss over micro-batch, then scale by 1/K
loss_micro = loss_fn(
logits_micro.reshape(-1, VOCAB_SIZE), y_micro.reshape(-1)
)
scaled_loss = loss_micro / ACCUM_STEPS # CRITICAL: divide by K
scaled_loss.backward() # Accumulates into gradient buffers
accum_loss += loss_micro.item()
# Record accumulated gradient norms (before optimizer step)
grad_norms_accum = {}
for name, param in model_accum.named_parameters():
if param.grad is not None:
grad_norms_accum[name] = param.grad.norm().item()
total_grad_norm_accum = sum(v**2 for v in grad_norms_accum.values()) ** 0.5
mean_loss_accum = accum_loss / ACCUM_STEPSAccumulation steps (K): 4 Micro-batch size: 8 Effective batch size: 32 Mean loss (accumulated): 7.0769 Total gradient norm (accumulated): 1.075834
With identical initial weights and the same data ordering, the accumulated gradient norm should match the baseline gradient norm closely. The small numerical differences arise from floating-point rounding order, which varies when summing gradients in a different sequence.
Comparing Gradient Equivalence
Let us verify the gradient equivalence numerically by comparing per-parameter gradient norms:
# Collect per-layer gradient differences
param_names = list(grad_norms_baseline.keys())
grad_diffs = {}
for name in param_names:
if name in grad_norms_accum:
diff = abs(grad_norms_baseline[name] - grad_norms_accum[name])
rel_diff = diff / (grad_norms_baseline[name] + 1e-12)
grad_diffs[name] = {"abs_diff": diff, "rel_diff": rel_diff}
max_rel_diff = max(v["rel_diff"] for v in grad_diffs.values())
mean_rel_diff = sum(v["rel_diff"] for v in grad_diffs.values()) / len(
grad_diffs
)Gradient norm comparison (baseline vs accumulated): Total norm baseline: 1.07583426 Total norm accumulated: 1.07583422 Absolute difference: 4.22e-08 Max relative diff (per layer): 1.16e-07 Mean relative diff: 2.82e-08
The relative differences should be on the order of or smaller, which corresponds to floating-point rounding errors rather than any systematic bias. This confirms the mathematical equivalence between full-batch training and gradient accumulation.
Correct Gradient Clipping Integration
When gradient clipping is part of the training recipe, it must be applied after accumulation but before the optimizer step:
model_clip = SmallTransformer(vocab_size=VOCAB_SIZE)
model_clip.load_state_dict(model_baseline.state_dict())
optimizer_clip = torch.optim.Adam(model_clip.parameters(), lr=1e-3)
MAX_GRAD_NORM = 1.0
loader_clip = DataLoader(dataset, batch_size=MICRO_BATCH, shuffle=False)
clip_iter = iter(loader_clip)
model_clip.train()
optimizer_clip.zero_grad()
for step in range(ACCUM_STEPS):
x_c, y_c = next(clip_iter)
logits_c = model_clip(x_c)
loss_c = loss_fn(logits_c.reshape(-1, VOCAB_SIZE), y_c.reshape(-1))
(loss_c / ACCUM_STEPS).backward()
# Clip AFTER accumulation, BEFORE optimizer step
grad_norm_before_clip = nn.utils.clip_grad_norm_(
model_clip.parameters(), max_norm=MAX_GRAD_NORM
)
optimizer_clip.step() # Now apply the clipped accumulated gradientsGradient norm before clipping: 1.0758 Max gradient norm allowed: 1.0000 Clipping applied: True Gradients were scaled by: 0.9295
The clip_grad_norm_ function returns the gradient norm before clipping, which is useful for logging. If the norm exceeds MAX_GRAD_NORM, it scales all gradients down proportionally. This also is a health monitor: persistently high pre-clip norms can signal the training instabilities discussed in the Training Stability chapter.
Distributed Training with no_sync
In multi-GPU training with PyTorch's DistributedDataParallel, you should suppress gradient all-reduce during the accumulation micro-batches and allow it only on the final step:
# Pattern for DDP + gradient accumulation (not executed: requires DDP setup)
# model must be wrapped: model = DDP(base_model, ...)
# optimizer and loss_fn as before
optimizer.zero_grad()
for step in range(ACCUM_STEPS):
x_micro, y_micro = next(micro_iter)
# Suppress all-reduce for all but the last micro-batch
if step < ACCUM_STEPS - 1:
with model.no_sync():
logits = model(x_micro)
loss = loss_fn(logits.reshape(-1, VOCAB_SIZE), y_micro.reshape(-1))
(loss / ACCUM_STEPS).backward()
else:
# Final micro-batch: allow the all-reduce to fire
logits = model(x_micro)
loss = loss_fn(logits.reshape(-1, VOCAB_SIZE), y_micro.reshape(-1))
(loss / ACCUM_STEPS).backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()Without no_sync(), each of the backward passes would trigger an all-reduce across all GPUs, performing times as many communication rounds as necessary. For a large model with , this would add 15 unnecessary all-reduce operations per optimizer step.
Memory Footprint Comparison
Let us measure the actual memory savings from reducing the micro-batch size:
import gc
def measure_peak_memory(
batch_size, model_fn, vocab_size=VOCAB_SIZE, seq_len=SEQ_LEN
):
"""Measure peak GPU memory (or estimate via tensor sizes) for one forward-backward."""
gc.collect()
torch.cuda.empty_cache() if torch.cuda.is_available() else None
model = model_fn()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
x = torch.randint(0, vocab_size, (batch_size, seq_len))
y = torch.randint(0, vocab_size, (batch_size, seq_len))
# Estimate activation memory: each layer stores O(batch * seq * d_model) floats
d_model = model.d_model
num_layers = 2 # from SmallTransformer definition
bytes_per_float = 4 # float32
# Rough activation estimate: 12 tensors per layer (QKV, attention, FFN intermediates)
activation_bytes = (
batch_size * seq_len * d_model * num_layers * 12 * bytes_per_float
)
return activation_bytes / (1024**2) # MB
batch_sizes = [1, 2, 4, 8, 16, 32]
memory_estimates = []
for bs in batch_sizes:
mem_mb = measure_peak_memory(
bs, lambda: SmallTransformer(vocab_size=VOCAB_SIZE)
)
memory_estimates.append(mem_mb)Estimated activation memory by micro-batch size: Micro-batch Activation (MB) Effective batch (K=32/bs) ---------------------------------------------------------- 1 0.4 32 2 0.8 32 4 1.5 32 8 3.0 32 16 6.0 32 32 12.0 32
The activation memory scales linearly with micro-batch size. A micro-batch of 1 uses 1/32 the activation memory of a micro-batch of 32, while gradient accumulation with recovers the same effective batch size of 32. Parameters, gradients, and optimizer states (not shown) remain constant across all rows.
Key Parameters
The key parameters for gradient accumulation are:
- accumulation_steps (): The number of micro-batches per optimizer step. Set to
effective_batch_size // micro_batch_size. Higher values reduce activation memory but increase the number of forward-backward passes per update. - micro_batch_size: The number of samples processed per forward-backward pass. Limited by GPU memory capacity. Smaller values consume less activation memory.
- loss_scaling: Always divide each micro-batch loss by before calling
.backward(). Failure to do this scales the effective learning rate by . - gradient_clipping: Apply
clip_grad_norm_()after the final backward pass and beforeoptimizer.step(). Never clip inside the accumulation loop. - no_sync(): In distributed training, wrap the first micro-batch backward passes in
model.no_sync()to suppress unnecessary all-reduce communication.