Unit 4.6 · about 4 h
316 / 324Counting parameters and FLOPs
Every cost claim in later sections multiplies these counts. A wrong parameter or FLOP formula silently moves every budget, schedule and scaling conclusion.
- derive
You can now
- Derive embedding, attention, feed-forward and output parameter counts for MHA, MQA and GQA, with tied versus untied heads.
- Derive forward, backward and training FLOPs from shapes, and state sequence, vocabulary and batch scaling terms.
- Fit a bounded IsoFLOPs curve from offline points and predict the compute-optimal size without running the training API.
Covers
- Embedding, attention, feed-forward and output parameter counts
- Forward, backward and training FLOP approximations
- Sequence length, vocabulary and batch scaling terms
- Bounded IsoFLOPs curve-fitting exercise
Study, in this order
5 steps
- Modern LLM Notebook Appendix B Section 1: parameter countingModern LLM Notebook - evaluation and FLOP accounting
Section 1 (\(D, L, V, H, K, d_h, F, B, S, P\) symbols, GQA \(Q, K, V, O\) shapes, SwiGLU \(3DF\) with \(8D/3\) rationale, per-layer \(12D^2\), global \(2VD\) plus \(12LD^2\), tied versus untied, SmolLM2-135M validation with
sum(p.numel())) - Modern LLM Notebook Appendix B Sections 3 and 4: FLOPs and training ruleModern LLM Notebook - evaluation and FLOP accounting
Section 3 (matmul \(2mkn\), attention \(8BSD^2\) plus \(4BS^2D\), FFN \(6BSDF\) as \(16BSD^2\), global \(2BSD(12LD + 2LS + V)\) with unembedding) and Section 4 (backward 2 times forward, \(6PD\) rule with SmolLM2 135M on 2T and 7B on 1.5T wall-clock)
- CS336 Lecture 2: napkin math, FLOPs and gradientsStanford CS336 - resource accounting, inference and evaluation
motivating_questions()with \(6 \times 70 \times 10^9 \times 15 \times 10^{12}\) at MFU 0.5 and the \(80 \times 10^9 \times 8/12\) bound,tensor_operations_flops()with \(2 \times B \times D \times K\) and MFU,gradients_flops()with forward \(2NP\), backward \(4NP\) and total \(6NP\) - Raschka FLOPS analysis: thop benchmarks and MFULLMs from Scratch companion - GPT FLOP analysis
Fixed-batch forward with thop profile and MACS times 2 on gpt-small to gpt-xl (vocab 50257, context 1024), plus the MFU section with observed over theoretical tokens per second and forward plus backward as 3 times forward
- CS336 Assignment 3: IsoFLOPs method and offline curvesStanford CS336 Assignment 3 - IsoFLOPs scaling
Section 2.1 Scaling Laws from IsoFLOPs profiles with \(C\) about \(6ND\), quadratic loss versus size per budget, optima per profile then power laws \(N_{\text{opt}}\) and \(D_{\text{opt}}\); Problem
chinchilla_isoflopswithdata/isoflops_curves.jsonfieldsparameters,compute_budgetandfinal_loss, fitted offline without the training API
Focus on
- Count each projection separately under GQA: \(Q\) stays \(D\) by \(H\) times \(d_h\), while \(K\) and \(V\) shrink to \(D\) by \(K\) times \(d_h\). MHA \(4D^2\) is the special case \(K\) equals \(H\).
- One matmul \([m, k]\) by \([k, n]\) costs \(2mkn\). Sum \(Q\), \(K\), \(V\), \(QK^\top\), \(AV\) and \(O\) plus gate, up and down plus unembedding \(2BSDV\), and do not drop the \(V\) term at large \(V\).
- Backward costs about twice forward, so training is about 6 times params times tokens for short contexts. Long \(S\) adds the \(2LS\) term.
- IsoFLOPs fixes \(C\) and varies \(N\), finds the loss minimum per budget, then fits \(N_{\text{opt}}\) and \(D_{\text{opt}}\) power laws. Fit the offline points, and do not burn the hosted budget.
Reproduce closed-book
Traps
- Counting MACs as FLOPs. One multiply-accumulate is two FLOPs, so double the thop MACS number.
- Comparing runs by steps at different batch sizes. Convert to tokens first: batch times steps times length.
- Fitting IsoFLOPs on the hosted API when the exercise gives offline curves. The bounded task uses
isoflops_curves.jsononly.
Practice
Modern LLM Notebook Appendix B Exercise 1: LLaMA-2 70B parameters and FP16 memory
Exercise 1 with \(D\)8192 \(L\)80 \(H\)64 \(K\)8 \(V\)32000; compute \(P\) as \(2VD\) plus \(12LD^2\) and FP16 bytes as \(2P\); the cells assert within 1 percent
CS336 Assignment 3: `chinchilla_isoflops` (bounded, offline)
Problem (chinchilla_isoflops) in Section 2.1; fit power laws from data/isoflops_curves.json offline (fields parameters, compute_budget, final_loss) without the training API
Self-test
-
What is the largest AdamW model on 8 H100s at 12 bytes per param?
\(80 \times 10^9 \times 8/12 = 53.3\text{B}\). It is an upper bound because activations are not counted. -
What does \(2VD\) plus \(12LD^2\) give for LLaMA-2 70B?
\(2 \times 32000 \times 8192 + 12 \times 80 \times 8192^2 = 0.52\text{B} + 64.42\text{B}\), about 64.9B. The real model has about 69B: GQA shrinks \(K\) and \(V\), but its SwiGLU width of 28672 is well above \(8D/3\), which more than makes up for it. The \(12LD^2\) shortcut assumes MHA and a \(4D\)-equivalent FFN. -
Why is the training rule \(6NP\)?
Forward is \(2NP\) and backward is \(4NP\) (one matmul for \(dW\) plus one for \(dX\) per op), so training is \(6NP\). Raschka measures forward with thop then triples it for forward plus backward. -
What are the IsoFLOPs fitting steps?
Group runs by \(C\), take the min-loss \(N\) per budget (or a quadratic fit), fit \(N_{\text{opt}}\) and \(D_{\text{opt}}\) power laws in \(C\), then extrapolate to the target budget such as 1e23 or 1e24 in exercise units.
Completion criterion
Stop when
You derive the 64.9B estimate and the 6.3e24 total from shapes alone, and your IsoFLOPs script predicts optimal sizes from the offline JSON.
Mark passed, then open 4.7 Only mark this when you can do it without notes.Leads to
- 4.7 Memory, arithmetic intensity and reconciling estimates Turn FLOP counts into memory budgets and intensity regimes.
- 5.4 KV memory arithmetic Carry KV and sequence terms into cache capacity math.
Comes after
- 2.7 Assembling and verifying a decoderExtend the component count into a full model arithmetic report.
Resources for this unit
- Stanford CS336 - resource accounting, inference and evaluation Use Spring 2026 Lectures 2, 10 and 12 for FLOPs, memory, arithmetic intensity, inference and evaluation. Reproduce the small calculations locally. Assignment 3 IsoFLOPs analysis is optional; do not require its hosted training API or costly experiments.
- LLMs from Scratch companion - GPT FLOP analysis Use ch04/02_performance-analysis/flops-analysis.ipynb as a concrete calculator, then rebuild the formulas from the model configuration and reconcile parameters and at least one memory or operation estimate with runtime state.
- Modern LLM Notebook - evaluation and FLOP accounting Use English notebook 25-evaluation and Appendix B FLOP and memory accounting. Build a 20-example evaluator with two slices, raw outputs and bootstrap intervals, and deliberately create one parser, prompt or judge failure.
- Stanford CS336 Assignment 3 - IsoFLOPs scaling Use Section 2.1 and Problem
chinchilla_isoflopswithdata/isoflops_curves.jsonfor the bounded curve-fitting exercise only. Do not use the hosted training API or costly experiments.
Also in this section
- LLM Foundations Checkpoint - evaluation and arithmetic slice Use the contamination, evaluator-bias, parameter, FLOP, activation, state and KV requirements as the assessment contract. Keep raw outputs and distinguish a verified task result from a broad model capability claim.
- Language Model Evaluation Harness After the small evaluator passes, inspect task versioning, request construction and result schemas to understand production evaluation plumbing. Do not use the harness as a substitute for owning the task and metric contract.