Batch size, accumulation and the token count

Three numbers that multiply

What the optimiser sees is the effective batch size:

effective = per_device_batch_size × gradient_accumulation_steps × num_devices

A config with per_device_train_batch_size=2, gradient_accumulation_steps=8 on one GPU has an effective batch of 16. The gradients from eight forward-backward passes are summed before a single optimiser step. Memory is paid for two sequences at a time; the update behaves like sixteen.

This is the standard trick for training on a small card, and it is close to free: the only cost is that you compute eight forward passes before making any progress, so wall-clock time per optimiser step goes up while total compute stays roughly the same.

The number that actually matters is tokens

Examples are a bad unit. An effective batch of 16 chat exchanges could be 2,000 tokens or 30,000 depending on how long they are. Gradient noise, and therefore the right learning rate, tracks tokens far better than examples.

Log tokens per optimiser step. Most training frameworks do not print it, and it takes one line to compute. A typical small fine-tune sits somewhere between 8,000 and 64,000 tokens per step; if yours is 900, your gradients are extremely noisy and no learning rate will make the run stable.

The accumulation bug that shipped for years

This is the detail no tutorial mentions, and it caught a great many people.

If each micro-batch computes a mean loss over its own tokens, and you then average those means across accumulation steps, you have not reproduced a larger batch: you have given every micro-batch equal weight regardless of how many real tokens it contained. A micro-batch with 40 unmasked tokens and one with 900 count the same. Since masked instruction data has extremely variable completion lengths, the discrepancy is real.

Hugging Face’s Trainer and TRL had exactly this behaviour, and it was fixed in late 2024 (transformers 4.46 and the corresponding TRL release) by normalising over the total number of unmasked tokens in the accumulation window. Two practical consequences:

• Configs and results from before that fix are not directly comparable to results after it, especially where gradient_accumulation_steps was large.
• If you write your own training loop, normalise by the total token count across the accumulation window, not by the number of micro-batches. It is three lines and it is correct:

total_tokens = sum(batch_token_counts)
for batch, n in zip(batches, batch_token_counts):
    loss = model(**batch).loss * n / total_tokens   # weight by real tokens
    loss.backward()
optimizer.step()

Packing: the other way to fix the token count

Rather than padding every sequence to the length of the longest in the batch (which on instruction data can mean half your compute is spent on padding tokens) packing concatenates short examples until the buffer is full.

The gain is substantial: on a dataset whose examples average 300 tokens with a 2,048-token window, packing can cut training time by more than half, because you stop multiplying matrices by padding.

The hazard is attention leaking across the boundary. Without a correct block-diagonal attention mask, the model reads the end of example A while predicting the start of example B, and learns a transition that does not exist. Modern implementations handle this with position ids reset per document and a FlashAttention variable-length path; packing=True in TRL with a recent version does the right thing, but this is worth verifying rather than assuming, because the failure is invisible in the loss curve.

If you cannot verify it, use length-grouped batching instead (group_by_length=True). It sorts similar-length examples into the same batch so padding is minimal, with no cross-contamination risk at all. Most of the speed, none of the subtlety.

Choosing the numbers in practice

A workable procedure on a single card:

1. Set per_device_train_batch_size to the largest value that does not run out of memory at your maximum sequence length. Find it by bisection in about four attempts.
2. Set gradient_accumulation_steps so that effective batch lands between 16 and 64 for a LoRA on instruction data.
3. Log tokens per step. If it is under 4,000, raise accumulation rather than the learning rate.
4. Switch on gradient checkpointing if step 1 forces you below a per-device batch of 1 at the sequence length you need. It costs roughly 30% more time per step and often lets you double the batch.

And one thing not to do: do not chase a large effective batch on a small dataset. With 800 examples and an effective batch of 64, one epoch is twelve optimiser steps. Twelve steps will not move an adapter anywhere useful, no matter how clean each gradient is.