How to run gradient accumulation in PyTorch

Large models and high-resolution inputs can exhaust accelerator memory before a preferred batch size fits. PyTorch can retain gradients across smaller microbatches so the optimizer still updates from the combined effective batch.

Each call to loss.backward() adds into parameter .grad buffers. Dividing each mean loss by accumulation_steps keeps the accumulated gradient on the same scale as one loss over the effective batch when every accumulation window contains equally sized microbatches.

A deterministic CPU run uses four microbatches and performs optimizer.step() after every pair. A second model trains on two full batches of four samples, so matching parameters, changed weights, and cleared gradients expose both the update and its accumulation boundary.

1. Create gradient_accumulation_demo.py with deterministic data and paired models.

   gradient_accumulation_demo.py

   import copy
    
   import torch
   from torch import nn
   from torch.utils.data import DataLoader, TensorDataset
    
    
   torch.manual_seed(7)
    
   features = torch.linspace(-1.5, 1.5, steps=24, dtype=torch.float32).reshape(8, 3)
   targets = features @ torch.tensor([[0.7], [-0.2], [0.4]]) + 0.15
    
   microbatch_size = 2
   accumulation_steps = 2
   effective_batch_size = microbatch_size * accumulation_steps
    
   train_loader = DataLoader(
       TensorDataset(features, targets),
       batch_size=microbatch_size,
       shuffle=False,
   )
    
   model = nn.Linear(3, 1)
   reference_model = copy.deepcopy(model)
   initial_parameters = [parameter.detach().clone() for parameter in model.parameters()]
    
   loss_fn = nn.MSELoss()
   optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
   reference_optimizer = torch.optim.SGD(reference_model.parameters(), lr=0.05)

   microbatch_size controls the samples processed by one forward pass. effective_batch_size counts the samples represented by one optimizer update.

2. Append the accumulation-window setup and scaled backward pass to gradient_accumulation_demo.py.

   if len(train_loader) % accumulation_steps != 0:
       raise ValueError("train_loader must contain complete accumulation windows")
    
   optimizer.zero_grad(set_to_none=True)
   backward_passes = 0
   optimizer_steps = 0
    
   for batch_index, (batch_features, batch_targets) in enumerate(train_loader, start=1):
       predictions = model(batch_features)
       loss = loss_fn(predictions, batch_targets) / accumulation_steps
       loss.backward()
       backward_passes += 1

   The loss division averages the gradients across two equal microbatches. A shorter final window needs scaling by its actual sample count rather than the configured full-window count.

3. Insert the optimizer boundary below backward_passes += 1 inside the microbatch loop.

       if batch_index % accumulation_steps == 0:
           optimizer.step()
           optimizer_steps += 1
           optimizer.zero_grad(set_to_none=True)
           gradients_cleared = all(
               parameter.grad is None for parameter in model.parameters()
           )
           print(
               f"optimizer_step={optimizer_steps} "
               f"completed_after_microbatch={batch_index} "
               f"gradients_cleared={gradients_cleared}"
           )

   optimizer.step() and optimizer.zero_grad() occur only after a complete accumulation window. Gradient clipping or mixed-precision scaler updates belong at this same boundary.
   Related: How to clip gradients in PyTorch
   Related: How to enable mixed precision in PyTorch

4. Append the full-batch reference updates and parameter checks after the microbatch loop.

   reference_loader = DataLoader(
       TensorDataset(features, targets),
       batch_size=effective_batch_size,
       shuffle=False,
   )
    
   for batch_features, batch_targets in reference_loader:
       reference_optimizer.zero_grad(set_to_none=True)
       reference_loss = loss_fn(reference_model(batch_features), batch_targets)
       reference_loss.backward()
       reference_optimizer.step()
    
   parameters_changed = any(
       not torch.equal(before, after)
       for before, after in zip(initial_parameters, model.parameters())
   )
   max_parameter_difference = max(
       (parameter - reference_parameter).abs().max().item()
       for parameter, reference_parameter in zip(
           model.parameters(), reference_model.parameters()
       )
   )
   matches_full_batch = max_parameter_difference < 1e-7
    
   print(f"effective_batch_size={effective_batch_size}")
   print(f"backward_passes={backward_passes}")
   print(f"optimizer_steps={optimizer_steps}")
   print(f"parameters_changed={parameters_changed}")
   print(f"max_parameter_difference={max_parameter_difference:.8f}")
   print(f"matches_full_batch={matches_full_batch}")
    
   if not parameters_changed or not matches_full_batch:
       raise SystemExit("gradient accumulation verification failed")

   The reference optimizer receives the same four samples per update in one batch. A nonzero parameter difference or unchanged accumulated model makes the script exit with an error.

5. Run the completed gradient accumulation script.

   $ python3 gradient_accumulation_demo.py
   optimizer_step=1 completed_after_microbatch=2 gradients_cleared=True
   optimizer_step=2 completed_after_microbatch=4 gradients_cleared=True
   effective_batch_size=4
   backward_passes=4
   optimizer_steps=2
   parameters_changed=True
   max_parameter_difference=0.00000000
   matches_full_batch=True