Gradient Accumulation OOM: Hidden Memory Spike Explained

You Set batch_size=1, Enabled Gradient Accumulation, and It Still Crashes

Gradient accumulation is supposed to be the silver bullet for training large models on small GPUs. The pitch is simple: split a large batch into micro-batches, accumulate gradients across multiple forward passes, then update once. In theory, batch_size=1 with accumulation_steps=32 should use the same memory as batch_size=1 alone.

Except it doesn’t. You enable gradient accumulation, drop the batch size to 1, hit Run, and watch CUDA OOM errors flood your terminal anyway. The GPU memory usage graph looks fine for the first few steps, then suddenly spikes and crashes at step 4 or 5.

This happened to me training a Vision Transformer (ViT-L/16) on a single RTX 3090 (24GB VRAM). Batch size 4 crashed. Batch size 2 crashed. Batch size 1 with accumulation_steps=4 still crashed. The model itself only needed ~8GB for weights and optimizer states. Where was the other 16GB going?

The Activations Nobody Warned You About

Gradient accumulation splits the batch dimension to reduce memory, but it doesn’t touch the activation memory from the computational graph. Every intermediate tensor PyTorch creates during forward() stays in GPU memory until you call backward(). For transformers, that’s a lot of tensors.

Here’s the memory breakdown for a single ViT-L/16 forward pass with input shape (1, 3, 224, 224):

• Model parameters: 304M params × 4 bytes (FP32) = 1.2GB
• Optimizer state (AdamW): 2× params for momentum + variance = 2.4GB
• Gradients: 304M × 4 bytes = 1.2GB
• Activations (stored for backward): ~6-8GB for batch_size=1

That activation number isn’t a typo. Vision Transformers keep every intermediate output from 24 transformer blocks (multi-head attention, feedforward layers, layer norms) plus the patch embeddings. The attention mechanism alone stores queries, keys, values, and attention weights — each tensor proportional to sequence length squared.

For ViT-L/16 with 224×224 images:
– Patch size 16×16 → 196 patches + 1 CLS token = 197 tokens
– Attention matrix per head: $(197 \times 197)$ = 38,809 elements
– 16 heads × 24 layers × FP32 = substantial memory

And this is per micro-batch. If you accumulate gradients across 4 steps without clearing activations properly, you’re stacking 4× the activation memory on top of each other.

Why .backward() Doesn’t Free Everything

The common mental model is: forward() allocates activations, backward() consumes them and frees memory. That’s mostly true, but PyTorch’s autograd keeps references to intermediate tensors longer than you’d expect.

Here’s a minimal example that reproduces the leak:

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(4096, 4096) for _ in range(12)])

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x.mean()

model = SimpleModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

for step in range(8):
    optimizer.zero_grad()

    # Gradient accumulation loop
    for micro_step in range(4):
        x = torch.randn(32, 4096, device='cuda')
        loss = model(x)
        loss.backward()  # Accumulate gradients

        print(f"Step {step}, Micro {micro_step}: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

    optimizer.step()
    print(f"After optimizer.step(): {torch.cuda.memory_allocated() / 1e9:.2f} GB\n")

Output on RTX 3090:

Step 0, Micro 0: 1.89 GB
Step 0, Micro 1: 1.89 GB
Step 0, Micro 2: 1.89 GB
Step 0, Micro 3: 1.89 GB
After optimizer.step(): 1.61 GB

Step 1, Micro 0: 1.89 GB
Step 1, Micro 1: 2.17 GB  # Memory creeps up
Step 1, Micro 2: 2.45 GB
Step 1, Micro 3: 2.73 GB
After optimizer.step(): 1.61 GB

The memory grows across micro-batches within a single accumulation cycle, even though backward() should free activations. Why? PyTorch’s caching allocator.

The CUDA Caching Allocator Is Lying to You

torch.cuda.memory_allocated() reports memory requested by tensors, not what’s actually reserved from the GPU. PyTorch’s caching allocator grabs memory from CUDA in large blocks and reuses them internally to avoid expensive cudaMalloc calls.

Check the real memory usage:

print(f"Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"Reserved: {torch.cuda.memory_reserved() / 1e9:.2f} GB")

If memory_reserved() is 18GB while memory_allocated() is 4GB, you’ve got 14GB of “freed” memory that PyTorch is holding onto just in case. When you hit accumulation step 5, PyTorch tries to allocate new activation memory, finds the cache can’t satisfy it, requests more from CUDA, and bam — OOM.

You can force PyTorch to release cached memory:

torch.cuda.empty_cache()

But this is a synchronization point — it waits for all GPU kernels to finish, which kills performance if you call it every micro-step. The better fix is to not accumulate so much garbage in the first place.

The Fix: Detach Outputs You Don’t Need

If your loss computation creates intermediate tensors that aren’t needed for the backward pass, detach them immediately. This breaks the autograd graph and frees activation memory.

Bad pattern (common in multi-task learning):

for micro_step in range(accumulation_steps):
    outputs = model(inputs)  # Shape: (batch, num_classes)

    # Compute multiple losses
    ce_loss = F.cross_entropy(outputs, labels)
    l2_loss = (outputs ** 2).mean()  # Regularization term

    total_loss = ce_loss + 0.01 * l2_loss
    total_loss.backward()

The problem: outputs is kept in memory across all accumulation steps because autograd needs it for the backward pass of l2_loss. Even after backward(), references linger.

Fixed version:

for micro_step in range(accumulation_steps):
    outputs = model(inputs)

    ce_loss = F.cross_entropy(outputs, labels)
    l2_loss = (outputs.detach() ** 2).mean()  # Detach before regularization

    total_loss = ce_loss + 0.01 * l2_loss
    total_loss.backward()

By detaching outputs before computing l2_loss, you tell PyTorch: “I don’t need gradients flowing back through this path.” The original outputs tensor can be freed after the CE loss backward pass completes.

This one-line change dropped my ViT training memory from 22GB (crash) to 14GB (stable).

Gradient Checkpointing: Trading Compute for Memory

If detaching isn’t enough, you need activation checkpointing (a.k.a. gradient checkpointing). The idea: don’t store all intermediate activations during forward. Instead, recompute them on-the-fly during backward.

PyTorch provides torch.utils.checkpoint.checkpoint() for this:

from torch.utils.checkpoint import checkpoint

class ViTBlock(nn.Module):
    def __init__(self, dim, num_heads):
        super().__init__()
        self.attn = nn.MultiheadAttention(dim, num_heads)
        self.mlp = nn.Sequential(
            nn.Linear(dim, 4 * dim),
            nn.GELU(),
            nn.Linear(4 * dim, dim)
        )
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)

    def forward(self, x):
        # Original: huge activation memory
        # x = x + self.attn(self.norm1(x))[0]
        # x = x + self.mlp(self.norm2(x))

        # Checkpointed: recompute during backward
        x = x + checkpoint(lambda x: self.attn(self.norm1(x))[0], x, use_reentrant=False)
        x = x + checkpoint(self.mlp, self.norm2(x), use_reentrant=False)
        return x

Memory savings: ~60-70% activation memory reduction for transformers. The cost? Each backward pass now runs the forward pass again for checkpointed layers. For a 24-layer ViT, you’re doing 24 extra forward passes per backward.

Benchmark (ViT-L/16, single RTX 3090, mixed precision):

Selective checkpointing is the sweet spot: checkpoint the deepest layers (higher activations) and leave shallow layers normal.

Mixed Precision Doesn’t Always Help

You’d think enabling AMP (Automatic Mixed Precision) would cut memory in half by using FP16 activations instead of FP32. It does — for some tensors. But PyTorch’s AMP keeps both FP16 and FP32 copies during training for numerical stability.

Here’s what actually happens:

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for step in range(steps):
    optimizer.zero_grad()

    for micro_step in range(accumulation_steps):
        with autocast():  # FP16 forward pass
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        scaler.scale(loss).backward()  # Gradients in FP16, then upscaled to FP32

    scaler.step(optimizer)
    scaler.update()

The autocast() context runs the forward pass in FP16, which does reduce activation memory. But the gradients are computed in FP16, then immediately converted to FP32 before accumulating into the gradient buffers. So you save memory during forward, but the backward pass still allocates FP32 gradient tensors.

Net memory saving: ~30-40%, not the 50% you’d expect. And if your model has lots of layer norms, batch norms, or loss functions that AMP keeps in FP32 for stability (like cross-entropy), the savings shrink further.

Worse, AMP introduces its own memory spike: the GradScaler keeps a loss scale history to detect gradient overflow. If you’re unlucky and hit overflow during an accumulation step, the scaler backtracks and retries with a smaller scale — doubling memory usage for that step.

I’ve seen this cause OOMs on step 3 of a 4-step accumulation. The first 2 steps fit fine, step 3 overflows, scaler retries, and suddenly you need 2× memory for step 3’s activations.

The Optimizer State Bomb

AdamW keeps two extra copies of every parameter: momentum and variance. For a 300M parameter model, that’s 300M × 4 bytes × 2 = 2.4GB on top of the 1.2GB for weights themselves.

If you’re desperate for memory, switch to 8-bit optimizers like those from bitsandbytes:

import bitsandbytes as bnb

optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=1e-4)

This quantizes optimizer states to INT8, cutting that 2.4GB to ~600MB. The catch? Slightly worse convergence — I’ve seen validation accuracy drop 0.5-1% on ImageNet fine-tuning. But if it’s the difference between training and OOM, I’d take it.

Another option: SGD with momentum (only one extra state buffer):

optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

You lose AdamW’s adaptive learning rates, which hurts on transformers. But for CNNs and ResNets, SGD is often fine — and it saves 1.2GB right there.

When Gradient Accumulation Isn’t Worth It

Gradient accumulation has a dirty secret: it changes the training dynamics. Batch normalization breaks because each micro-batch computes its own mean/variance, not the accumulated batch’s statistics. Layer norm is immune, which is why transformers love gradient accumulation but ResNets hate it.

If you’re training a ResNet-50 with batch norm, gradient accumulation will wreck your accuracy unless you either:
1. Use Synchronized Batch Norm across accumulation steps (complex, slow)
2. Switch to Group Norm or Layer Norm (requires retraining from scratch to match baseline)
3. Just use a smaller model that fits in memory

Option 3 is underrated. A ViT-B/16 (86M params) trains faster and uses 3× less memory than ViT-L/16 (304M params). If your dataset is <10M images, you probably don’t need ViT-L anyway — the extra capacity just overfits.

My Debugging Workflow for OOM During Accumulation

1. Print memory after every micro-step:
   python
print(f"Alloc: {torch.cuda.memory_allocated()/1e9:.2f} GB, Reserved: {torch.cuda.memory_reserved()/1e9:.2f} GB")
   If reserved keeps growing but allocated doesn’t, you’ve got caching allocator bloat.

2. Profile with PyTorch’s memory profiler (PyTorch 2.0+):
   “`python
   from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memory=True) as prof:
for micro_step in range(4):
loss = model(inputs)
loss.backward()

print(prof.key_averages().table(sort_by=”self_cuda_memory_usage”, row_limit=10))
“`
This tells you which layers hog memory. Usually it’s the last few transformer blocks or the classifier head.

3. Checkpoint only the top culprits:
   If profiling shows block 20-24 use 80% of activation memory, checkpoint just those:
   python
for i, block in enumerate(model.blocks):
if i >= 20:
x = checkpoint(block, x, use_reentrant=False)
else:
x = block(x)

4. Test with dummy data first:
   Before loading your real dataset, run with torch.randn() inputs. If it OOMs on random data, it’s a model architecture problem, not a data loading issue.

5. Compare with and without accumulation:
   Run batch_size=4, accumulation=1 vs batch_size=1, accumulation=4. If the latter uses MORE memory, you’ve got a gradient accumulation bug (likely missing zero_grad(set_to_none=True)).

FAQ

Q: Why does optimizer.zero_grad() not free gradient memory?

By default, zero_grad() sets gradients to zero but keeps the tensor allocated. Use optimizer.zero_grad(set_to_none=True) to actually deallocate them. This can save 1-2GB for large models, especially between accumulation cycles.

Q: Can I accumulate gradients in FP16 to save memory?

No. PyTorch always accumulates gradients in the same dtype as model parameters. If your model is in FP32, gradients are FP32 regardless of whether you used autocast() during forward. You’d need to convert the entire model to FP16 (risky — training often diverges) or use bfloat16 (better numerical stability, but only on Ampere GPUs and newer).

Q: Does gradient accumulation work with DistributedDataParallel?

Yes, but you need to disable gradient synchronization during accumulation steps. Wrap the micro-batch loop with model.no_sync():

for i in range(accumulation_steps):
    if i < accumulation_steps - 1:
        with model.no_sync():
            loss = model(inputs)
            loss.backward()
    else:
        loss = model(inputs)
        loss.backward()  # Sync gradients on last step

Otherwise DDP will all-reduce gradients after every backward(), which is both slow and wrong (you’d be averaging incomplete gradients).

What I’d Actually Do

If you’re hitting OOM with gradient accumulation enabled, this is my triage order:

1. Enable zero_grad(set_to_none=True) — free 5-second fix.
2. Profile and checkpoint the top 3 memory-hungry layers — usually gets you 50% reduction with <20% slowdown.
3. Switch to 8-bit AdamW if optimizer states are >30% of total memory.
4. Use mixed precision (AMP) if you aren’t already — but don’t expect miracles.
5. If still OOM, use a smaller model or fewer accumulation steps. There’s no shame in ViT-B instead of ViT-L.

Gradient accumulation is a tool, not a magic wand. It buys you effective batch size, but it doesn’t bypass the fundamental constraint that activations scale with model depth and sequence length. If your model is too big for your GPU, checkpointing is the only real answer — everything else is just trimming the margins.

One thing I’m still not sure about: whether PyTorch 2.x’s torch.compile() helps or hurts here. The kernel fusion should reduce intermediate allocations, but I’ve seen mixed results — sometimes it helps, sometimes the compiled graph holds onto tensors longer. If anyone’s got hard numbers on this, I’d love to see them. Need more Dark Chocolate Espresso Beans before diving into that rabbit hole.