Part I — Basics

Resource accounting — FLOPs and the memory wall

Lesson 04 handed you a model and a stable optimizer: AdamW, warmup→cosine, bf16, gradient clipping. You can now train a model on a GPU. What you cannot yet do is answer the only two questions that decide whether a run is even launchable: will this fit, and how long will it take? This lesson builds the two back-of-envelope models — compute (6ND) and memory (16 bytes/param) — and then runs them on a 7B model against one 80 GB H100. The answer is no. That single "no" is the hinge into the entire systems half of this course.

The previous step left this broken

Lesson 04 made the training step correct and stable — but it quietly assumed the step fits in the box. It does not. AdamW keeps a first moment, a second moment, and an fp32 master copy of every parameter; bf16 keeps gradients; the forward pass leaves a mountain of activations for the backward to consume. None of that was priced. Before you rent a single GPU-hour you must be able to predict, on paper, the bytes resident on each device and the wall-clock to burn through D tokens. Get this wrong and you discover it as an out-of-memory crash forty minutes into a multi-day run, or as a budget that runs dry at 60% of the planned tokens.

Linear position

Forced by: you must predict fit and runtime before spending GPU-hours — a stable optimizer is worthless if its state overflows the device or the run can't finish in budget.
New idea: two arithmetic models — compute (forward ≈ 2N/token, backward ≈ 4N/token ⇒ C = 6ND) and memory (the 16-bytes/param mixed-precision AdamW ledger + an activation term), plus gradient checkpointing as the dial that trades compute for activation memory.
Forces next: the ledger shows a 7B model's training state alone blows past one 80 GB GPU before a single activation is stored — so the work must be split across many GPUs, which is the systems half (starting with the hardware itself, lesson 6).

The plan

Five moves. (1) Derive 2N/token forward from a single matmul, then 6N/token for training, recovering C = 6ND from lesson 0. (2) Define MFU — what fraction of the peak you actually get, and what "good" looks like. (3) Build the mixed-precision AdamW memory ledger to ≈16 bytes/param and run it on 7B. (4) Add the activation term and show gradient checkpointing buying it back. (5) Assemble a fit/no-fit table and the runtime estimate, then drive the calculator until 7B overflows one GPU.

1 · Where the FLOPs go: 2N forward, 6N to train

Everything starts from one fact about matrix multiplication. A linear layer that maps an in-vector to an out-vector multiplies by a weight matrix of shape (in, out). Each of the out outputs is a dot product of length in: that is in multiplies and in adds, so 2·in floating-point operations. Over all out outputs, one token costs 2·in·out FLOPs — and in·out is exactly the number of parameters in that weight matrix.

one matmul, one token: FLOPs ≈ 2 · in · out = 2 · (params in this matrix)

That "2× the parameters" relationship is not special to one layer — it holds for every weight matrix in the network: QKV projections, the attention output projection, both MLP matrices, the un-embedding. Sum over all of them and the forward pass costs, per token,

forward ≈ Σ_matrices 2 · (params) = 2N FLOPs / token

where N is the total parameter count. (Embedding lookups are gathers, not matmuls, so they cost ≈0 FLOPs; the attention score and value matmuls add an O(T²) term we treat separately — see the caveat at the end of this section.)

Now the backward pass. For each weight matrix the backward computes two things: the gradient with respect to the input (to keep propagating the chain rule down the stack) and the gradient with respect to the weight (to update it). Each is a matmul of the same shape class as the forward, so each costs ≈2N. Backward is therefore ≈4N/token — twice the forward. Add them:

forward

2N FLOPs/token — one matmul per weight matrix, 2·in·out each.

backward

4N FLOPs/token — grad-w.r.t.-input plus grad-w.r.t.-weight, each ≈2N.

train step

6N FLOPs/token — the optimizer update is O(N) and rounds away.

C ≈ 6 · N · D total training FLOPs (D = tokens seen)

This is the budget identity from lesson 0, now derived rather than asserted. A worked instance: GPT-3 has N ≈ 175×10⁹ and trained on D ≈ 300×10⁹ tokens, so C ≈ 6 · 175×10⁹ · 300×10⁹ ≈ 3.15×10²³ FLOPs — matching the widely quoted ~3.14×10²³. A 7B model trained Chinchilla-style on D ≈ 1.4×10¹² tokens (≈200 tokens/param, over-trained for cheap serving) costs 6 · 7×10⁹ · 1.4×10¹² ≈ 5.9×10²² FLOPs. Hold that number; section 5 turns it into days.

When 6N under-counts

The 6N rule prices the dense linear layers and ignores attention's O(T²) score and value matmuls, which add roughly 12 · L · T · d FLOPs/token. For short context this is a rounding error; once T grows comparable to d (long-context training) the attention term becomes a real fraction of the total. Use 6ND for planning; keep the attention term when T ≳ d.

2 · MFU: the FLOPs you paid for vs the FLOPs you got

The 6ND number is FLOPs required. What the hardware delivers is the per-GPU peak (an H100 does ≈990 TFLOP/s of bf16 tensor-core math, ~1 PFLOP/s, with FP8 about double) times however many GPUs, times the fraction you actually achieve. That fraction is Model FLOPs Utilization:

MFU = (achieved model FLOP/s) / (GPUs · peak FLOP/s per GPU)

MFU is below 1 for unavoidable reasons you will spend the systems half attacking: memory-bound ops that can't saturate the tensor cores (lesson 6), kernel-launch and HBM round-trip overhead (lesson 7), communication stalls when work is split across devices (lessons 8–9), and pipeline bubbles. A well-tuned dense Transformer training run lands at roughly 40–55% MFU; 50% is a healthy planning assumption, and anything under ~35% is a signal something is wrong (too-small batch, un-fused kernels, comm not overlapped). MFU is the single honesty knob that turns "FLOPs required" into "wall-clock" — it appears in the runtime estimate in section 5 and is the reason the rest of Part II exists.

3 · The memory ledger: why "16 bytes per parameter"

Now the question that breaks the single GPU. Training state is not "the model"; it is everything AdamW + mixed precision must keep resident for one step. Walk the ledger one entry at a time, per parameter, for the standard bf16 + fp32-master AdamW recipe from lesson 4:

what's resident	precision	bytes / param	why it must exist
bf16 parameters	16-bit	2	the weights the forward/backward actually compute with
fp32 master parameters	32-bit	4	the precise copy the optimizer updates — bf16 has too few mantissa bits to accumulate tiny steps
Adam first moment m	fp32	4	running mean of the gradient
Adam second moment v	fp32	4	running mean of the squared gradient
subtotal (optimizer + weights)		14	the "always resident" core
bf16 gradients	16-bit	2	produced in backward, consumed by the optimizer
total with gradients		≈16	the planning number

So the model's training state costs about 16 bytes per parameter — and some accountings reach ≈18 when fp32 gradients or extra buffers are held. Note how counterintuitive this is: the bf16 weights you think of as "the model" are only 2 of those 16 bytes. The optimizer state and the fp32 master copy are 6× larger than the weights themselves (everything other than the bf16 weights together is 7×). This is the single fact most people get wrong when they reason about whether a model "fits."

training-state bytes ≈ 16 · N (2 bf16 w + 4 fp32 master + 4 m + 4 v + 2 grad)

Run it on 7B, before any activations:

7B · core14 bytes × 7×10⁹ = 98 GB just for weights + fp32 master + Adam m,v.

7B · +grads16 bytes × 7×10⁹ ≈ 112 GB resident training state.

7B · ≈1818 bytes × 7×10⁹ ≈ 126 GB with fp32 grads / buffers — the upper end of the ~112–130 GB range.

one H10080 GB of HBM. 112 GB > 80 GB. The optimizer state alone overflows the device.

A 7B model — small by frontier standards — cannot train on a single 80 GB H100 even if you store zero activations. The naive intuition "7 billion params × 2 bytes = 14 GB, plenty of room" is off by 8×, because it prices only one of the six things that must be resident. This is the memory wall.

4 · Activations — often the real ceiling — and checkpointing

The ledger above is the per-step state. On top of it sits activation memory: every intermediate tensor the forward pass produces and the backward pass needs to recompute its gradients. Activations scale with how much data is in flight, not with N:

activation bytes ∝ B · T · d · L (+ an attention term ∝ B · h · T² when scores are materialized)

The dependence on batch B and sequence length T is what makes activations dangerous: state is fixed once you pick N, but activations balloon the moment you raise B·T for throughput or stretch T for long context. For a large batch on a deep model, activation memory can exceed the 16N training state and become the dominant term — which is why FlashAttention (lesson 7) exists to kill the O(T²) attention activations specifically.

The dial that buys this back is gradient (activation) checkpointing. Normally you store every activation on the way forward so the backward can read them. Instead, store only a handful of checkpoints (typically one per layer boundary) and discard the rest; in the backward, recompute the discarded activations on demand from the nearest checkpoint. You trade memory for compute:

memory

activation memory drops from ∝L down to ∝√L (or to one layer's worth) — often a 5–10× cut.

compute

the forward is effectively run again inside the backward → roughly +30% FLOPs (one extra forward over the 6N step).

verdict

turn it on when activations, not state, are what's pushing you over; it lowers MFU slightly but unlocks larger B·T or longer T.

Checkpointing is the first appearance of the course's central move: spend the cheaper resource to relieve the scarcer one. Here compute is cheap and memory is the wall, so you pay ~30% compute to reclaim memory. The systems half is a catalog of trades exactly like this. Note the limit, though: checkpointing only attacks the activation term — it does nothing for the 112 GB of optimizer state, which is why even with aggressive checkpointing the 7B still does not fit, and why lesson 8 must shard that state across devices.

5 · Putting it together: fit/no-fit and runtime

Combine the two models. Per GPU you need 16N/G bytes if state is sharded across G data-parallel ranks (the ZeRO/FSDP story, lesson 8 — for now, on one GPU, G=1), plus the activation term. A model fits when that total is under 80 GB. Here is the picture across sizes on single 80 GB H100s with no sharding, modest activations:

model N	16N state	fits 1× H100 (80 GB)?	min H100s (state only, ÷G)
1.3B	≈21 GB	yes — room for activations	1
7B	≈112 GB	no — overflows before activations	2
13B	≈208 GB	no	3
70B	≈1,120 GB	no	14+

The 7B "no" is the hinge of the whole course. And runtime falls straight out of section 1 and 2 — total FLOPs over delivered FLOP/s:

time ≈ 6 · N · D / (GPUs · peak FLOP/s · MFU)

Worked, for the 7B at D = 1.4×10¹² tokens (C ≈ 5.9×10²² FLOPs from section 1), on 8 H100s at 50% MFU: denominator = 8 · 990×10¹² · 0.5 ≈ 3.96×10¹⁵ FLOP/s, so time ≈ 5.9×10²² / 3.96×10¹⁵ ≈ 1.49×10⁷ s ≈ 172 days. Halve the MFU and you double the calendar; double the GPUs and you halve it. This one formula is how you turn a dollar budget (GPU-hours) into a feasible model size — the planning loop that lesson 12 formalizes into scaling laws. The deeper accounting of the fused training step and its peak memory lives in gpu_kernels · 30 (the training step & peak memory).

6 · Drive the calculator

Set the model size, batch, context, precision, checkpointing, and GPU count. The stacked meter shows per-GPU bytes split into params / grads / optimizer / activations against the 80 GB H100 line — when the bar would cross it, the bar and the "fits?" KPI turn red. Start at the defaults (7B, 1 GPU) and watch it overflow with the optimizer state alone. Then prove three things to yourself: (a) checkpointing shrinks only the activation segment, never the optimizer one; (b) raising context T inflates activations fast; (c) you need at least 2 GPUs (sharding the state) before 7B even fits, and far more before it trains in reasonable time.

Memory & runtime calculator — does it fit, and how long?

The meter is per-GPU bytes vs the 80 GB H100 wall. It proves the hinge: at 7B on one GPU the optimizer state alone (params+grads+optim ≈ 112 GB) crosses the line before a single activation is added. Checkpointing only trims the activation segment. State is assumed evenly sharded across the GPUs (the lesson-08 preview); raise GPUs to make 7B fit at all.

N (params): B (batch/GPU): 1 T (context): 512 precision: gradient checkpointing #GPUs (G): 1

params grads optimizer (m,v + master) activations

80 GB · 1× H100

Total GB / GPU

—

Fits 80 GB?

—

Train FLOPs (6ND)

—

Est. days · D=1.4T · 50% MFU

—

The widget makes the course's spine visible in one screen: the optimizer-state segment alone is what pushes 7B over the 80 GB line, checkpointing can't touch it, and the only remedy is more devices — which is precisely what forces the systems half.

Failure modes & checklist

Failure modes

Pricing only the weights. "7B × 2 bytes = 14 GB, fits easily." Signal: instant OOM the moment the optimizer allocates m, v, and the fp32 master — you under-counted by 8×.
Forgetting activations exist. State fits, then a big B·T OOMs mid-step. Signal: training survives a few steps at small batch then crashes when you raise it for throughput.
Expecting checkpointing to fix everything. Turning it on and still OOMing. Signal: the optimizer-state segment was the overflow, not activations — checkpointing only trims activations and adds ~30% compute.
Planning runtime off peak FLOP/s. Quoting "it'll take X days" using 100% utilization. Signal: the real run is ~2× longer because actual MFU is ~50%, not 100%.
Ignoring the attention FLOP term at long T. Using 6ND for a 32k-context run. Signal: measured step time far exceeds the 6ND estimate because the O(T²) term is now material.

Checklist

Price the full 16 bytes/param before claiming a model fits — weights are 2 of them.
Add the activation term (∝ B·T·d·L) on top of state, and check it against your real batch and context.
Use 6ND for the FLOP budget, and keep the O(T²) attention term when T approaches d.
Apply MFU (~50%) to peak FLOP/s for any runtime estimate; never quote peak.
Reach for checkpointing only when activations are the bottleneck; reach for sharding (lesson 8) when state is.

Checkpoint

Try it

Open the calculator at 7B, 1 GPU, bf16, T=512, B=1, checkpointing off. Read the total and note that the params+grads+optimizer core alone is ~112 GB — the bar is red and "Fits?" says no, because that core crossed 80 GB before activations were even meaningful. Now (a) toggle checkpointing on and confirm the bar barely moves: which segment shrank, and why didn't it save you? (b) Drag #GPUs up until "Fits?" flips to yes — what's the minimum, and what does the bar tell you each GPU is holding? (c) With that GPU count, read off the estimated days, then halve the assumed MFU in your head and state the new figure. Finally, by hand: a 13B model on bf16 AdamW — how many bytes of resident state, and how many 80 GB GPUs just to hold it?

Where this points next

The accounting just proved that a 7B model — let alone a 70B — cannot train on one 80 GB GPU, because 16 bytes/param of training state overflows the device before activations are even counted. The forced consequence is that the work must be split across many GPUs. But you cannot reason about splitting it well — about MFU, about which ops stall, about the cost of moving bytes between devices — without first understanding the machine you are renting: its memory hierarchy, its tensor cores, and the gap between its TFLOP/s and its TB/s. That is the hardware foundation the rest of accounting silently assumed. Next: 06 · GPUs — the hardware you're renting.

Takeaway

Two pieces of arithmetic decide whether a training run is launchable. Compute: a matmul costs 2·in·out FLOPs, summing to ≈2N/token forward; backward is ≈2× that, so a training step is 6N/token and a whole run is C ≈ 6ND — the budget identity of lesson 0, now derived. Wall-clock is 6ND / (GPUs · peak · MFU), where MFU (~40–55%, plan 50%) is the honest fraction of peak you actually get. Memory: mixed-precision AdamW keeps 2 bytes of bf16 weights, 4 of fp32 master, 4+4 of Adam m,v, and 2 of grads — ≈16 bytes/param, of which the weights are only 2. For 7B that is ≈112 GB of resident state, which overflows one 80 GB H100 before a single activation. Activation memory (∝ B·T·d·L) is often the dominant term on top, and gradient checkpointing buys it back at ~30% extra compute — the course's recurring move of spending the cheap resource to relieve the scarce one. Checkpointing cannot touch the optimizer state, so the only fix for the 112 GB wall is more devices. That "no" is the hinge: it forces the entire systems half of CS336.

Interview prompts

Derive the 6ND training-FLOP rule from a single matmul. (§1 — a (in,out) matmul costs 2·in·out = 2·params per token; summed over all weights the forward is 2N, backward is ≈2× forward (grad-w.r.t.-input + grad-w.r.t.-weight) ≈4N, so a step is 6N/token and a run is 6ND.)
Why is training state ≈16 bytes/param when bf16 weights are only 2? (§3 — add the fp32 master (4), Adam m and v (4+4), and bf16 grads (2): 2+4+4+4+2 ≈ 16. The optimizer state + master dominate; the weights are the smallest piece.)
Walk the memory math for a 7B model on one 80 GB H100. (§3 — 16 × 7×10⁹ ≈ 112 GB of state, before any activations, > 80 GB. It cannot fit; you must shard state across GPUs.)
What is MFU, what's a good value, and why is it below 1? (§2 — achieved ÷ (GPUs·peak) FLOP/s; ~40–55% is healthy; below 1 due to memory-bound ops, kernel/HBM overhead, comm stalls, pipeline bubbles — the targets of Part II.)
What does gradient checkpointing trade, and what can't it fix? (§4 — it discards most activations and recomputes them in backward: ~30% extra compute to cut activation memory (often 5–10×). It does nothing for optimizer state, so it won't make 7B fit one GPU.)
Estimate the wall-clock to train 7B on 1.4T tokens on 8 H100s. (§5 — C ≈ 6·7×10⁹·1.4×10¹² ≈ 5.9×10²² FLOPs; ÷ (8·990×10¹²·0.5 MFU) ≈ 1.5×10⁷ s ≈ ~170 days.)
When does the 6ND rule under-count, and by how much? (§1 caveat — it ignores attention's O(T²) term (≈12·L·T·d/token); negligible for short context, material once T approaches d, e.g. long-context training.)