How to run gradient accumulation in PyTorch

Gradient accumulation in PyTorch lets a training loop build one larger effective batch from several smaller microbatches. It is useful when the model or input size fits only a small batch in device memory, but the optimizer should update after more samples.

Autograd adds each loss.backward() result into parameter .grad buffers until the optimizer clears them. Dividing the loss by the number of accumulation steps keeps the accumulated gradient on the same scale as one full-batch loss mean.

The smoke run uses CPU tensors, a tiny linear model, and two microbatches per optimizer update. It prints each backward pass, clears gradients with optimizer.zero_grad(set_to_none=True) after every step, and compares the final parameters with a full-batch reference update.

1. Create a gradient accumulation smoke script.

   gradient_accumulation_demo.py

   import copy
    
   import torch
   from torch import nn
    
    
   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
    
   model = nn.Sequential(nn.Linear(3, 1))
   reference_model = copy.deepcopy(model)
    
   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 = 2
   accumulation_steps = 2
   effective_batch_size = microbatch_size * accumulation_steps
    
   optimizer.zero_grad(set_to_none=True)
    
   microbatch_count = 0
   optimizer_steps = 0
    
   for start in range(0, len(features), microbatch_size):
       batch_features = features[start : start + microbatch_size]
       batch_targets = targets[start : start + microbatch_size]
    
       raw_loss = loss_fn(model(batch_features), batch_targets)
       scaled_loss = raw_loss / accumulation_steps
       scaled_loss.backward()
    
       microbatch_count += 1
       grad_norm = torch.sqrt(
           sum(
               parameter.grad.detach().pow(2).sum()
               for parameter in model.parameters()
               if parameter.grad is not None
           )
       ).item()
    
       print(
           f"microbatch {microbatch_count}: "
           f"raw_loss={raw_loss.item():.4f}, "
           f"scaled_loss={scaled_loss.item():.4f}, "
           f"grad_norm={grad_norm:.4f}"
       )
    
       if microbatch_count % accumulation_steps == 0:
           optimizer.step()
           optimizer_steps += 1
           optimizer.zero_grad(set_to_none=True)
           gradients_clear = all(parameter.grad is None for parameter in model.parameters())
           print(
               f"optimizer step {optimizer_steps}: "
               f"after {accumulation_steps} backward passes, "
               f"gradients_clear={gradients_clear}"
           )
    
   for start in range(0, len(features), effective_batch_size):
       batch_features = features[start : start + effective_batch_size]
       batch_targets = targets[start : start + effective_batch_size]
    
       reference_optimizer.zero_grad(set_to_none=True)
       reference_loss = loss_fn(reference_model(batch_features), batch_targets)
       reference_loss.backward()
       reference_optimizer.step()
    
   max_delta = max(
       (parameter - reference_parameter).abs().max().item()
       for parameter, reference_parameter in zip(model.parameters(), reference_model.parameters())
   )
    
   print(f"optimizer steps taken: {optimizer_steps}")
   print(f"full-batch max parameter difference: {max_delta:.8f}")

   microbatch_size is the device-sized batch. effective_batch_size is the number of samples represented by one optimizer update.

2. Run the smoke script.

   $ python gradient_accumulation_demo.py
   microbatch 1: raw_loss=1.7059, scaled_loss=0.8530, grad_norm=2.9950
   microbatch 2: raw_loss=0.2798, scaled_loss=0.1399, grad_norm=3.5828
   optimizer step 1: after 2 backward passes, gradients_clear=True
   microbatch 3: raw_loss=0.1112, scaled_loss=0.0556, grad_norm=0.3952
   microbatch 4: raw_loss=0.8146, scaled_loss=0.4073, grad_norm=2.4346
   optimizer step 2: after 2 backward passes, gradients_clear=True
   optimizer steps taken: 2
   full-batch max parameter difference: 0.00000000

3. Confirm the optimizer step boundary.

   The output shows four backward passes but only two optimizer steps. The full-batch max parameter difference line confirms that scaling each microbatch loss by accumulation_steps matched the reference full-batch updates.

4. Use the same boundary in the training loop.

   optimizer.zero_grad(set_to_none=True)
    
   for step, (batch_features, batch_targets) in enumerate(loader, start=1):
       predictions = model(batch_features)
       loss = loss_fn(predictions, batch_targets) / accumulation_steps
       loss.backward()
    
       if step % accumulation_steps == 0:
           optimizer.step()
           optimizer.zero_grad(set_to_none=True)

   This loop assumes the loader supplies a whole number of accumulation groups. Use drop_last=True or handle the final partial group separately when the dataset length does not divide evenly.

5. Keep gradient modifiers at the accumulation boundary.

   When using automatic mixed precision, call GradScaler.step() and GradScaler.update() only after all microbatches for the effective batch have called backward(). If clipping is needed, unscale or clip after the final microbatch and before the optimizer step.
   Related: How to clip gradients in PyTorch

6. Remove the temporary smoke script.

   $ rm gradient_accumulation_demo.py