A run has fixed HBM memory. Config A uses seq_len=4096 and global batch=256. The team doubles seq_len to 8192 without adding memory or activation checkpointing changes. What is the most likely immediate training system outcome?
At fixed memory, doubling sequence length about doubles activation memory per sequence, so sequences per step must drop by half or more, and tokens/sec falls as attention work per token grows.
Think of a delivery van with a fixed floor. Each box holds a row of items, and every box suddenly becomes twice as long. Only half as many boxes fit now. Each trip still carries about the same number of items, but packing long boxes takes more effort, so trips get slower. Try to load as many boxes as before and the doors will not shut: that is the out-of-memory error. You can make two loading rounds before each delivery (gradient accumulation) or fold the boxes flatter (activation checkpointing), but both cost time.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Memory on a GPU holds four things during training: the weights, the gradients, the optimizer state, and the activations saved for the backward pass. The first three depend only on the model. Activations depend on how many tokens the GPU processes at once, so they are the part that moves when sequence length changes.
The scenario fixes memory and doubles sequence length from 4,096 to 8,192. The outcome depends on how activation memory grows, which batch number has to give, and what happens to the compute per token.
Where activation memory goes
A standard estimate for one transformer layer, with 16-bit activations and no parallelism, is:
Here s is sequence length, b is the micro-batch size in sequences, h is the hidden size, and a is the number of attention heads. The 34sbh part covers the inputs to the linear layers, the normalizations, the activation function, and the dropout masks. It grows linearly with s. The 5as²b part is the attention score matrix, its softmax output, and its dropout mask. It grows with s squared.
The quadratic part is large. For a 7B-class layer (h = 4096, a = 32) at s = 4,096, the linear part is about 0.57 GB per sequence and the attention part is about 2.7 GB. So most of the memory is attention scores.
Memory-efficient attention kernels such as FlashAttention never store the full score matrix. They compute attention in tiles and recompute what they need in the backward pass. With such a kernel, the quadratic term is gone, and activation memory per sequence grows linearly with s. Doubling s then doubles memory per sequence, from 0.57 GB to 1.14 GB per layer. Without it, the total grows from 3.3 GB to 11.9 GB per layer, about 3.7x.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- Llama 3 started pretraining at 4M tokens per batch with 4,096-token sequences, then moved to 8M tokens with 8,192-token sequences after 252M tokens.
- Llama 3 extended context from 8K to 128K in six stages over about 800B tokens and used context parallelism to fit the long sequences.
What an interviewer would ask next. Try answering before peeking at the approach.
QWhich mitigation should be tried first: activation checkpointing or more sharding?
Find the dominant memory term. Activations scale with sequence length and micro-batch; ZeRO or FSDP sharding only shrinks weights, gradients and optimizer state.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Assuming longer context is free because the parameter count did not change, or that gradient noise must rise. Halving sequences per step still keeps about 1M tokens per step.
60 second bullets to scan on the way to the call.
Linear and quadratic terms in activation memory per layer
Effect of memory-efficient attention on memory per sequence
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.