Train LLM — Sequence Length vs Batch Size

Two training runs can consume exactly the same number of tokens per step and still have completely different memory profiles. Sequence length and batch size are not interchangeable knobs, because attention cost grows with the square of the sequence length.

The idea

The number of tokens processed per optimizer step is:

Optimization cares about that total, but the hardware cares about how it is split. Two shapes are worth separating:

• Linear and MLP activations scale with the token count, batch × seq. Doubling batch and halving sequence length leaves this unchanged.
• Attention score tensors scale with batch × heads × seq². Here the sequence length enters twice, so long sequences are disproportionately expensive.

That asymmetry is the whole trade-off. A shape made of few long sequences saves on padding and per-sequence overhead, but it pays quadratically for attention memory. Many short sequences pay padding but keep the attention matrices small.

Related levers:

• Padding wastes compute. A batch of mixed-length sequences is padded to the longest one, and every model still runs over the pad tokens. Packing several short sequences into one fixed-length row removes most of that waste, as long as attention is masked to keep sequences from attending across boundaries.
• Gradient accumulation raises the effective batch without raising memory. Note that the loss must be averaged over the accumulation window, not summed.
• Activation checkpointing recomputes activations during the backward pass, trading roughly 30% more compute for a large memory reduction, which is often the only way to keep long sequences in memory.

Worked example

Hold the token budget at 16,384 per step and assume 32 attention heads:

• Batch 16, sequence 1024: 16 × 32 × 1024² ≈ 5.4 × 10⁸ attention elements per layer.
• Batch 2, sequence 8192: 2 × 32 × 8192² ≈ 4.3 × 10⁹ attention elements per layer.

Same tokens, same linear-layer cost, about 8 times the attention memory. Reshaping toward shorter sequences is the standard first move when a long-context run runs out of memory.

In code

accum = 8
for step, micro_batch in enumerate(loader):
    loss = model(**micro_batch).loss / accum   # average over the window
    loss.backward()
    if (step + 1) % accum == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        opt.zero_grad(set_to_none=True)

micro_batch stays small, while the optimizer sees the gradient of an eight-times-larger batch.

Check yourself

1. With the same token budget, which costs less attention memory: 8 sequences of 2048, or 2 of 8192?
2. What does gradient accumulation change about the optimizer step, and what does it not change?
3. Why does packing waste less compute than padding, and what breaks if you pack carelessly?

Key takeaways

• Tokens per step is the product of micro-batch, sequence length, and accumulation.
• Attention is quadratic in sequence length, which is what makes long sequences expensive.
• Accumulate and pack to raise the effective batch without paying the padding or memory bill.