Inside a Training Step

Five lines that hide everything

Every transformer ever trained, at every scale, reduces to this loop:

for batch in loader:
    output = model(batch)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Six lines. Nothing in them tells you that model(batch) will allocate seventeen gigabytes of intermediate tensors, that loss.backward() costs roughly twice the arithmetic of the line above it, or that optimizer.step() — mathematically the most trivial line on the page — can eat a fifth of your step time because it reads and writes two hundred gigabytes of memory to do a handful of multiplications per parameter.

This article opens the loop up. We will walk one step from the moment bytes leave storage to the moment updated weights land back in HBM, keeping track of four quantities the whole way: what arithmetic is done, what memory is allocated, what data moves and how far, and what depends on what. That last one is the sleeper. Dependencies are what make training hard to parallelize, and they are the reason activation memory exists at all.

Throughout, I will use one concrete model so that every number is checkable rather than gestural. A dense decoder-only transformer, roughly Llama-2-7B shaped:

Everything below is arithmetic on those six numbers. If you can reproduce the numbers, you can re-derive them for your own model — which is the actual skill, because the model you care about is never the one in the blog post.

By the end you should be able to answer the question this whole series is building toward: what exactly are we trying to split across GPUs?

Before the GPU: a batch's journey

Training starts long before the first kernel launches. A batch has to be found on disk, decoded, turned into tensors, moved into a region of host memory the DMA engine is allowed to touch, and shipped across PCIe. Each of those hops has its own bandwidth, and they differ by three orders of magnitude.

Three details in that path matter more than the rest.

Worker processes. Decoding and tokenizing happen on CPU cores, in separate processes, because Python cannot do it concurrently in one. Those workers run ahead of the GPU: while step N computes, the workers are already preparing N+1. If they can't keep up, the GPU idles, and no amount of parallelism strategy will save you. This is input starvation, and it is the single most common cause of a training job running at half the throughput its owner expects.

Pinned memory. Ordinary host allocations are pageable — the OS can move them. A DMA engine cannot chase a moving target, so a copy from pageable memory is really two copies: kernel staging buffer, then device. Page-locking the buffer removes the staging copy and, more importantly, makes the transfer asynchronous: it can be issued on a side stream and overlap with compute.

Double buffering. Asynchrony is only useful if there is somewhere to put the incoming batch while the current one is still in use. Two device-side buffers, alternating, gives the copy engine and the SMs disjoint working sets, and the H2D transfer disappears behind the forward pass.

Inside one transformer block

Now the arithmetic. A decoder block is seven operations in a fixed order, and each one has a completely different profile: some are enormous matrix multiplications that saturate the tensor cores, some are elementwise passes that do almost no arithmetic and exist only to move memory around.

The figure below walks a single [1, 4096, 4096] activation through one block, showing at each step the tensor shape, the FLOPs, the HBM traffic, and the arithmetic intensity that follows from those two — the ratio that decides whether the step is limited by the tensor cores or by memory bandwidth.

Read the pattern rather than the individual numbers. The five projections — QKV, output, gate, up, down — carry 92% of the arithmetic and are comfortably compute-bound. The four cheap operations — two norms, the SiLU-multiply, two residual adds — carry under a twentieth of a percent of the arithmetic and are entirely memory-bound. They are not free. The SwiGLU activation alone reads and writes 270 MB per layer to perform one multiply per element; across 32 layers that is 8.6 GB of HBM traffic for arithmetic you could do on a phone.

This is why kernel fusion matters, and why torch.compile tends to buy 10–20% on transformer training without touching a single matrix multiply. It isn't making the GEMMs faster. It is deleting round trips to HBM in the gaps between them.

Reading shapes off an architecture diagram

The habit worth building is converting any block diagram into shapes without thinking. Everything downstream — parallelism, communication volume, memory — is a function of these.

Two rows in that table should make you uncomfortable. The attention score matrix is 1.07 GB per layer — larger than every other tensor in the block combined, and it grows quadratically. This is the single tensor FlashAttention exists to never materialize: it tiles the computation so the S × S matrix lives in SRAM and is discarded, turning attention's memory cost from O(S²) into O(S).

The logits are 262 MB in BF16, but the cross-entropy path typically upcasts to FP32 and keeps the softmax output plus a gradient buffer — call it 1.5 GB for one sequence. Double the vocabulary, as recent models have, and this single tensor becomes the peak-memory event of the entire step. It is the reason fused and chunked cross-entropy kernels have become standard equipment.

Why compute scales like N × D

A matrix multiply of an [m, k] by a [k, n] produces m × n outputs, each a dot product of length k: one multiply and one add per term.

Now apply it to a linear layer. Its weight is [H, H'], which is exactly its parameter count, and it is applied once to every token. So a forward pass through a linear layer costs 2 × (params) × (tokens), and summing over every linear layer in the model gives:

The whole scaling-law formalism rests on this identity, which is why compute budgets are quoted in N × D at all. But notice what the derivation assumed: every FLOP comes from a weight being used. Attention breaks that assumption. The QKᵀ and PV products multiply activations by activations; no parameters are involved, so they contribute nothing to N and are invisible in 2ND.

So at S = 4096, attention is about 17% of the block's arithmetic and the 2ND approximation is off by a tolerable margin. At S = 32k, attention has overtaken every projection in the model and 2ND under-counts your compute by more than a third. The approximation isn't wrong; it has a domain of validity, and long-context training sits outside it. Causal masking halves the attention term, which buys you roughly a factor of two of headroom and no more.

Backprop: one GEMM in, two out

The backward pass is not magic and it is not, despite the name, a reversal. It is the chain rule applied to a graph, and for the operation that dominates transformers — the matrix multiply — it produces a startlingly clean result.

Take Y = XW, with X of shape [S, H] and W of shape [H, H']. Given the gradient of the loss with respect to the output, dY, we need two things: the gradient to pass further back, and the gradient of the weight itself.

Three GEMMs where there was one gives the familiar estimate:

State the assumptions out loud, because each one fails somewhere real. Freeze a vision tower and its weight-gradient GEMM disappears — that branch costs 2ND for the frozen portion, and you no longer pay it. Turn on full activation checkpointing and you re-run the forward pass, pushing the coefficient toward 8. A mixture-of-experts model routes each token to a fraction of its parameters, so the N in the formula is active parameters, not total. And at long context, as we just saw, the attention term you dropped is no longer a rounding error.

Used with those caveats, 6ND is an excellent estimate. Used as a law, it will mislead you by 30% in either direction — which, on a run that costs six figures, is a real number.

The graph and the activation lifetime

Autograd builds a directed acyclic graph as the forward pass runs. Every operation that produces a tensor requiring gradients records a node holding: the backward function, references to its inputs' nodes, and — critically — any tensors it will need later. Calling .backward() traverses that graph in reverse topological order, and each node's backward function fires as soon as the gradient of its output is available.

The consequence is a memory profile shaped like a mountain. Look at what the dW = XᵀdY rule demands: to compute the weight gradient of layer 1, you need the input to layer 1 — a tensor produced at the very beginning of the forward pass and not consumed until the very end of the backward pass. Its lifetime spans the entire step. Layer 2's input spans slightly less. Layer 32's, barely any.

That nesting is why activation memory dominates, and why it scales with depth even though nothing about a single layer is large. It also explains a phenomenon everyone training large models has hit: the job survives the forward pass, reports a comfortable memory reading, and then OOMs milliseconds into the backward. Nothing went wrong in the backward pass. The peak was always going to be at the turn; you just hadn't reached it yet.

A useful mental model: the forward pass is a push, the backward pass is a pop. Memory grows monotonically through the forward, peaks at the loss, and drains through the backward as each saved tensor is consumed and freed. Anything that reduces what gets pushed — checkpointing, smaller microbatches, fused kernels that save one tensor instead of four — moves the whole mountain down.

The training memory ledger

We can now account for every byte an H100 holds during training. There are five categories and they behave completely differently: three scale with parameters, one scales with batch and sequence, and one is whatever the allocator happens to be doing.

The three that scale with parameters

In the standard mixed-precision recipe, each parameter carries: a BF16 copy used for the forward and backward GEMMs (2 bytes), a BF16 gradient (2 bytes), an FP32 master copy the optimizer actually updates (4 bytes), and Adam's two moment estimates in FP32 (4 + 4 bytes). Sixteen bytes per parameter, before a single activation exists.

An 80 GB H100 cannot hold the states of a 7B model. Not with a batch size of one; not with a sequence length of one. The model is 13.5 GB and the machine has 80 GB, and it still does not fit — because training a parameter costs eight times what storing it does.

The one that scales with the batch

Activations are a different animal. Counting the tensors a SwiGLU block must save for its backward — the block input, both norm outputs, Q/K/V, the attention output, the projection output, and the three MLP intermediates — comes to roughly 16H elements per token per layer with FlashAttention in play.

Note what is not in that expression: the parameter count. Activation memory depends on depth, width, batch, and sequence — you cannot predict it from model size, and two models with identical parameter counts can differ by 3× here depending on how wide and shallow they are and how many tensors their kernels save.

The one nobody budgets for

Workspace is everything the allocator holds that isn't yours: cuBLAS split-k buffers, FlashAttention's softmax statistics, NCCL's communication buffers, the temporaries inside a fused kernel. It is usually 2–5 GB, and it is peaky — it spikes inside a kernel and vanishes, which is exactly the behaviour that fragments a caching allocator.

Which brings up the distinction that actually matters when a job dies: peak memory versus steady-state memory. PyTorch's caching allocator never returns memory to the driver, so memory_allocated tells you what is live and memory_reserved tells you what has been claimed. A job with 60 GB allocated and 78 GB reserved is one fragmentation event from an OOM, and the error message will name a 200 MB allocation as the culprit. Watch the gap between the two numbers, not either one alone.

Four precisions, not one

"Training in BF16" is four independent decisions collapsed into one phrase, and conflating them is how people end up with silently diverging runs.

• Storage precision — how weights sit in HBM. Halving it halves parameter memory and halves the bytes a GEMM must read.
• Activation precision — how intermediates are saved for backward. This is the lever on the 17 GB above.
• Compute precision — what the tensor cores multiply. BF16 inputs are what unlock ~989 TFLOP/s on an H100 against ~67 for FP32.
• Accumulation precision — what the multiply-adds sum into. This is almost always FP32 inside the tensor core, and it is the reason low-precision training works at all.

The distinction between the last two is the whole trick. A BF16 multiply is fine; a BF16 sum of 4096 terms is not, because BF16's 8-bit mantissa means adding a small number to a large accumulator silently does nothing. Tensor cores multiply in BF16 and accumulate in FP32, so the dot product is accurate even though its inputs are not.

The same logic explains the FP32 master weights in the ledger. An Adam update at learning rate 3e-4 changes a weight by a relative amount around 10⁻⁴ — below BF16's resolution, which is about 2⁻⁸ ≈ 4×10⁻³. Update a BF16 weight directly and thousands of steps round to no change at all. Keep the master copy in FP32, apply the update there, and cast down for compute. Four bytes per parameter is the price of the updates not vanishing.

BF16 versus FP16 is a related choice: BF16 keeps FP32's 8-bit exponent and gives up mantissa bits, so it has the dynamic range to represent gradients without a loss scaler. FP16 has more precision and less range, and needs gradient scaling to keep small gradients from flushing to zero. FP8 goes further still — E4M3 for forward tensors, E5M2 for gradients — but needs per-tensor scaling factors and careful handling of outlier channels, which is why it lands first in the projections and last, if ever, in the norms.

Activation checkpointing

Seventeen gigabytes of activations for a single sequence is the problem. Activation checkpointing is the answer, and it is the cleanest compute-for-memory trade in all of training.

The idea follows directly from the lifetime picture. Instead of saving every intermediate tensor in the forward pass, save only a sparse set of checkpoints — typically one per transformer block — and discard everything between them. When the backward pass reaches a block, re-run its forward from the nearest checkpoint to regenerate the tensors it needs, use them, and free them immediately.

Checkpointing every block is the blunt version. The sharp version is selective: recompute only the operations that are cheap in FLOPs and expensive in memory, and keep the ones that are the reverse. The SiLU-multiply is nearly free to recompute and saves 90 MB per layer. A GEMM output is expensive to recompute and saves the same 33 MB. Recompute the first, keep the second, and you capture most of the memory saving for a fraction of the compute penalty — Korthikanti et al. report keeping the overhead near 4% instead of 33%.

The asymmetry gets sharper in multimodal models. In a VLM, the vision tower runs on far more tokens per image than the language model sees after pooling, and its activations are correspondingly larger relative to its parameter count. Full checkpointing on ViT blocks and selective checkpointing on LLM blocks is often the right split, precisely because the two towers sit at different points on this trade curve.

Gradient accumulation

Large-batch training is a statistical requirement — you want a global batch of a few million tokens for stable optimization at scale — and it collides with a hardware limit: activation memory scales linearly with microbatch size, and you only have 80 GB.

The resolution is to decouple the batch you optimize with from the batch you compute with. Run several microbatches in sequence, letting their gradients sum into the same buffer, and only then take an optimizer step.

The three factors are not interchangeable, and treating them as such is a classic source of confusion. N_devices buys you speed: the microbatches run concurrently, so wall-clock time per step stays flat. N_accum buys you batch size only: the microbatches run one after another, so step time scales linearly with it. Accumulation is not parallelism. It is a way to reach a target global batch on hardware you don't have enough of.

Two implementation details bite people. First, the loss must be scaled by 1/N_accum (or the gradients divided afterwards), or your effective learning rate silently scales with the accumulation count. Second, anything that averages over the batch — batch norm, and any loss normalized by token count — does not decompose across microbatches, so per-microbatch normalization gives a different answer than the equivalent large batch would.

There is also a forward reference here worth flagging. In data-parallel training, gradients must be all-reduced across devices before the step. With accumulation, you only need that all-reduce on the last microbatch — the earlier ones are pure local accumulation. That one observation is why no_sync() exists and why accumulation is a communication-reduction technique as much as a memory one.

The optimizer step is bandwidth

Adam looks like nothing. Per parameter: two exponential moving averages, a square root, a division. Maybe a dozen FLOPs. It should be instant.

Count the bytes instead. For one parameter, with BF16 compute weights and FP32 master and moments, the step must read the gradient (2), read the master weight (4), read and write m (8), read and write v (8), write the master weight (4), and write the downcast BF16 weight (2).

For 6.7B parameters that is 189 GB of HBM traffic per step. At a realistic 2.5 TB/s achieved, about 75 ms — against roughly 330 ms of forward and backward compute for a 4096-token step. The most arithmetically trivial line in the training loop is 20% of the step.

This is also why fused and foreach optimizers are not micro-optimizations. A naive implementation launches a handful of kernels per parameter tensor; with hundreds of tensors that is thousands of launches, each reading and writing HBM separately. A fused optimizer does the whole update in one pass over each tensor, and a multi-tensor implementation batches small tensors into single launches. Same math, a third of the traffic, and no launch overhead.

And it is the reason optimizer state sharding is such an effective idea. If the step is bandwidth-bound and each device only owns 1/k of the parameters, the step gets k× faster in addition to using k× less memory. That is ZeRO's second gift, and it is usually undersold next to the memory saving.

The step, end to end

Assemble everything. One training step, three concurrent lanes, real durations for our 6.7B model on one H100 with a 4096-token microbatch.

The shape to internalize is the offset between lanes. The dataloader is working on batch N+1 while the GPU computes batch N; the H2D copy for N+1 overlaps the forward pass of N on a side stream. Everything to the left of the GPU lane is supposed to be invisible. When it becomes visible, you have a bubble.

Utilization is not efficiency

nvidia-smi reporting 100% means one thing: at least one kernel was resident on the device during each sampling interval. It says nothing about how many SMs that kernel occupied, whether it used the tensor cores, or whether it spent its life waiting on HBM. A kernel that copies memory back and forth forever will report 100% utilization at approximately zero useful throughput.

The measurement that means something is model FLOPs utilization:

MFU deliberately excludes recomputation, because recomputation is a choice you made, not work the model needed. Hardware FLOPs utilization includes it. The gap between them is the price of your memory strategy: with full activation checkpointing, HFU ≈ 1.33 × MFU, and if you see that ratio in your logs it tells you recomputation is costing you a third of the machine.

Good large-scale transformer training lands somewhere around 35–55% MFU. Below 30%, something specific is wrong and it is nearly always one of: an input pipeline that can't keep up, unfused elementwise kernels between the GEMMs, a microbatch too small to fill the tensor cores, a communication pattern that isn't overlapping, or an optimizer that isn't fused. Every one of those is diagnosable from the timeline in the previous figure, which is why the timeline is the artifact worth building first when a run is slow.

Two walls, one question

Everything above collapses into two failure modes.

The time wall: the model fits, and training will take four months. This is an arithmetic statement. 6ND for a 7B model on 2T tokens is 8.1 × 10²² FLOPs; one H100 at 45% MFU delivers about 4.4 × 10¹⁴ FLOP/s, so the run takes roughly six years. There is no clever kernel that fixes six years. You need more chips, working on different data at the same time.

The memory wall: the model does not fit at all. 108 GB of model states plus 17 GB of activations on an 80 GB card. This one is not fixed by more chips working on different data — every device would still need the full 108 GB. It is fixed by making each device hold less of the model.

Two walls, two different escapes, and the second is the one that forces real architectural change. Which brings us to the question this whole dissection was for:

Look back at what we found. The batch divides trivially — different tokens are independent until the gradient all-reduce. The optimizer states divide cleanly, because each parameter's update depends only on that parameter. The parameters divide at the cost of gathering them back before every GEMM. The sequence divides everywhere except attention, which needs every token to see every other token. The layers divide, but they are strictly sequential, so dividing them creates idle time by construction.

Five axes, five different communication bills. That is The Parallelism Playbook, and it is the next article.

References

Grouped by topic. Every entry links to its primary source.