Gradient Accumulation Loss Bug: Why accum=8 Isn't Batch Size 8

You drop per_device_train_batch_size from 8 to 1 and set gradient_accumulation_steps=8 to fit the model on one GPU. Same math, less memory — that is the whole selling point. Then the fine-tune comes out measurably worse than the large-batch run, and the loss curve sits visibly above it from step one.

That gap is the gradient accumulation loss bug: the standard accumulation loop computes a mean of per-micro-batch means, not the mean over all tokens in the optimizer step. When your micro-batches contain different numbers of supervised tokens — which they always do in instruction tuning — the two are not equal, and short examples end up dominating the gradient.

TL;DR

• loss.backward() per micro-batch with reduction="mean", then dividing by G, gives each micro-batch weight 1/G instead of giving each token weight 1/N. That is only correct when every micro-batch has the same unmasked-token count.
• Each token in micro-batch g gets its gradient scaled by n̄ / n_g (mean token count over that micro-batch's count). A 20-token answer next to an 800-token answer gets ~40x the per-token pull it deserves.
• DDP repeats the same error one level up: all_reduce averages gradients across ranks, so uneven token counts between GPUs re-skew the batch.
• Fix: compute num_items_in_batch = total unmasked label tokens across all micro-batches and all ranks, use reduction="sum", and divide by that one number (times world_size to undo DDP's mean).
• Transformers ships the fix since v4.46 via num_items_in_batch plus average_tokens_across_devices. A custom compute_loss override that omits the kwarg silently reinstates the bug.

What is the gradient accumulation loss bug?

It is the mismatch between mean-of-means and the true token mean. The objective you think you are minimizing over an optimizer step is

L = (1 / N) * Σ_g Σ_t  ℓ(g, t)        N = Σ_g n_g

where n_g is the number of non-masked label tokens in micro-batch g and ℓ(g,t) is the per-token cross-entropy. The objective the naive loop actually minimizes is

L̂ = (1 / G) * Σ_g [ (1 / n_g) * Σ_t ℓ(g, t) ]

Pull out the weight on a single token in micro-batch g:

• true weight: 1 / N
• naive weight: 1 / (G · n_g)

Their ratio is N / (G · n_g) = n̄ / n_g. The gradient contribution of every token is inversely proportional to how many other supervised tokens happened to share its micro-batch. That is a data-ordering artifact, not a property of your loss.

How much gradient weight does a short example actually steal?

Concretely. Two micro-batches, G = 2, completion-only masking, n = [20, 800]. Then N = 820, n̄ = 410.

• Tokens in the 20-token answer: scaled by 410 / 20 = 20.5x
• Tokens in the 800-token answer: scaled by 410 / 800 = 0.51x
• Relative skew between the two examples: 40x

Instruction datasets are exactly this shape. "Yes, that's correct." sits in the same shuffle as a 900-token code explanation. With per_device_train_batch_size=1, every micro-batch is one sample, so the naive loop hands each sample equal weight regardless of length — you have silently switched from token-level to sequence-level averaging, without choosing to.

Sequence-level averaging is a defensible objective. It is just not the one your config claims, it is not what your bs=8 baseline computed, and it makes short, terse, low-information targets the loudest gradient signal in the run. In practice you see it as a model that gets curt, drops formatting on long generations, and emits EOS early.

Why does DDP make it worse?

Because DistributedDataParallel all-reduces gradients with mean, not sum. So even with gradient_accumulation_steps=1, eight ranks holding n = [12, 15, 640, 30, 22, 700, 18, 25] unmasked tokens produce a mean of eight independently normalized gradients. The two ranks that drew long samples get their tokens crushed by the same n̄ / n_g factor.

This is why the bug survives so long in multi-node setups: people validate their loss on a single GPU with accum=1, where mean-of-means is trivially correct, and never re-check after scaling out.

How do you fix the gradient accumulation loss bug?

Count tokens for the entire optimizer step first, then use reduction="sum" and normalize once. The counting has to happen before any backward(), which means materializing the micro-batches for the step up front.

import torch
import torch.distributed as dist
from torch.nn.functional import cross_entropy

def optimizer_step(model, micro_batches, optimizer, world_size):
    # 1) Count real supervised tokens across the WHOLE step.
    #    labels are shifted, so the last position has no target -> slice [..., 1:]
    num_items = sum(
        (mb["labels"][..., 1:] != -100).sum() for mb in micro_batches
    ).to(torch.float32).cuda()

    # 2) Every rank must divide by the SAME global denominator.
    if world_size > 1:
        dist.all_reduce(num_items, op=dist.ReduceOp.SUM)

    for i, mb in enumerate(micro_batches):
        logits = model(input_ids=mb["input_ids"],
                       attention_mask=mb["attention_mask"]).logits

        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = mb["labels"][..., 1:].contiguous()

        # sum, not mean: normalization is global, not per micro-batch
        loss = cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)).float(),
            shift_labels.view(-1),
            ignore_index=-100,
            reduction="sum",
        )

        # DDP averages grads over ranks -> multiply by world_size to undo it
        loss = loss * world_size / num_items

        # only sync gradients on the final micro-batch
        if i < len(micro_batches) - 1 and world_size > 1:
            with model.no_sync():
                loss.backward()
        else:
            loss.backward()

    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

Three details that bite people writing this by hand:

1. The shift. labels[..., 1:] — if you count (labels != -100).sum() on the unshifted tensor you over-count by one per sequence that ends on a supervised token. Small, but it makes your bs-vs-accum equivalence test fail at 1e-3 and you will spend an afternoon on it.
2. .float() on the logits. In bf16 the summed loss over 8k tokens loses low bits fast. Upcast before the softmax, or accept noise larger than the effect you are measuring.
3. world_size multiply. Skip it and your effective learning rate is 1/W of what you set. This one hides beautifully — the run still trains, just slower, and you blame the LR schedule.

If you use Hugging Face Trainer or TRL's SFTTrainer on a recent version, this is already handled: the trainer computes num_items_in_batch for the accumulation window and threads it into the model's loss function, with average_tokens_across_devices gathering the count across ranks. The trap is subclassing:

# reintroduces the bug — kwarg dropped, model falls back to per-batch mean
class MyTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        ...

# correct — accept it and pass it through
class MyTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False,
                     num_items_in_batch=None):
        outputs = model(**inputs, num_items_in_batch=num_items_in_batch)
        ...

Custom model classes have the same failure: if your forward does not accept and honor num_items_in_batch, the kwarg is swallowed by **kwargs and you are back to mean-of-means with no warning.

Why do some people never see this?

Because packing hides it. If you train on packed sequences — concatenate documents to a fixed max_seq_length — then n_g is nearly constant across micro-batches, n̄ / n_g ≈ 1, and mean-of-means is approximately right. Pretraining pipelines are packed, which is why this bug is an SFT problem, not a pretraining problem.

It also disappears when every sample is a fixed-format short answer (classification-style tuning), and it is small when you compute loss over the full sequence including the prompt, since prompt length dominates and varies less than completion length. Completion-only masking, variable answer lengths, and per_device_train_batch_size=1 is the worst case — and it is the default recipe for QLoRA on a single 24 GB card.

How do you verify your trainer is correct?

One test, two minutes, and it is exact. Fix the seed, disable dropout, use fp32, take 8 samples with deliberately lopsided label counts, and compare:

# A: one batch of 8
loss_a = step_and_get_grad_norm(model, [batch_of_8])
# B: 8 micro-batches of 1, same samples, same order
loss_b = step_and_get_grad_norm(model, [b0, b1, b2, b3, b4, b5, b6, b7])

assert abs(loss_a - loss_b) / loss_a < 1e-5

Correct normalization gives agreement to float32 round-off. The buggy path will be off by percent-level or worse on skewed data — and the sign of the gap tells you which examples are being over-weighted. Run the same check with torchrun --nproc_per_node=2 against the single-GPU result to catch the DDP half.

Does this affect DPO and RL fine-tunes?

Yes, and it is more consequential there because the loss is a difference. In DPO, chosen and rejected completions rarely have the same length; a preference pair normalized per-sequence versus per-token changes what the implicit reward is measuring, which is why length-normalized DPO variants exist as a separate knob. Policy-gradient fine-tunes have the identical decision at the group-mean level: normalizing each rollout by its own length weights short rollouts more per token. The mechanism is the same n̄ / n_g factor.

The difference is that in DPO/RL the normalization is an explicit modeling choice you can defend. In SFT accumulation, nobody chose it — it fell out of a memory-saving config flag.

So why isn't accum=8 the same as batch size 8?

Because gradient accumulation reproduces a larger batch only if you normalize the loss over the total token count of the whole optimizer step. The default loop takes the mean inside each micro-batch and then averages those means, which weights each micro-batch equally instead of each token equally — so every token's gradient gets multiplied by n̄ / n_g, and short supervised targets can carry tens of times their fair share. The fix is mechanical: sum the per-token losses, divide once by the global unmasked-token count gathered across micro-batches and ranks, and multiply by world_size to undo DDP's mean all-reduce. Verify with a bs=N versus bs=1 × accum=N gradient-norm equality test in fp32; if they do not match to round-off, your accumulated run is optimizing a different objective than the one you benchmarked.