AI Infra Interviews logo
🧮 Napkin Math & Capacity
Foundational

Training FLOPs: 6ND

The compute needed to train a language model is six floating-point operations per parameter per token: two for the forward pass and four for the backward. Multiply by the parameter count and the token count and you have the whole run's compute, which is the number every fleet-sizing, time-to-train and cost question starts from. This page derives the 6, states the attention correction and when it matters, and shows where the 2N of inference comes from, so the reader can rebuild the formula rather than recall it.

TL;DR: Training compute C ≈ 6 × N × D FLOPs, for N parameters and D training tokens. The forward pass costs 2 FLOPs per parameter per token (a multiply and an add for every weight); the backward pass costs twice that, because it computes two gradients (with respect to the activations, to keep propagating, and with respect to the weights, to update them), each a matmul of the same size. Llama 3.1 70B on 15.6 trillion tokens: 6 × 70.6e9 × 15.6e12 ≈ 6.6e24 FLOPs. Inference is the forward pass alone: about 2N per generated token. Attention adds a term that grows with context and is a small correction at 8k and a large one at 128k.

Where the 2 comes from

A dense layer multiplies an input vector by a weight matrix. For every weight in the matrix, the multiplication does one multiply and one add (multiply the input element by the weight, add it into the running sum). Two floating-point operations per weight, per token that passes through. A model with N parameters, almost all of them in weight matrices, therefore does about 2N FLOPs to push one token through the forward pass.

forward pass, one token, one weight matrix of shape (d_in × d_out):
  output = input(1 × d_in) × W(d_in × d_out)
  FLOPs = 2 × d_in × d_out     (one multiply-add per weight, and a multiply-add is 2 FLOPs)
  the matrix has d_in × d_out parameters, so FLOPs per parameter = 2

whole model: sum over all matrices = 2 × N FLOPs per token

Embeddings, norms, activations and softmax are rounding errors next to this: the embedding lookup does no multiplies, and the elementwise ops are a few FLOPs per element rather than per weight.

Where the 6 comes from

Backpropagation through the same layer computes two things. It needs the gradient of the loss with respect to the layer's input, to pass backward to the previous layer; that is a matmul with W transposed, the same size as the forward one, so 2 FLOPs per weight. And it needs the gradient with respect to W itself, which is an outer product of the input and the incoming gradient, also 2 FLOPs per weight. So the backward pass costs 4N, and forward plus backward is 6N per token.

rendering diagram…

The first layer does not need the gradient with respect to its input (nothing is before it), so the true count is a hair under 6N, and the interviewer does not care. What they do care about is that you can say which two matmuls make up the backward pass, because "backward is twice forward" stated without the reason sounds memorized.

The attention correction

Attention has a cost that 6ND does not count: the score matrix, queries against keys over the whole context, which grows with the sequence length rather than with the parameter count. Per layer and per token, the attention scores and the weighted sum of values cost about 4 × s × d FLOPs in the forward pass for context length s and model width d (2 for QKᵀ, 2 for the product with V), and three times that including backward. The PaLM paper's version of the full formula is:

FLOPs per token ≈ 6N + 12 × n_layers × d_model × s
                    (first term: parameters; second: attention over the context)
Llama 3.1 70B, n_layers 80, d_model 8,192:
  s = 8,192:   6 × 70.6e9 = 4.24e11;   12 × 80 × 8,192 × 8,192 = 6.4e10  → attention adds ~15%
  s = 128k:    attention term = 1.03e12 → attention is 2.4x the parameter term
sanity: at pre-training context the correction is a modest markup; at long-context stages it dominates,
        which is why the long-context phase of a run is budgeted separately.

For the estimation round, 6ND is the answer and "plus an attention term that matters at long context" is the sentence that shows you know its limits.

Worked runs

Llama 3.1 70B, 15.6T tokens (the paper's figure):
  C = 6 × 70.6e9 × 15.6e12 = 6.6e24 FLOPs

Llama 3.1 405B, 15.6T tokens:
  C = 6 × 405e9 × 15.6e12 = 3.8e25 FLOPs
  Meta's model card lists about 30.8 million cumulative H100-hours for the 405B;
  check: 3.8e25 ÷ (30.8e6 h × 3,600 s/h) = 3.4e14 FLOP/s per GPU = 340 TFLOPS,
         which is 34% of the 989 dense peak. The paper reports 38 to 43% MFU for the
         pre-training phases; the cumulative hour count includes restarts, evaluations and
         the long-context stage, so the two numbers agree once that overhead is allowed for.
  sanity: a formula, a published hour count and a published MFU land within a few points
          of each other; that is the check to do out loud.

an 8B model on 15T tokens:
  C = 6 × 8e9 × 15e12 = 7.2e23 FLOPs, about a ninth of the 70B run for the same data

The training calculator turns any of these into a fleet size and a duration; GPU-Hours and Time to Train does it by hand.

Inference: 2N, and why it is not 6N

Serving runs only the forward pass, so each generated token costs about 2N FLOPs (plus the attention term over the context, which at long context is again the bigger piece). A 70B model spends 1.4e11 FLOPs per token; at 989 TFLOPS that is 0.14 ms of pure compute, and yet a single-stream decode step takes about 42 ms on an H100. The gap is the point: decode is not compute-bound, it is bound by reading 141 GB of weights, and the FLOP count is the wrong lens for it (Bandwidth-Bound Decode). Prefill, which processes thousands of tokens against the same weights in one pass, is where the 2N per token FLOP count becomes the limit.

Working it in the room

"How much compute did it take to train X?" wants three sentences: 2 per parameter per token forward because a multiply-add per weight; 4 backward because two gradients; 6ND with the numbers substituted. Then the correction and the check against a published hour count if one is known. The follow-up held back is "why is inference 2N and not 6N?", and the second follow-up is "so why is decode slow if it is only 2N?", which is the bandwidth answer. The answer that sounds right and fails is stating 6ND as a fact with no derivation: it is the most memorized formula in the field and interviewers know it.

What to remember

  • Forward: 2 FLOPs per parameter per token (a multiply-add per weight). Backward: 4 (two gradients, each a matmul of the same size). Training: 6ND.
  • 70B on 15.6T tokens ≈ 6.6e24 FLOPs; 405B ≈ 3.8e25, which checks against Meta's 30.8M cumulative H100-hours at an MFU in the mid-thirties once restart and evaluation overhead is allowed for.
  • Attention adds 12 × layers × d_model × s per token: ~15% at 8k for a 70B, dominant at 128k.
  • Inference is 2N per token; decode is slow anyway because it is bandwidth-bound, not FLOP-bound.
  • Derive it out loud; the formula is too well known to state bare.
RELATED CONCEPTS
LESSONS THAT TEACH THIS
PRACTICE THIS IN REAL QUESTIONS