Why a Training Step Costs Three Times Its Forward Pass

A training step over $D$ tokens on a model with $N$ parameters performs approximately $6ND$ floating-point operations. The number is quoted often enough to look like a definition. It is an estimate, derived from counting multiply-accumulate operations in the forward and backward passes and ignoring everything that is not a matrix multiply, and knowing which assumptions it rests on determines when it will be wrong by a large factor.

The estimate is useful because compute, not memory, sets the wall-clock time of a well-configured training run, and because the same accounting exposes the ratio that decides whether the hardware is being used at all.

Counting one multiply-accumulate

A matrix-vector product of a $M \times K$ matrix with a $K$-vector needs $MK$ multiplications and about as many additions. Convention counts a fused multiply-add as two floating-point operations, so the product costs $2MK$ FLOPs. For a batch of $B$ vectors the cost is $2MKB$; the operation is unchanged, the count scales with the number of vectors.

A transformer layer is mostly such products. Projections from width $H$ to width $H$ cost $2H^2$ FLOPs per token each, and there are several per layer. The feed-forward block typically expands to $4H$ and contracts back, costing $2 \cdot H \cdot 4H$ and $2 \cdot 4H \cdot H$, or $16H^2$ per token. Summing the projections and the feed-forward gives roughly $24H^2$ per token per layer, and with $L$ layers the whole forward pass is about $24 L H^2$ FLOPs per token. The parameter count of the same model, excluding embeddings, is approximately $12 L H^2$: each $H \times H$ projection contributes $H^2$ weights and each feed-forward block contributes $8H^2$. Substituting gives forward cost $\approx 2N$ FLOPs per token.

Multiply by the number of tokens $D$: forward is $2ND$. This is the origin of the first coefficient.

Where the second factor of two comes from

Backward computes, for each parameter, the derivative of the loss with respect to that parameter. Two quantities are needed: the gradient with respect to the layer input, which propagates the error to earlier layers, and the gradient with respect to the weights, which is accumulated for the optimiser.

For a linear layer $y = Wx$, the input gradient is $W^\top \dot{y}$, costing the same $2MK$ per vector as the forward product. The weight gradient is the outer product $\dot{y} x^\top$, also $2MK$ per vector. So backward through a linear layer costs twice what forward through it cost. This holds generally: every matrix multiply in the forward pass appears twice in the backward pass, once transposed against the incoming gradient and once as the parameter gradient. The only forward operations that are cheaper in backward are elementwise ones, and they are a small fraction of the total for wide models.

Backward is therefore $4ND$, and the step total is $6ND$. The update itself is elementwise over $N$ parameters and negligible by comparison.

Attention, which the estimate omits

The $6ND$ formula counts only the parameter-bearing matrix multiplies. Attention has products between activations, with no parameters, and they are invisible to the formula. For a sequence of length $S$ with $A$ heads of dimension $d = H/A$, the score computation $QK^\top$ costs $2S^2 d$ per head per layer, and the product with the values costs the same, for $4S^2 d A = 4S^2 H$ FLOPs per token per layer in the forward pass. Multiplied by $L$ layers and tripled for the backward pass, attention adds roughly $12 L S H$ FLOPs per token.

The ratio of attention to parameter cost is therefore about $S/(6H)$ per token. For $S = 2048$ and $H = 4096$ that is roughly eight percent, which is why the parameter-only formula is a reasonable approximation at moderate context lengths. At $S = 32768$ it is about 1.3 times the parameter term: the estimate no longer describes the step. Long-context training is attention-bound, and the fix is a different attention algorithm rather than a different sharding scheme.

Embeddings are the other omission. An embedding table of $V \times H$ entries is large in parameter count but receives one lookup per token, so its FLOP cost is a gather rather than a multiply. Models with very large vocabularies have parameter counts that overstate their compute.

The same arithmetic as a memory-traffic problem

An accelerator has a peak arithmetic rate and a peak memory bandwidth, and the ratio of the two is its arithmetic intensity in FLOPs per byte. Current accelerators sit in the range of a few hundred FLOPs per byte of off-chip bandwidth: peak throughput in the hundreds of teraFLOP/s against bandwidth in the low terabytes per second, so the crossover is on the order of one hundred to several hundred FLOPs per byte.

A matrix multiply of shape $M \times K$ against $K \times B$ touches $MK + KB + MB$ elements and performs $2MKB$ FLOPs, giving intensity $\frac{2MKB}{MK + KB + MB}$ FLOPs per element. When $M$ and $B$ are both small relative to $K$, the intensity approaches $2 \cdot \frac{MB}{M+B}$, which is small. A batch of 8 sequences of length 2048 with hidden width 4096 gives a per-layer matmul with $M = 16384$ rows and $K = B = 4096$: intensity around 3600 FLOPs per element, comfortably above the crossover. A decode-style product with one row and $K = 4096$ has intensity near 2 and is entirely bandwidth-bound.

Training sits on the good side of that line as long as the micro-batch is large enough that the weight matrix is reused across many samples. This is the reason micro-batching below a certain size stops saving time even though it saves memory, and the reason gradient accumulation cannot recover throughput lost to an undersized micro-batch.

Using the number

The estimate supports three calculations. Total training compute is $6ND$ for a run over $D$ tokens; dividing by achievable throughput gives an elapsed time. Achievable throughput is a fraction of peak, and for large transformer training the fraction is governed by whether the step is compute- or communication-bound rather than by the arithmetic rate itself.

The second calculation is the compute-optimal token count. The scaling-law literature fits a relationship between model size and the number of tokens that minimises loss for a fixed compute budget, and the fitted exponents are what make “train a smaller model on more tokens” a defensible strategy rather than a guess. The third is the recompute tax: any mechanism that recomputes activations instead of storing them adds forward FLOPs, and each recomputed layer adds $2ND_{\text{layer}}$ to the step. Activation checkpointing is exactly that trade, and the $6ND$ accounting says how much it costs before any measurement is taken.

def step_flops(n_params, n_tokens, n_layers, seq_len, hidden):
    """Approximate forward+backward FLOPs for one training step."""
    param_term = 6 * n_params * n_tokens
    # Attention scores and value product: 4 * S^2 * H per token per layer
    # in forward, tripled for backward.
    attn_term = 3 * 4 * n_layers * seq_len * hidden * n_tokens
    return param_term + attn_term

The function returns a number in the range of $10^{15}$ to $10^{18}$ for realistic training configurations, which is the useful scale: it says whether a run is hours or weeks of device time before any code is written, and it makes the attention term visible at the point where it stops being ignorable.

Two corrections are worth applying before the number is used for a decision. The first is that the update step, which the formula treats as negligible, is not negligible in a different currency: it reads and writes every parameter and every optimiser moment, so its memory traffic is on the order of $20N$ bytes per step against the matmuls’ arithmetic. That traffic does not show up in the FLOP count at all, and for a model whose step is already communication-bound it can add percent-level time.

The second is that the formula counts a step with no recomputation. Checkpointed training, gradient accumulation that shortens the micro-batch, and any other mechanism that changes the shape of the matmuls alter either the FLOP total or the efficiency at which those FLOPs execute. The formula is an estimate of work, not of time; converting it to time requires the achieved fraction of peak, which is a property of the shapes and the interconnect rather than of the model.