Warmup Guide — The FLOPs & Memory Algebra of Pre-Training

How to read this. Nothing here assumes you know what a FLOP, a matmul, a transformer, or a GPU is. Every term is built from nothing: what it is → why it exists → how it works underneath → what it costs in production → the misconception people carry. If you already know transformers, skip to Chapter 5 and do not skip Chapter 7 or Chapter 9.

Table of Contents

• Chapter 1: What a FLOP Is, and Why Anyone Counts Them
• Chapter 2: What a Matrix Multiply Costs
• Chapter 3: A Transformer, Assembled From Matmuls
• Chapter 4: Counting the Parameters
• Chapter 5: The 6ND Derivation
• Chapter 6: The Exact Per-Step Count
• Chapter 7: Where 6ND Breaks
• Chapter 8: Training Memory — The Thing That Actually Stops You
• Chapter 9: The KV Cache — The Serving Wall
• Chapter 10: From Chips and Days to (N, D)
• Chapter 11: Money, Watts and the Other Units
• Lab Walkthrough
• Success Criteria
• Interview Q&A
• Tips & Takeaways
• References

Chapter 1: What a FLOP Is, and Why Anyone Counts Them

What it is

FLOP = one floating-point operation. A single multiply, or a single add, on decimal numbers.

3.7 × 2.1   →  1 FLOP
3.7 + 2.1   →  1 FLOP
3.7 × 2.1 + 0.5  →  2 FLOPs   (one multiply, one add)

That last one — multiply then add — is so common that hardware implements it as a single instruction called an FMA (fused multiply-add). It still counts as 2 FLOPs, because it did two operations' worth of arithmetic. This convention matters enormously and trips up beginners constantly: an FMA is one instruction and two FLOPs.

Two related terms that are easy to confuse:

Work divided by speed equals time. That is the entire discipline of capacity planning.

Why it exists

Because you need a hardware-independent unit for "how much computation does this cost?" Wall time depends on your chip, your compiler, your batch size, and how many other jobs are running. FLOPs don't. They let you say "this experiment is 100× the work of that one" and be right on any hardware.

How it works underneath

Floating-point numbers store a sign, an exponent, and a mantissa. The bit width determines both precision and speed — and this becomes central in Phase 08.

Why BF16 beat FP16 for training: BF16 keeps FP32's 8 exponent bits and sacrifices mantissa bits. FP16 does the opposite. Gradients span an enormous dynamic range — some are 10⁻⁸ — so range matters more than precision, and FP16 silently flushes small gradients to zero. BF16 does not. This is why the loss-scaling machinery that FP16 training needed largely disappeared.

Peak FLOP/s is quoted per format, and the numbers roughly double each time you halve the bits:

H100 SXM (approximate, dense, no sparsity):
  FP32 :    67 TFLOP/s
  TF32 :   495 TFLOP/s
  BF16 :   990 TFLOP/s
  FP8  :  1979 TFLOP/s

Trap: vendors sometimes quote "with sparsity," which doubles the number again and requires a 2:4 structured-sparse model you almost certainly do not have. Always check.

Production significance

FLOPs are the currency the whole field trades in. "GPT-4-class" implicitly means "~10²⁵ FLOPs." Regulatory thresholds (the EU AI Act, the US executive order) are written in FLOPs. Compute budgets are allocated in FLOPs. When Feinberg asks "if I give you 1000 H100 for 30 days," the first move is always to convert that into a FLOP number.

The misconception

"More FLOPs means a slower model."

Not necessarily — it depends on whether you are compute-bound or memory-bound. During decode (generating one token at a time), a model can be doing almost no FLOPs and still be slow, because it is waiting on memory. Chapter 9 and Phase 05 make this precise. FLOPs measure work, not time.

Chapter 2: What a Matrix Multiply Costs

What it is

A matrix is a grid of numbers. Multiplying an (m × k) matrix A by a (k × n) matrix B gives an (m × n) matrix C, where:

$$ C_{ij} = \sum_{p=1}^{k} A_{ip} B_{pj} $$

In words: each output entry is the dot product of a row of A with a column of B.

The cost, derived

Count it directly:

• The output has m × n entries.
• Each entry is a sum of k products.
• Each product-and-accumulate is 2 FLOPs (one multiply, one add).

$$ \text{FLOPs} = 2 \cdot m \cdot n \cdot k $$

def matmul_flops(m, k, n):
    """FLOPs for (m x k) @ (k x n). The 2 is multiply + add."""
    return 2 * m * k * n

# One token (m=1) through a 4096 -> 16384 projection
print(matmul_flops(1, 4096, 16384))     # 134,217,728  = 2 * 4096 * 16384
# Which is exactly 2 x the number of weights in that layer:
print(2 * 4096 * 16384)                 # 134,217,728

This is the whole trick. For a layer applied to one token, m = 1, so the cost is 2 × k × n — and k × n is exactly the number of weights in that layer. Therefore:

A linear layer costs 2 × (its parameter count) FLOPs per token.

Hold onto that sentence. 6ND is three applications of it.

Why it exists (why neural nets are matmuls at all)

A neural network layer computes "every output is a weighted combination of every input." That is a matrix multiply, by definition. And it is a spectacular fit for hardware: matmuls are massively parallel (every output entry is independent), have high arithmetic intensity (lots of FLOPs per byte read), and have a regular access pattern. Chips are built around them — NVIDIA's tensor cores and Google's TPU systolic arrays exist to do exactly this one operation.

How it works underneath

A systolic array (the TPU's matmul unit) is a physical grid of small multiply-accumulate cells, e.g. 128×128. Data flows through it rhythmically:

        B (weights stay resident)
        ↓ ↓ ↓ ↓
A →  [·][·][·][·]  → partial sums flow right/down
A →  [·][·][·][·]
A →  [·][·][·][·]
A →  [·][·][·][·]
        ↓ ↓ ↓ ↓
     accumulated results out

Each cell does one multiply-accumulate per clock. A 128×128 array at 1 GHz does 128 × 128 × 2 = 32,768 FLOPs per cycle = 32.8 TFLOP/s from one array. Chips have several.

The consequence you must remember: the array is a fixed size. If your matrix dimension is not a multiple of 128, the hardware pads it, and you pay for the padding.

import math

def padded_matmul_flops(m, k, n, tile=128):
    """FLOPs the hardware ACTUALLY performs, including tile padding."""
    def up(x): return math.ceil(x / tile) * tile
    return 2 * up(m) * up(k) * up(n)

useful = matmul_flops(1024, 4000, 4000)
actual = padded_matmul_flops(1024, 4000, 4000)
print(f"useful {useful:,}  actual {actual:,}  efficiency {useful/actual:.1%}")
# useful 32,768,000,000  actual 34,359,738,368  efficiency 95.4%

Nearly 5% lost to a dimension being 4000 instead of 4096. This is one of the reasons production models use dimensions like 4096, 8192, 11008 — they are chosen to tile cleanly.

Production significance

~99% of a transformer's arithmetic is matmuls. Everything else — activations, norms, softmax — is a rounding error in FLOPs, though not in time (Phase 05 explains why: those operations run on a much slower unit and move a lot of memory).

The misconception

"The 2 in 2mkn is because of forward and backward."

No. The 2 is multiply-plus-add, in the forward pass alone. The forward/backward factor is a separate 3×, derived in Chapter 5. Confusing these gives you 6ND for the wrong reason and falls apart the moment someone asks about inference (2N, not 6N).

Chapter 3: A Transformer, Assembled From Matmuls

What it is

A decoder-only transformer — the architecture behind every modern LLM — is a stack of identical blocks. Each block has two sub-layers:

                    ┌─────────────────────────────┐
    input x ───────►│  RMSNorm                    │
                    │       │                     │
                    │       ▼                     │
                    │  ATTENTION                  │
                    │   Q = x·W_q  (d → n_h·d_h)  │  matmul
                    │   K = x·W_k  (d → n_kv·d_h) │  matmul
                    │   V = x·W_v  (d → n_kv·d_h) │  matmul
                    │   A = softmax(QKᵀ/√d_h)     │  matmul (seq-dependent!)
                    │   O = A·V                   │  matmul (seq-dependent!)
                    │   out = O·W_o (n_h·d_h → d) │  matmul
                    │       │                     │
                    │       + ◄─── residual       │
                    │       ▼                     │
                    │  RMSNorm                    │
                    │       ▼                     │
                    │  MLP / FFN                  │
                    │   h  = x·W_up   (d → d_ff)  │  matmul
                    │   g  = x·W_gate (d → d_ff)  │  matmul  (SwiGLU only)
                    │   a  = silu(g) * h          │  elementwise
                    │   out = a·W_down (d_ff → d) │  matmul
                    │       │                     │
                    │       + ◄─── residual       │
                    └───────┼─────────────────────┘
                            ▼   (repeat n_layers times)

The symbols (memorize these — every paper uses them)

Why it exists

Two ideas, each solving a specific failure of what came before:

Attention solves the RNN's problem. An RNN reads left to right through a fixed-size hidden state, so information from token 1 must survive 1,000 overwrites to influence token 1,000. Attention lets every token directly read every earlier token. The cost is O(T²) work instead of O(T) — which is why long context is expensive and why Chapter 7 matters.

The MLP is where knowledge lives. Attention routes information between positions; the MLP transforms it. Empirically, most factual knowledge sits in MLP weights — which is why the MLP is ~2/3 of the parameters, and why MoE (Phase 03) replaces the MLP and not attention.

How it works underneath — what each matmul is doing

• Q, K, V projections: turn the residual vector into a query ("what am I looking for?"), a key ("what do I offer?"), and a value ("what do I contribute?").
• QKᵀ: every query dots with every key → an attention score matrix of shape (T × T). Divided by √d_h to keep variance stable (without this, large d_h produces enormous logits and a saturated softmax with vanishing gradients).
• softmax: turns scores into weights that sum to 1 per row.
• A·V: each position's output is the weighted average of all values.
• W_o: projects the concatenated heads back to d.
• MLP: expand to d_ff, apply a nonlinearity, project back. The nonlinearity is the only reason depth helps — without it, a stack of linear layers collapses to a single linear layer.

The misconception

"Attention is where all the computation is."

At typical training context lengths (2k–8k), attention's sequence-dependent matmuls are roughly 5–15% of FLOPs. The projections and MLP dominate. Attention only takes over at very long context — and even then, the memory traffic of attention was the problem FlashAttention solved, not the FLOPs.

Chapter 4: Counting the Parameters

Before FLOPs, count weights. Every parameter is one learned number.

Per layer

def params_per_layer(d_model, d_ff, n_heads, n_kv_heads, d_head, gated=True):
    """Parameter count for one transformer block (biases omitted — modern LLMs drop them)."""
    # Attention projections
    w_q = d_model * n_heads * d_head
    w_k = d_model * n_kv_heads * d_head        # fewer under GQA/MQA
    w_v = d_model * n_kv_heads * d_head
    w_o = n_heads * d_head * d_model
    attn = w_q + w_k + w_v + w_o

    # MLP. Gated (SwiGLU) needs THREE matrices, not two.
    n_mats = 3 if gated else 2
    mlp = n_mats * d_model * d_ff

    # Norms: 2 per block, d_model each. Negligible but real.
    norms = 2 * d_model

    return {"attn": attn, "mlp": mlp, "norms": norms, "total": attn + mlp + norms}


p = params_per_layer(d_model=4096, d_ff=11008, n_heads=32,
                     n_kv_heads=8, d_head=128, gated=True)
for k, v in p.items():
    print(f"{k:6s} {v:>14,}")

attn      41,943,040
mlp      135,266,304
norms          8,192
total    177,217,536

Notice: the MLP is 76% of the block. That is why MoE targets it.

Whole model

def total_params(n_layers, d_model, d_ff, n_heads, n_kv_heads, d_head,
                 vocab, tied_embeddings=False, gated=True):
    per = params_per_layer(d_model, d_ff, n_heads, n_kv_heads, d_head, gated)["total"]
    body = n_layers * per
    embed = vocab * d_model                       # input embedding table
    unembed = 0 if tied_embeddings else vocab * d_model
    final_norm = d_model
    return {"body": body, "embed": embed, "unembed": unembed,
            "total": body + embed + unembed + final_norm}


# Llama-2-7B-ish shapes
m = total_params(n_layers=32, d_model=4096, d_ff=11008, n_heads=32,
                 n_kv_heads=32, d_head=128, vocab=32000)
print(f"{m['total']/1e9:.2f}B params   (body {m['body']/1e9:.2f}B, "
      f"embeddings {(m['embed']+m['unembed'])/1e9:.2f}B)")
# 6.74B params   (body 6.48B, embeddings 0.26B)

The embedding trap

At small N with a large vocabulary, embeddings are a huge fraction — and this is a classic source of wrong scaling-law fits.

for d, L, name in [(512, 8, "tiny"), (1024, 12, "small"),
                   (4096, 32, "7B"), (8192, 80, "70B")]:
    m = total_params(L, d, 4*d, d//128, d//128, 128, vocab=256000)
    emb_frac = (m["embed"] + m["unembed"]) / m["total"]
    print(f"{name:6s} total={m['total']/1e9:6.2f}B  embeddings={emb_frac:5.1%}")

tiny   total=  0.30B  embeddings=88.6%
small  total=  0.73B  embeddings=72.3%
7B     total= 10.69B  embeddings=19.6%
70B    total= 90.09B  embeddings= 4.7%

89% of a "tiny" model with a 256k vocabulary is embeddings. If you run a scaling ladder from tiny to 70B and use total parameters as N, your smallest points are measuring something almost entirely unlike your largest. This is why serious scaling work reports non-embedding parameters — a convention that looks pedantic until you see this table.

Chapter 5: The 6ND Derivation

Now the main event. Feinberg's slide states it and footnotes the reasoning; here it is in full.

Step 1 — Forward: 2N FLOPs per token

From Chapter 2: a linear layer costs 2 × (its parameters) FLOPs per token. A transformer is (almost entirely) a collection of linear layers. Sum over all of them:

$$ \text{forward FLOPs per token} = \sum_{\text{layers}} 2 \cdot (\text{params}) = 2N $$

Step 2 — Backward: 4N FLOPs per token

This is the step people get wrong, so go slowly. Consider one linear layer, Y = X · W.

During backprop, you arrive with dY (the gradient of the loss w.r.t. this layer's output) and you need two things:

(a) dX = dY · Wᵀ     ← the gradient to hand to the PREVIOUS layer.
                        Without this, backprop cannot continue.

(b) dW = Xᵀ · dY     ← the gradient of THIS layer's weights.
                        Without this, this layer never learns.

Both are matrix multiplies. Check the shapes: if X is (T × k) and W is (k × n), then Y and dY are (T × n).

• dX = dY · Wᵀ is (T × n) @ (n × k) → 2·T·n·k FLOPs.
• dW = Xᵀ · dY is (k × T) @ (T × n) → 2·k·T·n FLOPs.

Each is exactly the same size as the forward matmul (2·T·k·n). Two of them:

$$ \text{backward FLOPs per token} = 2 \times 2N = 4N $$

def verify_backward_is_2x(T, k, n):
    fwd = 2 * T * k * n          # Y = X @ W
    dX  = 2 * T * n * k          # dY @ W.T
    dW  = 2 * k * T * n          # X.T @ dY
    return fwd, dX + dW, (dX + dW) / fwd

print(verify_backward_is_2x(1024, 4096, 11008))
# (92341796864, 184683593728, 2.0)   <- backward is exactly 2x forward

Why this is exactly 2 and not approximately 2: because the two backward matmuls are the transposes of the forward one. Matmul cost 2mkn is symmetric in which operand you transpose — you touch the same number of elements either way. There is no hand-waving here; it is an identity.

Step 3 — Add

$$ C = \underbrace{2N}{\text{forward}} + \underbrace{4N}{\text{backward}} = 6N \text{ FLOPs per token} $$

Over D tokens:

$$ \boxed{C = 6ND} $$

The three numbers to keep separate

The 2N is as important as the 6N. Phase 02's lifetime-cost model is 6N·D_train + 2N·D_inference, and mixing them up by 3× wrecks the crossover analysis.

Sanity checks

def training_flops(n_params, n_tokens): return 6 * n_params * n_tokens
def inference_flops(n_params, n_tokens): return 2 * n_params * n_tokens

# Published runs — check the order of magnitude against reality.
runs = [
    ("GPT-3",        175e9,  300e9),
    ("Chinchilla",    70e9,  1.4e12),
    ("Llama-3-70B",   70e9, 15.0e12),
    ("Llama-3-8B",     8e9, 15.0e12),
]
for name, N, D in runs:
    print(f"{name:14s} N={N/1e9:6.1f}B  D={D/1e12:5.1f}T  C={training_flops(N,D):.2e} FLOPs")

GPT-3          N= 175.0B  D=  0.3T  C=3.15e+23 FLOPs
Chinchilla     N=  70.0B  D=  1.4T  C=5.88e+23 FLOPs
Llama-3-70B    N=  70.0B  D= 15.0T  C=6.30e+24 FLOPs
Llama-3-8B     N=   8.0B  D= 15.0T  C=7.20e+23 FLOPs

Two things to notice, both of which are the entire story of Phase 01:

1. Chinchilla used ~2× GPT-3's compute with a model 2.5× smaller — and beat it. That is the Kaplan-vs-Chinchilla result, visible in one table.
2. Llama-3-8B used more compute than GPT-3 at 1/22 the size. That is deliberate overtraining for serving efficiency — Phase 02's inference-aware scaling, in the wild.

Chapter 6: The Exact Per-Step Count

Feinberg's slide gives the precise identity for a training step:

$$ 18BTDF + 24BTDNH = 6 \cdot BT \cdot (3DF + 4DNH) $$

where B = batch, T = sequence length, D = d_model, F = d_ff, N = number of heads, H = head dimension. Let us verify it term by term — this is exactly the kind of derivation his hiring bar asks for.

The MLP term: 18BTDF

A gated MLP (SwiGLU) has three matrices, each D × F, so 3DF parameters per layer.

• Forward: 2 × 3DF = 6DF FLOPs per token
• Backward: 2 × that = 12DF
• Total: 18DF per token per layer

Over B·T tokens: 18BTDF. ✓

The attention-projection term: 24BTDNH

Four projections — W_q, W_k, W_v, W_o — each D × (N·H), so 4·D·N·H parameters.

• Forward: 2 × 4DNH = 8DNH
• Backward: 2 × that = 16DNH
• Total: 24DNH per token per layer

Over B·T tokens: 24BTDNH. ✓

The factoring

$$ 18BTDF + 24BTDNH = 6BT(3DF + 4DNH) $$

And 3DF + 4DNH is precisely the parameter count per layer (three MLP matrices + four attention projections). So the identity reads:

$$ \text{FLOPs per step} = 6 \times (\text{tokens per step}) \times (\text{params per layer}) $$

which, summed over layers, is 6ND. The slide identity is 6ND, written out in shapes.

def exact_step_flops(B, T, d_model, d_ff, n_heads, d_head, n_layers,
                     n_kv_heads=None, include_attention_matmuls=True, gated=True):
    """Per-optimizer-step training FLOPs, decomposed."""
    if n_kv_heads is None:
        n_kv_heads = n_heads                     # MHA -> reproduces the 24BTDNH form
    n_mats = 3 if gated else 2
    mlp = n_mats * 6 * B * T * d_model * d_ff * n_layers
    # W_q and W_o scale with n_heads; W_k and W_v scale with n_kv_heads.
    proj = 6 * B * T * (2 * d_model * n_heads * d_head
                        + 2 * d_model * n_kv_heads * d_head) * n_layers

    # The sequence-dependent attention matmuls: QK^T and A@V.
    # Forward: 2 * (2 * B * n_heads * T * T * d_head); backward doubles it again.
    attn = 0
    if include_attention_matmuls:
        attn = 6 * 2 * B * n_heads * T * T * d_head * n_layers

    return {"mlp": mlp, "attn_proj": proj, "attn_seq": attn,
            "total": mlp + proj + attn}


cfg = dict(B=8, T=8192, d_model=8192, d_ff=28672,
           n_heads=64, d_head=128, n_layers=80)
r = exact_step_flops(**cfg)
for k, v in r.items():
    print(f"{k:10s} {v:>22,}  ({v/r['total']:5.1%})")

mlp        22,166,154,415,964,160  (63.6%)
attn_proj   8,444,249,301,319,680  (24.2%)
attn_seq    4,222,124,650,659,840  (12.1%)
total      34,832,528,367,943,680  (100.0%)

12% is in the sequence-dependent attention matmuls at T = 8192. That is the part 6ND throws away — and Chapter 7 shows what happens when you push T further.

Note on the causal mask. A causal model only attends to earlier positions, so in principle the attn_seq term could be halved. Most published FLOP accounting (including Kaplan's) does not halve it, because the dense implementation computes the full matrix and masks. FlashAttention does skip the masked blocks. Pick one convention, state it, and be consistent — this is a common source of two people's MFU numbers disagreeing by 5%.

Chapter 7: Where 6ND Breaks

Four regimes. Know all four and the error in each; this is a standard interview probe.

Break 1 — Attention at long context

The attention matmuls (QKᵀ and A·V) scale with T², not with N.

$$ \text{attention FLOPs per token} \approx 12 \cdot L \cdot T \cdot d_h \cdot n_h / n_h = 12 \cdot L \cdot T \cdot d_{\text{model}} $$

(using n_h · d_h ≈ d_model). Ratio to the 6N term:

def attention_fraction(n_params, n_layers, d_model, seq_len):
    """Fraction of training FLOPs in the sequence-dependent attention matmuls."""
    per_token_body = 6 * n_params
    per_token_attn = 6 * 2 * n_layers * seq_len * d_model     # QK^T + A@V, fwd+bwd
    return per_token_attn / (per_token_body + per_token_attn)

for T in (2048, 8192, 32768, 131072, 1048576):
    f = attention_fraction(70e9, 80, 8192, T)
    print(f"T={T:>9,}  attention = {f:6.1%} of FLOPs   "
          f"(6ND error {f/(1-f):7.1%})")

T=    2,048  attention =   3.7% of FLOPs   (6ND error    3.8%)
T=    8,192  attention =  13.3% of FLOPs   (6ND error   15.3%)
T=   32,768  attention =  38.0% of FLOPs   (6ND error   61.4%)
T=  131,072  attention =  71.1% of FLOPs   (6ND error  245.4%)
T=1,048,576  attention =  95.2% of FLOPs   (6ND error 1963.4%)

At 2k context 6ND is ~4% low. At 1M context it is off by 20×. Verdict: use 6ND freely below ~8k, add the attention term above that, and never use it at all for long-context work.

Break 2 — Mixture of Experts

For an MoE, N in 6ND must be the active parameter count — what a single token actually routes through — not the total.

def moe_params(n_layers, d_model, d_ff, n_experts, top_k,
               n_heads, n_kv_heads, d_head, shared_experts=0):
    """Total (memory) vs active (FLOPs) parameters for an MoE transformer."""
    attn = params_per_layer(d_model, d_ff, n_heads, n_kv_heads,
                            d_head, gated=True)["attn"]
    one_expert = 3 * d_model * d_ff
    router = d_model * n_experts

    total_per_layer = attn + router + (n_experts + shared_experts) * one_expert
    active_per_layer = attn + router + (top_k + shared_experts) * one_expert
    return {"total": n_layers * total_per_layer,
            "active": n_layers * active_per_layer}

m = moe_params(n_layers=60, d_model=7168, d_ff=2048, n_experts=256, top_k=8,
               n_heads=128, n_kv_heads=128, d_head=128, shared_experts=1)
print(f"total  {m['total']/1e9:7.1f}B   <- what you must STORE (memory, HBM)")
print(f"active {m['active']/1e9:7.1f}B   <- what you must COMPUTE (FLOPs, 6ND)")
print(f"sparsity ratio: {m['total']/m['active']:.1f}x")

total    707.4B   <- what you must STORE (memory, HBM)
active    52.1B   <- what you must COMPUTE (FLOPs, 6ND)
sparsity ratio: 13.6x

Use active for 6ND; use total for memory. Getting this backwards is the single most common MoE arithmetic error, and here it is off by 13.6×. This is the audience question on Feinberg's slide: "What About MoEs?"

Break 3 — Embeddings

The output unembedding (d_model → vocab) is a real matmul. 6ND counts it if you included embeddings in N, and misses it if you did not — and Chapter 4 showed embeddings can be 91% of a small model. Fix: report N as non-embedding parameters, and add the unembedding term explicitly:

def unembed_flops_per_token(d_model, vocab):
    return 6 * d_model * vocab       # fwd 2 + bwd 4

Break 4 — Activation checkpointing (the MFU/HFU distinction)

To save memory, training frameworks discard intermediate activations and recompute them during the backward pass. Full recomputation adds an extra forward pass:

model FLOPs    (what 6ND counts) : 6N per token
hardware FLOPs (what you pay for): 8N per token with full recompute
                                   ~6.5N with selective recompute

This is exactly the gap between MFU (uses 6ND, the honest number) and HFU (uses 8ND, the flattering one). When someone quotes a utilization figure, ask which. HFU is always higher.

Chapter 8: Training Memory — The Thing That Actually Stops You

You almost never run out of FLOPs. You run out of HBM. Here is every term.

The four consumers

def training_memory_bytes(n_params, n_activations_bytes=0,
                          optimizer="adam", precision="mixed_bf16",
                          zero_stage=0, dp_degree=1):
    """Bytes of HBM for a training step, itemized.

    Mixed-precision Adam, the standard recipe:
      - bf16 weights for the forward/backward matmuls        : 2 bytes/param
      - bf16 gradients                                       : 2 bytes/param
      - fp32 master weights (for numerically stable updates) : 4 bytes/param
      - fp32 Adam first moment  (m)                          : 4 bytes/param
      - fp32 Adam second moment (v)                          : 4 bytes/param
                                                       total : 16 bytes/param
    """
    weights = 2 * n_params
    grads = 2 * n_params
    if optimizer == "adam":
        opt = 12 * n_params          # fp32 master + m + v
    elif optimizer == "sgd_momentum":
        opt = 8 * n_params           # fp32 master + momentum
    elif optimizer == "adafactor":
        opt = 4 * n_params           # factored second moment: ~O(sqrt) not O(n)
    else:
        raise ValueError(f"unknown optimizer: {optimizer}")

    # ZeRO/FSDP shards these across data-parallel replicas.
    if zero_stage >= 1: opt //= dp_degree
    if zero_stage >= 2: grads //= dp_degree
    if zero_stage >= 3: weights //= dp_degree

    return {"weights": weights, "grads": grads, "optimizer": opt,
            "activations": n_activations_bytes,
            "total": weights + grads + opt + n_activations_bytes}


N = 70e9
for stage, dp in [(0, 1), (1, 64), (2, 64), (3, 64)]:
    m = training_memory_bytes(N, zero_stage=stage, dp_degree=dp)
    print(f"ZeRO-{stage} (dp={dp:2d}): {m['total']/1e9:8.1f} GB/device "
          f"[w {m['weights']/1e9:6.1f} | g {m['grads']/1e9:6.1f} | "
          f"o {m['optimizer']/1e9:6.1f}]")

ZeRO-0 (dp= 1):   1120.0 GB/device [w  140.0 | g  140.0 | o  840.0]
ZeRO-1 (dp=64):    293.1 GB/device [w  140.0 | g  140.0 | o   13.1]
ZeRO-2 (dp=64):    155.3 GB/device [w  140.0 | g    2.2 | o   13.1]
ZeRO-3 (dp=64):     17.5 GB/device [w    2.2 | g    2.2 | o   13.1]

Look at ZeRO-0: 1.12 TB for a 70B model. An H100 has 80 GB. The model cannot be trained without sharding — which is why Phase 04 exists. And notice that the optimizer states are 75% of it, which is why ZeRO-1 (shard the optimizer only — cheap, almost no extra communication) is the highest-value single change you can make.

Activations

Activations are the intermediate tensors kept for the backward pass. Without checkpointing:

$$ \text{activation bytes} \approx B \cdot T \cdot d_{\text{model}} \cdot L \cdot c \cdot \text{bytes} $$

where c ≈ 10–30 depending on what the framework stores.

def activation_bytes(B, T, d_model, n_layers, c=16, bytes_per=2, checkpointing=None):
    raw = B * T * d_model * n_layers * c * bytes_per
    if checkpointing == "full":
        # Store only layer boundaries; recompute the rest. ~sqrt-ish saving in practice.
        return B * T * d_model * n_layers * bytes_per
    if checkpointing == "selective":
        return raw * 0.3
    return raw

for ck in (None, "selective", "full"):
    b = activation_bytes(8, 8192, 8192, 80, checkpointing=ck)
    print(f"checkpointing={str(ck):10s} {b/1e9:8.1f} GB")

checkpointing=None       1374.4 GB
checkpointing=selective   412.3 GB
checkpointing= full        85.9 GB

Activations can exceed the weights. Checkpointing trades ~30% more FLOPs for ~16× less activation memory — almost always the right trade at scale, and the reason MFU numbers look "low."

Chapter 9: The KV Cache — The Serving Wall

What it is

During generation, the model produces one token at a time. Each new token attends to every previous token, so it needs their keys and values. Recomputing them every step would be O(T²); instead you cache them. That cache is the KV cache.

The size

$$ \text{KV bytes} = 2 \cdot L \cdot n_{kv} \cdot d_h \cdot T \cdot B \cdot b $$

The 2 is K and V. b is bytes per element (2 for bf16).

def kv_cache_bytes(n_layers, n_kv_heads, d_head, seq_len, batch, bytes_per=2):
    return 2 * n_layers * n_kv_heads * d_head * seq_len * batch * bytes_per

# A 70B-class model: 80 layers, 64 query heads, head_dim 128.
for name, kv in [("MHA (64 kv)", 64), ("GQA-8 (8 kv)", 8), ("MQA (1 kv)", 1)]:
    for B in (1, 32):
        gb = kv_cache_bytes(80, kv, 128, 8192, B) / 1e9
        print(f"{name:14s} batch={B:3d}  {gb:8.1f} GB")

MHA (64 kv)    batch=  1     21.5 GB
MHA (64 kv)    batch= 32    687.2 GB
GQA-8 (8 kv)   batch=  1      2.7 GB
GQA-8 (8 kv)   batch= 32     85.9 GB
MQA (1 kv)     batch=  1      0.3 GB
MQA (1 kv)     batch= 32     10.7 GB

Read the MHA row again: 687 GB of KV cache for 32 concurrent users at 8k context — on top of the 140 GB of weights. That is nine H100s of pure cache.

Why this is the serving constraint

Weights are a fixed cost — 140 GB for a 70B model in bf16, no matter how many users. The KV cache is a per-request, per-token cost. So:

Available for KV = (total HBM) − (weights) − (workspace)
Max concurrent   = Available / (KV bytes per request)

def max_concurrent(total_hbm_gb, n_params, kv_per_request_gb,
                   bytes_per_param=2, workspace_gb=10):
    weights_gb = n_params * bytes_per_param / 1e9
    free = total_hbm_gb - weights_gb - workspace_gb
    if free <= 0:
        return 0
    return int(free / kv_per_request_gb)

# 8x H100 = 640 GB, serving a 70B model at 8k context
kv_mha  = kv_cache_bytes(80, 64, 128, 8192, 1) / 1e9
kv_gqa8 = kv_cache_bytes(80,  8, 128, 8192, 1) / 1e9
print("MHA  :", max_concurrent(640, 70e9, kv_mha),  "concurrent requests")
print("GQA-8:", max_concurrent(640, 70e9, kv_gqa8), "concurrent requests")

MHA  : 22 concurrent requests
GQA-8: 182 concurrent requests

One architecture decision — 8 KV heads instead of 64 — is an 8× difference in serving throughput, at essentially no quality cost. This is the cleanest single example of what Feinberg means by "inference co-design," and it must be decided before pre-training starts, because it changes the weights.

Under the hood — why decode is memory-bound

Generating one token requires reading every weight and the whole KV cache, to do a tiny amount of arithmetic:

def decode_arithmetic_intensity(n_params, kv_bytes, batch):
    flops = 2 * n_params * batch          # 2N per token, times batch
    bytes_moved = n_params * 2 + kv_bytes # weights (shared!) + KV (per request)
    return flops / bytes_moved            # FLOPs per byte

for B in (1, 8, 64, 256):
    kv = kv_cache_bytes(80, 8, 128, 8192, B)
    ai = decode_arithmetic_intensity(70e9, kv, B)
    print(f"batch={B:4d}  arithmetic intensity = {ai:6.1f} FLOP/byte")

batch=   1  arithmetic intensity =    1.0 FLOP/byte
batch=   8  arithmetic intensity =    6.9 FLOP/byte
batch=  64  arithmetic intensity =   28.7 FLOP/byte
batch= 256  arithmetic intensity =   43.3 FLOP/byte

An H100's balance point is ~990e12 / 3.35e12 ≈ 295 FLOP/byte. Every one of these is far below it — decode is memory-bound at every realistic batch size. That single fact explains batching, PagedAttention, GQA, quantization, and speculative decoding. Phase 05 formalizes it with the roofline.

Chapter 10: From Chips and Days to (N, D)

Now answer Feinberg's opening question end to end.

HARDWARE = {
    # name:          (peak bf16 FLOP/s, HBM bytes, HBM bandwidth B/s, watts)
    "H100":          (990e12,  80e9, 3.35e12, 700),
    "A100-80":       (312e12,  80e9, 2.03e12, 400),
    "TPU v5e":       (197e12,  16e9, 0.819e12, 170),
    "TPU v5p":       (459e12,  95e9, 2.77e12,  600),
}

def budget_to_flops(chip, n_chips, days, mfu=0.4):
    peak, *_ = HARDWARE[chip]
    return n_chips * peak * mfu * days * 86400

def chinchilla_split(C):
    """Chinchilla: N and D scale as C^0.5 each, with D ~= 20*N."""
    # C = 6ND and D = 20N  =>  C = 120 N^2  =>  N = sqrt(C/120)
    N = (C / 120) ** 0.5
    return N, 20 * N

C = budget_to_flops("H100", 1000, 30)
N, D = chinchilla_split(C)
print(f"Budget: 1000 H100 x 30 days @ 40% MFU")
print(f"  C = {C:.3e} FLOPs")
print(f"  Chinchilla-optimal: N = {N/1e9:.1f}B params, D = {D/1e12:.2f}T tokens")
print(f"  check: 6ND = {6*N*D:.3e}")

Budget: 1000 H100 x 30 days @ 40% MFU
  C = 1.026e+24 FLOPs
  Chinchilla-optimal: N = 92.5B params, D = 1.85T tokens
  check: 6ND = 1.026e+24

That is the answer to his opening question. But a senior answer does not stop there — it adds four sanity checks:

def sanity_checks(N, D, chip, n_chips):
    peak, hbm, bw, watts = HARDWARE[chip]
    checks = {}

    # 1. Does the optimizer state even fit across the cluster?
    need_gb = training_memory_bytes(N)["total"] / 1e9
    have_gb = n_chips * hbm / 1e9
    checks["memory fits (sharded)"] = (need_gb < have_gb * 0.7, f"{need_gb:.0f} GB / {have_gb:.0f} GB")

    # 2. Do we HAVE that many unique tokens?
    checks["data available"] = (D < 15e12, f"needs {D/1e12:.2f}T tokens")

    # 3. Will it be cheap enough to serve?
    checks["servable"] = (N < 100e9, f"{N/1e9:.0f}B params, {2*N/1e9:.0f} GB in bf16")

    # 4. Energy
    joules = n_chips * watts * (6 * N * D) / (n_chips * peak * 0.4)
    checks["energy"] = (True, f"{joules/3.6e6/1e3:.1f} MWh, ~${joules/3.6e6*0.12:,.0f} at $0.12/kWh")
    return checks

for k, (ok, detail) in sanity_checks(N, D, "H100", 1000).items():
    print(f"  [{'OK ' if ok else 'FAIL'}] {k:24s} {detail}")

  [OK ] memory fits (sharded)    1480 GB / 80000 GB
  [OK ] data available           needs 1.85T tokens
  [OK ] servable                 92B params, 185 GB in bf16
  [OK ] energy                   504.0 MWh, ~$60,480 at $0.12/kWh

And then the senior judgment on top: at 92B parameters this model needs 3 H100s just to hold its weights for serving. If the product is a real-time assistant, you should deliberately undershoot Chinchilla — train a 30B model on 5.5T tokens for the same C — and accept slightly worse loss for 3× cheaper serving. That is Phase 02's inference-aware scaling, and it is the actual job.

Chapter 11: Money, Watts and the Other Units

The conversions you will be asked for, in one place.

def cost_report(chip, n_chips, days, price_per_chip_hour, kwh_price=0.12, pue=1.2):
    peak, hbm, bw, watts = HARDWARE[chip]
    hours = days * 24
    rental = n_chips * hours * price_per_chip_hour
    energy_kwh = n_chips * watts * pue * hours / 1000
    return {
        "rental $": rental,
        "energy kWh": energy_kwh,
        "energy $": energy_kwh * kwh_price,
        "FLOPs @40% MFU": n_chips * peak * 0.4 * hours * 3600,
        "$ per 1e21 FLOPs": rental / (n_chips * peak * 0.4 * hours * 3600) * 1e21,
    }

for k, v in cost_report("H100", 1000, 30, price_per_chip_hour=2.50).items():
    print(f"{k:20s} {v:>14,.2f}")

rental $               1,800,000.00
energy kWh               604,800.00
energy $                  72,576.00
FLOPs @40% MFU  1.026e+24
$ per 1e21 FLOPs           1,753.65

$ per 1e21 FLOPs is the number to carry in your head. It converts any research proposal into dollars instantly: "that ablation is 3e21 FLOPs" → "about $5,300 of compute." Suddenly you can reason about a research portfolio (Phase 12) in a currency executives understand — and you can immediately see that a ladder of small runs (Phase 01) costs a rounding error compared to the flagship it de-risks.

Note the ratio: rental is 25× the raw electricity here. That is not a contradiction of Feinberg's "99% of TCO is power" — a rental price bundles amortized silicon, datacenter capital, cooling, networking, and margin. His claim is about the total cost of operating the hardware over its life, where energy (chips + cooling + power delivery) dominates. When you own the fleet, the electricity bill is the thing that scales with usage; when you rent, it is hidden inside the hourly rate. Phase 08 builds the owner's-view model properly.

Lab Walkthrough

Lab 01 — Transformer FLOPs, Memory & Budget Calculator

You will implement, in order:

1. matmul_flops(m, k, n) — the 2mkn primitive. Everything else calls this.
2. params_per_layer(...) / total_params(...) — with gated, GQA, and tied-embedding handling, returning the non-embedding count separately.
3. training_flops_exact(...) — reproducing 18BTDF + 24BTDNH, with the sequence-dependent attention term as a separate line item so you can see Break 1.
4. moe_parameter_split(...) — total vs active, which is Break 2.
5. training_memory(...) — weights, grads, optimizer states, activations, ZeRO stages.
6. kv_cache_bytes(...) / max_concurrent_requests(...) — the serving wall.
7. budget_to_flops(...) / chinchilla_split(...) — the round trip.
8. budget_report(...) — the whole thing, plus the four sanity checks.

Start with matmul_flops and the invariant test (backward == 2 × forward). If that passes, the rest is bookkeeping. If it does not, re-read Chapter 5.

The trap in this lab is units. Bytes vs GB, FLOPs vs FLOP/s, active vs total parameters, per-token vs per-step vs per-run. The tests check all four; the docstrings name the unit in every signature.

Success Criteria

• LAB_MODULE=solution pytest test_lab.py -v — all green.
• Your lab.py passes after filling the TODOs.
• python solution.py prints a full budget report.
• You can derive 6ND on a whiteboard in under two minutes, including why backward is exactly 2× forward.
• Given (B, T, d_model, d_ff, n_heads, d_head, n_layers) you can produce per-step FLOPs without the calculator, to within 10%.
• You can state all four breaks of 6ND and estimate the error in each.
• You can compute a KV cache and explain the GQA saving in terms of concurrent requests.
• You have run the calculator on a cluster you might plausibly be given and written a paragraph on what you would train, including where you would deliberately deviate from Chinchilla.

Interview Q&A

Q: Derive the training FLOPs of a transformer. Forward: a linear layer costs 2 × params FLOPs per token, because each output entry is a dot product of length k and each term is a multiply-add. Summing over layers gives 2N. Backward: each forward matmul becomes two backward matmuls — dX = dY·Wᵀ to propagate, dW = Xᵀ·dY to learn — each the same size as forward, giving 4N. Total 6N per token, C = 6ND. Excludes sequence-dependent attention, which adds roughly 12·L·T·d_model per token.

Q: Why exactly 2× for the backward pass, not approximately? Because both backward matmuls are transposes of the forward one, and matmul cost 2mkn is invariant to which operand is transposed — you touch the same number of elements. It is an identity, not an empirical rule.

Q: A 400B-parameter MoE with 8 of 128 experts active. What is its training compute per token? 6 × active, not 6 × 400B. Compute the active count: attention + router + top_k experts per layer. If active is ~35B, it costs like a 35B dense model to train — but you must store 400B parameters, which is ~800 GB in bf16 and dictates your sharding. Total for memory, active for FLOPs.

Q: You have 1,000 H100s for 30 days. What do you train? At 40% MFU that is ~1.03e24 FLOPs. Chinchilla-optimal is ~92B params on ~1.85T tokens. But I would not do that: a 92B model needs 3 H100s just to hold weights at serve time. If this model serves real traffic I would deliberately overtrain a smaller one — say 30B on 5.5T tokens for the same budget — trading ~0.02 nats of loss for ~3× cheaper serving, and I would justify the trade with a lifetime-cost calculation over projected served tokens.

Q: Why is MFU only 40% in that estimate — is that bad? No. MFU is the fraction of peak matmul throughput achieved, and a transformer is not pure matmul: it also runs vector ops (norms, activations, softmax), moves activations to and from HBM, runs collectives, and executes the optimizer step. 35–55% is the normal band for large-scale training. Also check whether you are being shown MFU or HFU — HFU counts recomputation from activation checkpointing as useful work and is always higher.

Q: Your 70B model needs 1.1 TB of memory to train but your GPUs have 80 GB. What now? Shard. Optimizer states are ~75% of that, so ZeRO-1 (shard optimizer states across data-parallel ranks) is the cheapest first move — it adds almost no communication. ZeRO-2 adds gradient sharding, ZeRO-3/FSDP shards parameters too and costs an all-gather per layer. Then add tensor parallelism within a node (high-bandwidth NVLink) and pipeline parallelism across nodes, plus activation checkpointing. Phase 04 covers picking the combination.

Q: Why is decode memory-bound but prefill compute-bound? Prefill processes T tokens at once against the same weights, so arithmetic intensity is high — you amortize each weight read across many tokens. Decode processes one token per sequence, so you read every weight to do 2N FLOPs — an intensity around 1–35 FLOP/byte, far below an H100's ~295 balance point. That asymmetry is why they get different parallelism strategies (Phase 06) and why batching is the primary decode optimization.

Tips & Takeaways

Tips

• Say the unit out loud at every step. "Two N per token." "One-forty gigabytes." Most errors at this level are unit errors, not algebra errors.
• Memorize 6ND, 2N, and 16 bytes/param for Adam. Those three cover most napkin math.
• Sanity-check every answer against a known run. GPT-3 ≈ 3.1e23, Llama-3-70B ≈ 6.3e24. If your number is not between them, you slipped a factor.
• Always ask "total or active?" the moment MoE appears.
• Always ask "MFU or HFU?" the moment a utilization number appears.
• Non-embedding parameters is the right N for scaling work. Say so explicitly; it signals you have actually done this.
• Carry $ per 1e21 FLOPs in your head. It turns research proposals into money.

Takeaways

1. A matmul costs 2 × params per token. Every other number in this phase descends from that.
2. Backward is exactly 2× forward, for a structural reason.
3. 2N inference, 6N training. Never mix them.
4. 6ND excludes attention: ~3% error at 2k context, 22× at 1M.
5. For MoE: active for FLOPs, total for memory.
6. Memory, not compute, is the binding constraint. Adam mixed-precision is ~16 bytes/param.
7. The KV cache is the serving wall; n_kv_heads is the lever, and it is set before training.
8. Decode is memory-bound at every realistic batch size. That one fact explains most of serving.
9. Chinchilla-optimal is a starting point, not an answer. Serving economics move it.

References

Primary

• Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides (the C = 6ND slide and its footnote; the 18BTDF + 24BTDNH identity)
• Austin et al., How To Scale Your Model — https://jax-ml.github.io/scaling-book/ (the definitive treatment of this arithmetic; do its exercises)

Papers

• Kaplan et al., Scaling Laws for Neural Language Models, 2020 — https://arxiv.org/abs/2001.08361 (Appendix: the FLOP accounting)
• Hoffmann et al., Training Compute-Optimal Large Language Models, 2022 — https://arxiv.org/abs/2203.15556
• Vaswani et al., Attention Is All You Need, 2017 — https://arxiv.org/abs/1706.03762
• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (MQA), 2019 — https://arxiv.org/abs/1911.02150
• Ainslie et al., GQA, 2023 — https://arxiv.org/abs/2305.13245
• Rajbhandari et al., ZeRO, 2019 — https://arxiv.org/abs/1910.02054
• Chen et al., Training Deep Nets with Sublinear Memory Cost (checkpointing), 2016 — https://arxiv.org/abs/1604.06174
• Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022 — https://arxiv.org/abs/2205.05198
• Pope et al., Efficiently Scaling Transformer Inference, 2022 — https://arxiv.org/abs/2211.05102

Hardware

• NVIDIA H100 architecture whitepaper — peak FLOP/s per format, HBM bandwidth
• Google Cloud TPU documentation — v5e / v5p specifications
• Jouppi et al., In-Datacenter Performance Analysis of a Tensor Processing Unit, 2017 — https://arxiv.org/abs/1704.04760 (systolic arrays)