Frontier Pre-Training Lead
The track that teaches you the job Vlad Feinberg does: turn a fixed pile of compute into the best model that can actually be served — forecast it before you spend the money, build it with an architecture the hardware likes, compress it until it fits a real latency budget, and keep the run alive for forty days while it happens.
Start here
| Document | What it gives you |
|---|---|
| The Role | Who Vlad Feinberg is, the exact career path, what a pre-training area lead owns, the synthesized job description, the skill matrix, and his own hiring advice |
| The Transcript, Dissected | Every claim from his interview and his Princeton talk, decoded from zero background — MoE, scaling laws, MFU, distillation, quantization, pipeline prefill, the Flash 2.0 war story — with runnable code for each |
| Lab Standard | What every lab in this track guarantees |
| Glossary | Every term, one line each |
| Cheat Sheet | The numbers and formulas to have memorized |
Then work the phases in order. Phase 00 gates everything — if C = 6ND is not something
you can derive on a napkin, nothing downstream will land.
Why this track exists
Every other track in this hub teaches you to use models — serve them, fine-tune them, chain them into agents. This one teaches you to decide what model should exist in the first place, which is a different discipline with different math and a different failure mode.
The defining constraint: you get one shot. A flagship pre-training run costs eight or nine figures, takes one to three months, and is by construction larger than anything you have ever run. You cannot A/B it. You cannot roll it back. So the entire field is organized around one question — how do you make a defensible prediction about an experiment you have never run? — and everything else (scaling laws, IsoFLOPs ladders, ablations, forecasting) is machinery for answering it.
Then a second constraint lands on top: the model has to be servable. Feinberg's own framing is that Chinchilla-optimal is the wrong objective for a model that will serve trillions of tokens, because model size is a serving cost paid forever. So the objective becomes quality per served token per watt — and that pulls distillation, quantization, MoE, sharding, and roofline analysis into what would otherwise be a pure modelling job.
Nobody teaches this. There is no course. The knowledge lives in about a dozen papers, a handful of talks, and the heads of maybe a few hundred people. This track assembles it.
The roadmap
BUILT ─────────────────────────────────────────────────────────┐
┌─────────────────────────────────────────┐ │
FOUNDATIONS │ 00 FLOPs & memory algebra C = 6ND │ │
│ 01 Scaling laws Kaplan → Chinchilla │ │
└────────────────┬────────────────────────┘ │
│ │
┌────────────────▼────────────────────────┐ │
ARCHITECTURE │ 02 Mixture of Experts from scratch │ │
& HARDWARE │ 03 Roofline, MFU & inference co-design │ │
└────────────────┬────────────────────────┘ │
───────────────────────────────── │ ───────────────────────────┘
PLANNED │
┌────────────────▼────────────────────────┐
SYSTEMS │ Sharding & parallelism DP/TP/PP/EP │
│ Pipeline prefill (the Flash 2.0 fix) │
└────────────────┬────────────────────────┘
│
┌────────────────▼────────────────────────┐
COMPRESSION │ Scaling laws II U/R, inference-aware │
│ Distillation at trillion-token scale │
│ Quantization & power economics │
└────────────────┬────────────────────────┘
│
┌────────────────▼────────────────────────┐
OPERATIONS │ Training stability & the 40-day run │
& CRAFT │ Post-training: SFT, RLHF, PPO │
│ Kernels & tile DSLs │
│ Research as an MDP (taste) │
└────────────────┬────────────────────────┘
│
┌────────────────▼────────────────────────┐
CAPSTONE │ Mini frontier lab: budget → recipe → │
│ forecast → compress → serve → the memo │
└─────────────────────────────────────────┘
Phase index
Built and test-verified
| # | Phase | What you build | Maps to |
|---|---|---|---|
| 00 | FLOPs & Memory Algebra | A transformer FLOP/memory/time calculator; the 6ND derivation; MoE active-vs-total; KV-cache sizing; budget → (N, D) → dollars | "C = 6ND is a very good approximation" · 75 tests |
| 01 | Scaling Laws I — Kaplan → Chinchilla | IsoFLOPs ladder, parabola/power-law/parametric fitting, bootstrap forecasts, optimal ladder design, the baseline-vs-candidate crossover | The six-step IsoFLOPs slide; "UNDERTRAINED!" · 64 tests |
| 02 | Mixture of Experts From Scratch | Router, top-k gating, load-balance & z-loss, capacity/drop/pad, shared experts, total-vs-active, the all-to-all wall, a router-collapse simulator | "MoE… uses a lot more parameters" · 53 tests |
| 03 | Roofline, MFU & Inference Co-Design | Roofline with ridge points, MFU/HFU accounting and budget decomposition, prefill/decode latency models, the chip-count solver, GQA and tile levers | "shapes… that fully saturate all of those hardware units" · 50 tests |
Roadmap — the remaining phases
These are specified but not yet built. Each will follow the same contract: README.md +
WARMUP.md + a runnable, test-verified lab. The material for all of them is already covered
conceptually — with code — in The Transcript, Dissected, which is
the fastest way to get the ideas now.
| Phase | What it will build | Covered today in |
|---|---|---|
Scaling Laws II — beyond L(N, D) | L(N, U, R) data-constrained optimizer, inference-aware lifetime cost, routed & distillation laws | Claims 15–17 |
| Sharding & Parallelism | Collective cost models, DP/FSDP/TP/PP/EP planner, the 4×4 mesh question | Claim 12 |
| Pipeline Prefill & Serving | All-to-all vs pipelined-prefill simulator, bubble analysis, prefill/decode disaggregation | Claim 12 |
| Distillation at Scale | KL/temperature losses, top-k teacher-logit store, teacher-compute allocation | Claims 3, 17 |
| Quantization & Power Economics | Affine/group/vector quantizers, outlier handling, the energy & TCO model | Claim 5 |
| Training Stability & SRE | Goodput model, spike detector, checkpoint-policy optimizer | Claim 13 |
| Post-Training: SFT, RLHF, PPO | SFT masking, Bradley-Terry reward model, PPO step with KL control | Claim 1 |
| Kernels & Tile DSLs | A tiny tile-DSL interpreter, online softmax, memory-traffic accounting | Claims 1, 19 |
| Research as an MDP | Portfolio planner, value of information, calibration scoring | Claim 2 |
| Capstone — Mini Frontier Lab | Budget → recipe → forecast → compress → serve → the go/no-go memo | the whole track |
What every phase contains
| File | Voice | What it gives you |
|---|---|---|
README.md | the syllabus | why the phase exists, the concept map, the lab spec, deliverables, key takeaways |
WARMUP.md | the professor | zero background → principal level. Every term built from nothing: what it is → why it exists → how it works underneath (mechanism, diagrams, math, runnable code) → what it costs in production → the common misconceptions. Then a lab walkthrough, interview Q&A, tips & takeaways, and primary-source references |
lab-*/ | the workbench | a runnable, test-verified miniature: README.md, lab.py (TODOs), solution.py, test_lab.py, requirements.txt |
Labs are pure Python stdlib, offline, deterministic, free. No GPU, no network, no model downloads. You implement the mechanism — the FLOP counter, the parabola fit, the router, the collective cost model, the pipeline scheduler, the quantizer, the PPO step — because that is what makes the knowledge defensible in an interview and useful at 2 a.m. See LAB-STANDARD.md.
Suggested schedule
Assumes ~10 focused hours a week. Compress aggressively if you already have the background.
| Weeks | Work | Milestone |
|---|---|---|
| 1 | Read jd.md + TRANSCRIPT-DISSECTED.md | You can explain MoE, MFU, 6ND, distillation and the Flash 2.0 fix to a colleague |
| 2 | Phase 00 | You can size any model on any cluster on a napkin |
| 3–4 | Phase 01 | You can fit a scaling law and defend a flagship forecast with error bars |
| 5–6 | Phase 02 | You can implement an MoE layer and debug a collapsing router |
| 7–8 | Phase 03 | You can do the latency napkin and defend an architecture shape |
| 9+ | The four artifacts below | The portfolio |
The parallel track: build the artifacts he asks for
While working the phases, produce the four things he explicitly says he screens on (see jd.md §8):
- Handwritten How To Scale Your Model exercises, on video.
- A transformer from scratch — with a FLOP counter and a memory profile, not just a forward pass.
- A real, benchmarked contribution to vLLM / SGLang / TensorRT.
- A citation-tree writeup for one topic, ending with the paper that contradicts the consensus.
Phases 00, 05, and 11 feed directly into #1 and #2. Phase 12 is #4 as a method.
Prerequisites
Required: Python (comfortable), and the ability to sit with algebra. That is genuinely it — every other concept is built from scratch in the WARMUPs.
Helpful but not assumed: having seen a transformer before; having trained any model; basic probability. If you have none of these, read TRANSCRIPT-DISSECTED §Part 0 first — it defines the nine words the rest depends on.
Explicitly not required: a GPU, a cloud account, a PhD, or a job at a lab. Feinberg himself did not finish the PhD, and his slide "Future Pretrain Research Ideas – Without Big Costs!" is a list of open problems you can work on with a laptop.
Related tracks in this hub
- Senior AI Engineer — the from-scratch implementation track: tokenizer, transformer, autograd, LoRA, serving internals. Best companion; this track assumes the mechanism and focuses on the decision.
- LLM Inference Engineer — deeper on serving stacks.
- Model Accuracy & AI Performance — deeper on quantization tooling and compilers.
- GPU Engineering — deeper on CUDA and the hardware itself.
Key takeaways for the whole track
- You get one shot. Everything upstream of the flagship run exists to make one irreversible decision defensible.
- A scaling law describes your recipe, not the universe. Change the recipe, refit the law.
- Loss forecasting is recipe selection. Compare fitted curves at the target FLOP count.
- Optimize lifetime cost, not training cost — while knowing the inference term is unforecastable.
- Nothing is free; costs move. MoE trades compute for memory, communication, instability, and data hunger.
- Prefill and decode are different machines. Compute-bound vs bandwidth-bound.
- Power is the bill. Quantization is an energy lever before it is a memory lever.
- Distillation is variance reduction, and its hard part is storage.
- Goodput is a research deliverable. 78% vs 94% on the same hardware is days of a 40-day run.
- The artifact is the signal. Build things other people can check.
The Role — Frontier Pre-Training Lead
Modelled on: Vlad Feinberg — Distinguished Engineer, Google DeepMind; area lead for Gemini pre-training (Flash / Flash-Lite). Sources: his April 2025 Princeton talk "Gemini Pretraining" (slides), his personal site, and his interview on Developing Dev — "Google DeepMind Pre-Training Lead: How To Get a Job at a Frontier Lab" (video).
This document is not a copy of a job posting. It is a reconstruction of what the job actually is — the scope, the decisions, the failure modes, the artifacts — assembled from what he has said publicly about doing it. Everything in the rest of this track exists to make you able to do the things listed here.
Table of Contents
- 1. One-Paragraph Summary of the Job
- 2. Career Timeline — What the Path Actually Looked Like
- 3. What "Pre-Training Area Lead" Owns
- 4. The Three Research Verticals
- 5. The Products That Depend On You
- 6. Synthesized Job Description
- 7. The Skill Matrix, Graded
- 8. How He Says To Get Hired
- 9. What This Track Does About It
1. One-Paragraph Summary of the Job
You decide how to convert a fixed pile of compute into the best possible model — and "best" is not test loss, it is quality per served token per watt. You own the recipe (architecture, data mixture, optimizer, schedule), the forecast (a scaling law that says what loss you will land on before you spend $10M of TPU time), the compression pipeline that turns a giant teacher into a small servable student, and — when the run is live — the pager. You are simultaneously a researcher (navigating ideas that may not work), a distributed-systems engineer (a training job across many datacenters), an economist (power is the bill), and an SRE (the run must not die). Feinberg's framing of the two halves is exact: software engineering is a deterministic DAG where you "just make monotone progress"; research is a stochastic one, where "some of the nodes... may or may not work out."
2. Career Timeline — What the Path Actually Looked Like
Taken from the "About Me" slide of his own talk. Study the shape, not the brand names — the shape is the lesson.
| When | Where | What he actually did |
|---|---|---|
| →2017 | Princeton, COS + SML (CS + Statistics & Machine Learning) | 3D CNN for MRI segmentation, for connectome reconstruction |
| 2017–2018 | UC Berkeley RISE Lab, ML Systems PhD (with Ion Stoica, Joey Gonzalez, Mike Franklin) | Model-based deep RL on MuJoCo (with Sergey Levine) |
| 2018–2019 | Dropped out | Left the PhD for a startup |
| 2019–2021 | Sisu Data (with Peter Bailis), Head of ML | Efficient database cubing with FDR control, via a custom distributed lasso engine |
| 2021–2022 | Google Cerebra | Quantizing the Ads DNN for pCTR serving efficiency |
| 2022→ | Google Brain → Google DeepMind | Optimizer work (with Elad Hazan) → inference-efficient LLMs → Flash pre-training lead → Distinguished Engineer, area lead for Gemini pre-training |
Four things to extract from that table
- The through-line is "make the expensive thing cheap." Distributed lasso for DB cubing → quantizing an ads DNN → inference-efficient LLMs → Flash. That is one skill, re-applied at four scales. It is not a career of unrelated hops.
- The unglamorous job was the on-ramp. He describes doing SFT for an early version of Bard — "hyperparameter tuning and engineering work to get this model running on some really old TPUs" — and getting a spot bonus from Jeff Dean for it. He had been in pure research optimizing for first-author papers at NeurIPS/ICML/ICLR; his manager Rohan Anil pushed him toward the product-shaped work. That pivot is the whole story.
- He did not finish the PhD. The credential is not the gate. What the PhD teaches — research taste, literature traversal, mathematical maturity — is the gate. This track teaches those directly.
- Ads pCTR quantization is the unsexy ancestor of Flash. His own advice: "chase the problems that people are facing in the world today... don't be afraid to tackle a smaller part of this problem or maybe a more menial sounding part."
3. What "Pre-Training Area Lead" Owns
Think of it as five distinct hats worn by the same person in the same week.
Hat 1 — The recipe owner
A recipe is the full parameterized specification of a training run: architecture family
(dense vs MoE, depth/width ratio, attention shape), tokenizer, data mixture and its
curriculum, optimizer and its hyperparameters, learning-rate schedule, batch-size schedule,
numerics (bf16/fp8, what stays in fp32), and the parallelism plan. You do not tune these
per-run — you tune the function that maps compute budget C to all of the above, because
you only get to run the big one once.
Hat 2 — The forecaster
Before the flagship run, you must be able to state the final test loss with error bars. His framing of why this is the problem:
Old ML: iterate on CIFAR-10, then apply to ImageNet. LR searches by doing multiple "final runs" — "the last data point is our test set!"
Now: "every single time you go up for a pre-training run, you're about to put in more FLOPs into this run than you've ever done before." Every next run requires extrapolation.
So you run a ladder of small models, fit a law, and extrapolate. Phase 01 and Phase 02 are this, end to end.
Hat 3 — The inference co-designer
Chinchilla-optimal is the wrong objective for a model that will serve billions of tokens. You trade training FLOPs against serving FLOPs, and you choose architecture shapes that saturate the specific hardware units you will serve on. His talk's napkin math is the canonical example (Phase 05 rebuilds it exactly).
Hat 4 — The compressor
Distillation, quantization, and serving-friendly architecture changes. He calls the fact that you can drop weights from FP32 down to 4 bits and keep quality "kind of a miracle" — and the reason it matters is economic: "99% of the total cost of operation for AI hardware comes from the power that it takes to run these chips."
Hat 5 — The SRE
When the flagship run is live, someone has to keep it alive. For Flash 2.0 that was 40 days of a five-person rotation, handing off day by day across Paris and Mountain View, watching data iterators, chasing indexing bugs, avoiding "wasting all of this GPU time." He is blunt: "We did not do a lot of sleeping."
4. The Three Research Verticals
His team's stated research surface. Each one gets its own phase in this track.
| Vertical | The one-line problem | Why it is hard | Phase |
|---|---|---|---|
| Distillation | Push "statistics about the underlying dataset" from a teacher into a student | The teacher must be run over "trillions and trillions of tokens" — "millions and millions of dollars"; every op is "multiplied by such a large factor". Forces storage and multi-datacenter engineering. | Phase 07 |
| Inference co-design | Pick "shapes of the matrices", "attention shapes, num heads" so the net "fully saturates" every hardware unit | Quality and MFU pull in opposite directions; the right answer differs per chip generation. | Phase 05 |
| Quantization | Shrink the bits used to represent weights (and then activations) | Quality cliffs are non-uniform; activation quantization multiplies the win but is far more fragile. Power is the bill. | Phase 08 |
5. The Products That Depend On You
From his talk's "Small Model Customers" section — the reason Flash and Flash-Lite exist is that Google's biggest surfaces need (1) high volume and (2) real time:
- Free-tier Gemini app (the chatbot)
- AIO — AI Overviews in the search bar
- AIM — AI Mode in the search bar
- Vertex AI — customer fine-tuning and deployment
- AI Studio — the generation API
- Astra (real-time multimodal assistant) and Mariner (web-interaction agent)
- Plus models for Ads and YouTube, and technical work on the Google–Apple partnership
And the internal cadence he names: the Gemini "tick-tock" — this generation's Flash should match last generation's Pro. That single sentence is a compression target, a distillation target, and a scaling-law target all at once.
6. Synthesized Job Description
What the posting would say if it were honest.
Title: Research Engineer / Research Scientist — Pre-Training (Frontier LLM)
You will:
- Own a parameterized pre-training recipe: architecture, data mixture, optimizer, schedule, numerics, and parallelism plan, as a function of the compute budget.
- Design and run scaling-law ladders (IsoFLOPs and/or parametric
L(N, D)fits) and publish forecasts with uncertainty for flagship runs, including the decision rule for baseline-vs-candidate recipe comparison. - Extend scaling laws along non-classical axes: unique tokens and repeat count
L(N, U, R), routed/MoE capacity, inference-aware total-cost objectives, and distillation (teacher-size / teacher-compute allocation). - Co-design architectures for the serving target: choose
d_model,d_ff, head count/shape, KV-head sharing, and expert count against a roofline model of the target accelerator; defend the choice with MFU and latency arithmetic, not vibes. - Build and operate the compression pipeline: teacher-logit generation at trillion-token scale (storage format, multi-datacenter data movement), distillation objectives, and post-training / quantization-aware paths down to 4-bit.
- Partner with serving to make architectures deployable — expert parallelism, pipelined prefill, prefill/decode disaggregation, KV-cache budgets.
- Keep flagship runs alive: goodput, checkpoint/restart, loss-spike diagnosis and mitigation, data-iterator correctness, silent-data-corruption detection, on-call rotation.
- Make and defend go/no-go calls under uncertainty; run "a very transparent technical process" when a large bet (e.g. dense → MoE) is on the table.
You must be able to:
- Derive
C ≈ 6NDand the full per-step FLOP count of a transformer from the layer shapes, by hand, and say exactly where the approximation breaks (attention, MoE, embeddings). - Read the primary literature and traverse a citation tree to find "the high-value papers."
- Implement a technique from a paper yourself, efficiently — including at the kernel level.
- Reason about collectives (all-reduce/all-gather/reduce-scatter/all-to-all) and their cost in bytes and in latency, and about how a sharding decision changes convergence.
- Explain why an MFU number is what it is, unit by unit.
Signals that get you the interview (his words, verbatim in spirit): "intent, mathematical maturity, grit."
7. The Skill Matrix, Graded
Grade yourself 0–3 (0 = never heard of it, 3 = I have shipped this). The right column is where this track fixes the gap.
| # | Skill | Phase |
|---|---|---|
| 1 | Derive transformer FLOPs and memory from shapes; 6ND; KV-cache sizing | 00 |
| 2 | Fit and extrapolate a scaling law; IsoFLOPs; know why Kaplan and Chinchilla disagreed | 01 |
| 3 | Data-constrained (L(N,U,R)), inference-aware, routed and distillation scaling laws | 02 |
| 4 | Implement an MoE layer: router, top-k, auxiliary losses, capacity factor, drop/pad | 03 |
| 5 | DP / FSDP / TP / PP / EP; collective cost models; GSPMD-style sharding annotations | 04 |
| 6 | Roofline & MFU; arithmetic intensity; prefill vs decode; the latency napkin math | 05 |
| 7 | Pipelined prefill for MoE; disaggregated serving; hiding comms behind compute | 06 |
| 8 | Distillation objectives; teacher-logit storage at scale; capacity gap; variance reduction | 07 |
| 9 | Quantization from FP32 → 4-bit; activation quant; power/TCO arithmetic | 08 |
| 10 | Training stability: loss spikes, goodput, checkpointing, iterator bugs, SDC | 09 |
| 11 | SFT and RLHF/PPO mechanics well enough to debug them on old hardware | 10 |
| 12 | Kernel/tile-DSL thinking: why ThunderKittens-style abstractions exist | 11 |
| 13 | Research as an MDP: value-of-information, kill criteria, portfolio planning | 12 |
| 14 | Put it together: budget → recipe → forecast → run → compress → serve → memo | 13 |
8. How He Says To Get Hired
He is unusually concrete about this, so it is worth listing exactly.
- Produce "actual evidence that you've created something of use to other people." Contribute to vLLM, SGLang, TensorRT — real improvements to real serving stacks. He calls this "an extremely positive signal."
- Do the exercises in How To Scale Your Model ("The Scaling Book") — handwritten, detailed — and record yourself doing them. He publicly offered interviews (and referrals when he had no headcount) to people who did this plus the transformer exercise.
- Implement a transformer from scratch. It demonstrates "willingness to get into the weeds engineering-wise" and fluency in the "bread-and-butter math that we use every day to size these LLMs."
- Be able to traverse a citation tree. Knowing which papers are the high-value ones for a topic is itself the skill.
- If you are already inside a big company: don't try to teleport. Become the person who owns LLM adoption for your product area. You then become "the partner that we work with on the research side" — and the transfer becomes a formality. (He names Nate Lintz, who transferred in and now owns much of their inference work.)
- Be someone people want to see succeed. His anti-Machiavelli argument: build projects that leverage other people's complementary skills so they shine, and they will show up for your next project. He credits Todd Lipkin with first getting him into CS.
On the doom: he calls the discourse "FUD everywhere." His counter is accountability — "you can't hand off blame to AI." A lawyer can't be replaced by a model that "can't be disbarred." Someone must be responsible for the output and for how resources get allocated. His practical version: "we all have agency over our future and we can start investing in skills that matter for tomorrow today."
9. What This Track Does About It
Every phase is: a README.md (why the phase exists + lab spec), a WARMUP.md
(zero-background → principal-level explanation of every term, with mechanism, math, code,
diagrams, misconceptions, tips and takeaways), and at least one runnable lab — pure
Python stdlib, offline, deterministic, test-verified — that builds the mechanism itself
rather than calling a framework.
You will finish having personally implemented: the FLOP counter, the IsoFLOPs fitter, the
L(N,U,R) and inference-aware optimizers, a top-k MoE router with capacity and auxiliary
losses, a collective-cost simulator, a roofline/MFU model, a pipelined-prefill scheduler, a
distillation loss with temperature and a teacher-logit store, a 4-bit quantizer with a power
model, a training-stability watchdog, a PPO step, a tile-DSL interpreter, an MDP research
planner, and a capstone that chains all of it into one defensible technical memo.
Start here: Transcript, Dissected — every claim he made, unpacked for someone with no background — then Phase 00.
References
- Vlad Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides (PDF)
- Vlad Feinberg, personal site and about
- Ryan Peterman, Developing Dev: Google DeepMind Pre-Training Lead: How To Get a Job at a Frontier Lab · video
- Kaplan et al., Scaling Laws for Neural Language Models, 2020 — https://arxiv.org/abs/2001.08361
- Hoffmann et al., Training Compute-Optimal Large Language Models (Chinchilla), 2022 — https://arxiv.org/abs/2203.15556
- Sardana et al., Beyond Chinchilla-Optimal, 2024 — https://arxiv.org/abs/2401.00448
- Muennighoff et al., Scaling Data-Constrained Language Models, 2023 — https://arxiv.org/abs/2305.16264
- Busbridge et al., Distillation Scaling Laws, 2025 — https://arxiv.org/abs/2502.08606
- Clark et al., Unified Scaling Laws for Routed Language Models, 2022 — https://arxiv.org/abs/2202.01169
- Pope et al., Efficiently Scaling Transformer Inference, 2022 — https://arxiv.org/abs/2211.05102
- Austin et al., How To Scale Your Model — https://jax-ml.github.io/scaling-book/
- Xu et al., GSPMD, 2021 — https://arxiv.org/abs/2105.04663
- DeepSeek-AI et al., DeepSeek-V3 Technical Report, 2024 — https://arxiv.org/abs/2412.19437
- Jacob Steinhardt, Research as a Stochastic Decision Process — https://cs.stanford.edu/~jsteinhardt/ResearchasaStochasticDecisionProcess.html
The Transcript, Dissected
What this document is. Vlad Feinberg's interview and his Princeton talk are dense with claims that are obvious to a frontier-lab engineer and opaque to everyone else. This document takes every substantive claim, in order, and does four things with each:
- Quote it (briefly) or state it plainly.
- Decode it for someone with zero background — every term defined from nothing.
- Go under the hood — the actual mechanism, with math and runnable code.
- Zoom out — why it matters, what it implies, what the common misreading is.
Nothing here assumes you know what a transformer is, what a FLOP is, or what "sharding" means. If you already do, the Under the hood subsections are where the value is.
Sources: the Developing Dev interview (video) and his Gemini Pretraining talk, Princeton, Apr 2025. Quotes are short and attributed; everything else is this track's own explanation.
Table of Contents
- Part 0 — The Vocabulary You Need Before Anything Else
- Part 1 — "Pre-Training" Is Not "Training"
- Claim 1 — Kernel and low-level work is in "voracious demand"
- Claim 2 — Research is a stochastic DAG; engineering is a deterministic one
- Claim 3 — Vertical one: distillation
- Claim 4 — Vertical two: inference co-design
- Claim 5 — Vertical three: quantization, and the 99%-power fact
- Claim 6 — MFU, and why 100% is not the goal
- Claim 7 — Pre-training is a one-shot extrapolation problem
- Claim 8 —
C = 6ND, derived - Claim 9 — Kaplan said scale parameters; Chinchilla said scale both
- Claim 10 — Scaling laws are brittle and dataset-dependent
- Claim 11 — MoE: what it is, why it wins, why it hurts
- Claim 12 — The MoE serving bottleneck, and Geng Yan's pipeline-prefill fix
- Claim 13 — The Flash 2.0 war story: 40 days, five people, two continents
- Claim 14 — Real-time products force small models (the napkin math)
- Claim 15 — Chinchilla ignores inference cost
- Claim 16 — "We're running out of internet"
- Claim 17 — Distillation scaling laws and the capacity gap
- Claim 18 — The Gemini tick-tock
- Claim 19 — You can do pre-training research without a supercomputer
- Claim 20 — Hiring: intent, mathematical maturity, grit
- Claim 21 — The internal-transfer play
- Claim 22 — On AI replacing engineers
- Claim 23 — The career philosophy
- Master Takeaways
- References
Part 0 — The Vocabulary You Need Before Anything Else
If you have no background, read this once. Nine words unlock the rest of the document.
FLOP — one floating-point operation: a single multiply or a single add on decimal numbers. "FLOPs" (plural) counts work; "FLOP/s" measures speed. An NVIDIA H100 does roughly 10¹⁵ FLOP/s (about a petaflop) on the number formats used for training. If a task needs 10²⁴ FLOPs and your machine does 10¹⁵ FLOP/s, it takes 10⁹ seconds ≈ 31 years — which is why you use 10,000 machines.
Parameter — one number inside the model that gets learned. N is the count. A "70B
model" has 70 × 10⁹ parameters. At 2 bytes each (bf16) that is 140 GB just to hold the
weights.
Token — a chunk of text (roughly ¾ of an English word). Models read and write tokens, not
characters. D is the number of tokens the model is trained on. Modern runs use 10¹³–10¹⁴
tokens ("trillions and trillions").
Matmul (matrix multiply) — the operation that is ~99% of a transformer's arithmetic.
Multiplying an (m × k) matrix by a (k × n) matrix costs 2·m·k·n FLOPs (each of the
m·n outputs is a sum of k products, and each product-plus-add is 2 FLOPs).
HBM (High Bandwidth Memory) — the fast memory physically attached to an accelerator chip. An H100 has 80 GB; a TPU v5e has 16 GB. This is the scarcest resource in serving. Two numbers matter: capacity (GB) and bandwidth (TB/s — how fast you can read it).
Accelerator / chip — a GPU (NVIDIA H100, B200) or a TPU (Google's v5e, v6e). It contains a matmul unit (systolic array / tensor cores) that is enormously fast, a vector unit that is much slower, and HBM that is slower still relative to the matmul unit's appetite.
Pre-training — the first and most expensive stage: predict the next token over a huge corpus. This produces a "base model" that knows language and facts but is not yet a helpful assistant.
Post-training — everything after: supervised fine-tuning (SFT) on demonstrations, preference optimization (RLHF/DPO), reasoning training. Turns a base model into a product.
Inference / serving — running the trained model to answer requests. Split into prefill (process the user's whole prompt at once — compute-heavy) and decode (generate output one token at a time — memory-bandwidth-heavy). This split governs everything in Phases 05–06.
Scaling law — an empirical formula predicting model quality (test loss) from compute, parameters, and data. The tool that makes a $50M training run a calculated bet instead of a prayer.
Tip. Write these nine on a card. Every claim below is a combination of them.
Part 1 — "Pre-Training" Is Not "Training"
Before the claims, the shape of the job, because the interview assumes it.
A normal ML job looks like: get data → train → check validation → adjust → retrain → ship. The loop runs dozens of times a week.
A frontier pre-training job looks like this:
┌──────────────────────── months ────────────────────────┐
recipe design → scaling-law ladder → forecast → GO/NO-GO → the run → compress → serve
(weeks) (many small runs) (a number) (one meeting) (40 days) (weeks) (forever)
▲
this happens exactly ONCE
Everything upstream of "the run" exists to make one irreversible decision defensible. That is why scaling laws are the intellectual centre of the job, why forecasting is a research area, and why the SRE hat exists. Hold that picture; every claim below is a piece of it.
Claim 1 — Kernel and low-level work is in "voracious demand"
He says there is "voracious demand... across all the different labs" for "kernel development and low-level engineering to improve the runtime for these LLMs," because when you change the architecture "you just need to be able to implement these new techniques in efficient ways."
Decode it
A kernel is a small program that runs on the accelerator and does one job extremely fast —
"multiply these two matrices," "apply softmax to this row," "do attention for this block." Your
Python code (torch.nn.Linear) is a thin wrapper; the actual work happens in a kernel written
in CUDA (NVIDIA), Triton, Pallas (TPU), or an assembly-adjacent language.
Why the demand exists. A researcher proposes a new attention variant on a whiteboard. To test it, someone must write a kernel for it. If no kernel exists, the idea runs 10× slower than the baseline in naive PyTorch, loses the comparison, and gets discarded — even if it was better. Missing kernels silently kill good research. That is the demand.
Under the hood — why naive code is slow
The classic example is attention. Naive attention materializes an n × n score matrix:
# Naive attention. n = sequence length, d = head dimension.
# The killer: `scores` is n×n and must be WRITTEN to HBM, then READ back. Twice.
scores = Q @ K.T / math.sqrt(d) # (n, n) <- n² floats hit memory
weights = softmax(scores, axis=-1) # (n, n) <- read n², write n²
out = weights @ V # (n, d) <- read n² again
At n = 8192, n² = 67M floats = 134 MB in bf16 — per head, per layer. You blow past cache
and pay HBM bandwidth three times.
FlashAttention never materializes it. It walks over blocks of keys/values, keeping a running softmax in fast on-chip memory using the online-softmax trick:
# Online softmax: fold a new block into a running (max, sum, weighted-output) triple.
# This is the whole idea behind FlashAttention, in 8 lines of pure Python.
def online_softmax_update(m_prev, l_prev, o_prev, new_scores, new_values):
m_new = max(m_prev, max(new_scores)) # running max, for stability
rescale = math.exp(m_prev - m_new) # how much to shrink the old state
exp_new = [math.exp(s - m_new) for s in new_scores]
l_new = l_prev * rescale + sum(exp_new) # running denominator
o_new = [o * rescale for o in o_prev] # rescale old numerator
for e, v in zip(exp_new, new_values): # add the new block's contribution
o_new = [a + e * b for a, b in zip(o_new, v)]
return m_new, l_new, o_new
# Final output = o_new / l_new. Never stores an n×n matrix.
Same mathematical result. Memory traffic drops from O(n²) to O(n). Speedups of 2–4× on
long sequences, and it is why 100k+ context is affordable at all.
Takeaway. "Kernel work" is not plumbing beneath the science. It is the science: the set of architectures you can honestly evaluate is exactly the set someone has written a fast kernel for. Phase 11 builds a tiny tile-DSL so this stops being magic.
Claim 2 — Research is a stochastic DAG; engineering is a deterministic one
He frames "research as an MDP" (crediting Jacob Steinhardt's essay). Engineering: "the DAG is more or less deterministic... you can just make monotone progress." Research: "some of the nodes which might be research ideas... may or may not work out." The core skill is "building an intuition of how likely an approach is to work out without having yet done that approach" — research taste.
Decode it
A DAG (directed acyclic graph) is a to-do list with arrows: "must do A before B." Building a web service is a DAG — write the database layer, then the API, then the frontend. Each box you finish stays finished. Progress is monotone: it only goes up.
An MDP (Markov Decision Process) is the mathematical model for "make decisions when outcomes are random." It has states (what you know now), actions (which experiment to run), transition probabilities (this experiment works with probability p), rewards, and costs. The famous fact about MDPs: the greedy choice — highest immediate reward — is often wrong, because a cheap experiment that teaches you which branch to take can be worth more than an expensive one that might directly succeed.
Feinberg's point is that research planning is the MDP, not the DAG. You must "factor in the success rate and the time investment... as well as a priori estimating what those different rates are."
Under the hood — make "taste" a computation
Taste feels mystical. It is not: it is a prior over p(success) plus a habit of computing
expected value per unit cost. Here it is in code you can actually run:
# A research portfolio, scored the way a pre-training lead scores it.
# value = what you gain if it works (e.g. % quality improvement on the flagship)
# p = your prior that it works
# cost = engineer-weeks (or TPU-days) to find out
experiments = [
# name, value, p, cost
("new attention variant", 100, 0.15, 8),
("data mixture reweighting", 30, 0.70, 2),
("bigger LR + warmup tweak", 15, 0.50, 1),
("switch dense -> MoE", 250, 0.40, 20),
]
for name, value, p, cost in sorted(
experiments, key=lambda e: -(e[1] * e[2] / e[3])):
print(f"{name:30s} EV/cost = {value * p / cost:6.2f}")
data mixture reweighting EV/cost = 10.50
bigger LR + warmup tweak EV/cost = 7.50
switch dense -> MoE EV/cost = 5.00
new attention variant EV/cost = 1.88
Now the three refinements that separate a senior researcher from this toy:
- Value of information. An experiment whose result changes what you do next is worth more than its direct payoff. Running a cheap 400M-parameter MoE ablation has low direct value but tells you whether to spend the 20 weeks. Buy information before you buy outcomes.
- Kill criteria set in advance. "If the 1B ablation is not within 0.01 nats of baseline by 20B tokens, we stop." Written before you start, because after you start you will be attached to it. This is the single highest-leverage habit in research.
- Correlated failures. Five variants of one idea are not five independent bets. If the underlying premise is wrong, all five die together. Diversify across premises, not across implementations.
Takeaway. When someone says a researcher has "great taste," they mean the person's internal
pvalues are well-calibrated and they instinctively buy information first. Both are trainable. Phase 12 makes you build the planner.
Claim 3 — Vertical one: distillation
Transferring "knowledge or some form of statistics about the underlying dataset through a teacher model into the student model." At scale this means running the teacher over "trillions and trillions of tokens," costing "millions and millions of dollars," so every operation matters "because every operation... is multiplied by such a large factor."
Decode it
You have a huge, excellent, expensive model (the teacher — say Gemini Pro). You want a small, fast, cheap model (the student — Flash) that behaves as much like it as possible.
The naive approach: train the student on the same text the teacher saw. The distillation approach: train the student to match the teacher's probability distribution over the next token.
Why that is dramatically better — this is the key insight and it is easy to miss:
Suppose the context is "The capital of France is". The ground-truth label is a single token:
" Paris". That is one bit of supervision — "this one, not the other 255,999."
The teacher instead says:
" Paris" 0.92
" Lyon" 0.03
" Marseille" 0.02
" located" 0.01
" the" 0.005
... (a full distribution over the whole vocabulary)
That is a rich, dense signal. It doesn't just say the answer — it says which wrong answers are reasonable. It encodes that Lyon is a French city while "banana" is not. Each token now carries hundreds of bits of supervision instead of one. This is the "some form of statistics about the underlying dataset" phrase: the teacher's distribution is a compressed summary of everything it learned from the corpus, delivered per token.
Under the hood — the actual loss
import math
def kl_distillation_loss(student_logits, teacher_logits, temperature=2.0):
"""KL(teacher_T || student_T) * T^2 — the standard Hinton distillation loss.
Temperature T > 1 flattens both distributions, which amplifies the information
carried by the small probabilities (the 'dark knowledge'). The T^2 factor
restores the gradient magnitude, which otherwise scales like 1/T^2.
"""
def softmax_t(logits, T):
m = max(logits) # max-subtraction: never overflow
exps = [math.exp((z - m) / T) for z in logits]
s = sum(exps)
return [e / s for e in exps]
p_teacher = softmax_t(teacher_logits, temperature)
p_student = softmax_t(student_logits, temperature)
kl = sum(pt * math.log(pt / ps)
for pt, ps in zip(p_teacher, p_student) if pt > 1e-12)
return kl * temperature ** 2
# Toy vocabulary of 4 tokens. The teacher is confident but not absolute.
teacher = [4.0, 1.0, 0.5, -1.0]
good_student = [3.6, 1.2, 0.4, -0.8] # similar shape -> low loss
bad_student = [4.0, -5.0, -5.0, -5.0] # right argmax, wrong shape -> higher loss
print(round(kl_distillation_loss(good_student, teacher), 4)) # 0.031
print(round(kl_distillation_loss(bad_student, teacher), 4)) # 2.1126
Note what that demonstrates: bad_student gets the answer right (same argmax) and still
takes 68× the loss, because it has thrown away the shape. Distillation grades the shape,
not the answer.
The engineering problem nobody warns you about
Do the arithmetic on "trillions of tokens":
- 10 trillion tokens × a vocabulary of 256,000 × 2 bytes per logit
- = 5 × 10¹⁸ bytes = 5 exabytes.
You cannot store that. Nobody can. So the real job is choosing an approximation:
| Strategy | Storage for 10T tokens | Cost | Fidelity |
|---|---|---|---|
| Full logits | ~5 EB | impossible | perfect |
| Top-k logits (k=64) + renormalized tail mass | ~10 TB | cheap | very good |
| Online distillation (teacher runs alongside student) | 0 | teacher FLOPs every step, forever | perfect |
| Sampled sequences (teacher generates text, student trains on it) | ~20 TB text | one-time | weakest |
This is precisely the "infrastructure investments in storage systems and multi-datacenter operations" he refers to. The teacher may run in one datacenter while the student trains in another, so those top-k logits cross continents. The research idea is three lines; the system around it is a year of work. Phase 07 builds the top-k store and measures the fidelity loss.
Takeaway. Distillation is where "research" and "storage engineering" become the same job. If you want a differentiating skill, own that seam.
Claim 4 — Vertical two: inference co-design
Building "neural architectures that are efficient to run inference on" by choosing network topology, "shapes of the matrices," and "attention shapes, num heads" that "fully utilize the hardware."
Decode it
Two models with identical parameter counts can differ by 3× in serving speed, purely because of shape choices. Inference co-design is picking shapes with the chip in mind, before you train — because after you train, the shape is frozen forever.
Under the hood — three concrete shape decisions
Decision A: matmul dimensions should be multiples of the hardware tile. A TPU's matmul
unit is a systolic array (e.g. 128×128). A GPU's tensor cores want multiples of 8/16/64. If
your d_ff is 11,000, the hardware pads to 11,008 or worse, and you pay for arithmetic you
throw away.
def tile_efficiency(dim, tile=128):
"""Fraction of the padded matmul that is real work."""
padded = math.ceil(dim / tile) * tile
return dim / padded
for d in (11000, 11008, 8192, 4096, 4097):
print(f"d_ff={d:6d} padded={math.ceil(d/128)*128:6d} eff={tile_efficiency(d):.1%}")
d_ff= 11000 padded= 11008 eff=99.9%
d_ff= 11008 padded= 11008 eff=100.0%
d_ff= 8192 padded= 8192 eff=100.0%
d_ff= 4096 padded= 4096 eff=100.0%
d_ff= 4097 padded= 4224 eff=97.0%
Small here — but the same logic applied to head dimension, expert count, and shard boundaries compounds multiplicatively across dozens of layers, and the pathological cases are much worse than 97%.
Decision B: KV heads. During decode, you re-read the KV cache for every generated token. The cache size is:
KV bytes = 2 (K and V) × layers × kv_heads × head_dim × seq_len × batch × bytes_per_elem
Multi-Head Attention (MHA) gives every query head its own KV head. Grouped-Query Attention (GQA) shares one KV head across a group of query heads; Multi-Query Attention (MQA) uses exactly one. Quality barely moves; the cache shrinks by the group factor.
def kv_cache_gb(layers, kv_heads, head_dim, seq, batch, bytes_per=2):
return 2 * layers * kv_heads * head_dim * seq * batch * bytes_per / 1e9
# A 70B-class model, 8k context, batch of 32:
print(round(kv_cache_gb(80, 64, 128, 8192, 32), 1)) # MHA : 687.2 GB -> nine H100s of cache
print(round(kv_cache_gb(80, 8, 128, 8192, 32), 1)) # GQA-8: 85.9 GB -> just over one H100
print(round(kv_cache_gb(80, 1, 128, 8192, 32), 1)) # MQA : 10.7 GB -> trivially fits
That is one architecture decision moving memory by 64×. It has to be made before training starts. It is the single clearest example of what "inference co-design" means.
Decision C: depth vs width at fixed N. Deeper models are often slightly better per
parameter, but depth is serial — layer k+1 cannot start until layer k finishes — so
depth directly costs decode latency and creates more pipeline stages to sync. Width is
parallel and matmul-friendly. Co-design usually means: as wide as quality allows, as
shallow as quality tolerates.
Takeaway. Ask of every architecture choice: "what does this do to bytes-read-per-token at decode?" That single question is 80% of inference co-design.
Claim 5 — Vertical three: quantization, and the 99%-power fact
Reducing "the size that the neural nets take up in order to represent their weights" from FP32 down to four bits — which he calls "kind of a miracle." And the reason it matters: "99% of the total cost of operation for AI hardware comes from the power that it takes to run these chips."
Decode it
A number in a computer is stored in bits. FP32 = 32 bits per parameter. BF16 = 16. FP8 = 8. INT4 = 4. Quantization means storing the same model with fewer bits per number, accepting a small rounding error.
A 70B model: 280 GB in FP32, 140 GB in BF16, 70 GB in FP8, 35 GB in INT4. The last one fits in a single H100. The first needs four.
Under the hood — how you actually do it
The simplest scheme, affine (asymmetric) quantization, per group of weights:
def quantize_affine(values, n_bits=4):
"""Map a group of floats onto 2^n_bits integer levels. Returns (codes, scale, zero)."""
qmin, qmax = 0, 2 ** n_bits - 1
lo, hi = min(values), max(values)
if hi == lo: # degenerate group: everything identical
return [0] * len(values), 1.0, lo
scale = (hi - lo) / (qmax - qmin) # how much real value one code step is worth
zero = lo # what code 0 means
codes = [min(qmax, max(qmin, round((v - zero) / scale))) for v in values]
return codes, scale, zero
def dequantize_affine(codes, scale, zero):
return [c * scale + zero for c in codes]
w = [0.12, -0.45, 0.88, -0.03, 0.51, -0.77, 0.20, 0.05]
codes, scale, zero = quantize_affine(w, n_bits=4)
recon = dequantize_affine(codes, scale, zero)
err = max(abs(a - b) for a, b in zip(w, recon))
print(codes) # [8, 3, 15, 7, 12, 0, 9, 7]
print(round(err, 4)) # 0.05 -> max error is half a step, as it must be
Note the error bound: with 16 levels spanning a range of 1.65, one step is 0.11, and rounding
can never be off by more than half a step (0.055). That is the guarantee you test for —
max_error <= scale/2 is an invariant, not an empirical observation.
Why it works at all (the "miracle"): neural network weights are massively redundant and roughly bell-shaped. The network was trained with noise (dropout, stochastic gradients, bf16 rounding) so it is already robust to small perturbations. What matters is not the precision of any single weight but the statistical structure of the whole layer.
Why it stops working: outliers. A handful of weights (and far more importantly, a handful
of activation channels) are 100× larger than the rest. In a per-group scheme, one outlier
stretches scale and destroys the resolution of the other 63 values. Every serious method —
GPTQ, AWQ, SmoothQuant — is fundamentally an outlier-handling strategy.
The 99% fact, unpacked
This claim surprises people, so let us be careful about what it means. It is a claim about the total cost of operation — the ongoing cost of running a fleet, amortizing chips over years — and it lumps together electricity for the chips, electricity for cooling, and the power-delivery infrastructure sized to that draw. Under that accounting, the marginal cost of serving is energy, not silicon.
Now, why quantization is such an outsized lever on energy — this is a hardware fact worth memorizing. Moving a number costs vastly more energy than computing with it. Approximate energies at 45nm (Horowitz, ISSCC 2014 — the canonical reference; ratios still hold):
| Operation | Energy | Relative |
|---|---|---|
| 8-bit integer add | 0.03 pJ | 1× |
| 32-bit float add | 0.9 pJ | 30× |
| 32-bit float multiply | 3.7 pJ | 123× |
| Read 32 bits from on-chip SRAM | 5 pJ | 167× |
| Read 32 bits from off-chip DRAM | 640 pJ | ~21,000× |
Read that last row again. A DRAM read costs ~20,000× an integer add. So when you halve the bits, you halve the dominant term. Going FP16 → INT4 cuts weight-movement energy ~4×, and weight movement is most of decode.
# Rough marginal-energy model for decoding one token from an N-parameter model.
# Decode is memory-bound: the cost is dominated by reading every weight once.
PJ_PER_BIT_DRAM = 640 / 32 # ≈ 20 pJ per bit moved
def decode_energy_joules(n_params, bits_per_param):
return n_params * bits_per_param * PJ_PER_BIT_DRAM * 1e-12
for bits, name in [(32, "FP32"), (16, "BF16"), (8, "FP8"), (4, "INT4")]:
j = decode_energy_joules(70e9, bits)
# 1e9 tokens/day at $0.12/kWh
daily = j * 1e9 / 3.6e6 * 0.12
print(f"{name}: {j*1000:8.1f} mJ/token ${daily:9,.0f}/day at 1B tokens/day")
FP32: 44800.0 mJ/token $ 1,493/day at 1B tokens/day
BF16: 22400.0 mJ/token $ 747/day at 1B tokens/day
FP8 : 11200.0 mJ/token $ 373/day at 1B tokens/day
INT4: 5600.0 mJ/token $ 187/day at 1B tokens/day
The absolute numbers are a deliberate over-estimate in one direction (real systems reuse weights across a whole batch, which is exactly why batching exists) and an under-estimate in another (this counts only DRAM traffic, not cooling, power delivery, or the rest of the chip). The ratios are the point: FP32→INT4 is an 8× reduction in the dominant cost term. Scale that to Google's actual token volumes and it is a nine-figure line item.
He also notes the compounding: reducing activation precision on top of weights "multiplies" the gain — because now you shrink the other thing crossing the memory bus. It is also much harder, because activations have far worse outliers than weights.
Takeaway. Quantization is not a compression trick. It is the primary energy lever in AI, and energy is the bill. Phase 08 builds the quantizer and the power model.
Claim 6 — MFU, and why 100% is not the goal
He explicitly addresses Twitter confusion about "low" MFU numbers. To hit 100% you would need to be "doing a bunch of matmuls in a loop without reading any memory," which is not a neural network, because real nets "have to apply activation functions or do attention or write intermediate outputs."
Decode it
MFU = Model FLOPs Utilization.
MFU = (useful model FLOPs performed per second) / (chip's peak FLOP/s)
If your run does 3 × 10¹⁴ useful FLOP/s on hardware rated at 10¹⁵, your MFU is 30%. Published large-scale numbers commonly land in the 30–55% band. People see "35%" and conclude someone is incompetent. They are wrong, and here is exactly why.
Under the hood — where the other 65% goes
A chip is not one unit. It is several, with wildly different throughput:
┌──────────────────────────────────────────────┐
│ ACCELERATOR │
│ │
│ ┌───────────────┐ very fast (the peak) │
│ │ MATMUL UNIT │ ~1000 TFLOP/s │
│ └───────────────┘ │
│ ┌───────────────┐ ~50-100× slower │
│ │ VECTOR UNIT │ (gelu, softmax, norms) │
│ └───────────────┘ │
│ ┌───────────────┐ ~3 TB/s │
│ │ HBM │ (weights, activations) │
│ └───────────────┘ │
│ ┌───────────────┐ ~0.05-0.9 TB/s │
│ │ INTERCONNECT │ (chip-to-chip) │
│ └───────────────┘ │
└──────────────────────────────────────────────┘
"Peak FLOP/s" is the matmul unit's number alone. But your model must also:
- apply GELU/SwiGLU — vector unit, no matmul FLOPs credited
- compute softmax in attention — vector unit, plus exponentials
- compute layer norms / RMS norms — vector unit, plus a reduction (which serializes)
- write and read intermediate activations — HBM traffic, no FLOPs credited
- all-reduce gradients — interconnect, no FLOPs credited
- pay the optimizer step (Adam: several element-wise passes over all parameters)
Every one of those is time during which the matmul unit is idle. So:
def mfu_budget(matmul_s, vector_s, hbm_s, comms_s, optimizer_s):
total = matmul_s + vector_s + hbm_s + comms_s + optimizer_s
return {
"MFU (matmul busy fraction)": matmul_s / total,
"lost to vector ops": vector_s / total,
"lost to memory traffic": hbm_s / total,
"lost to collectives": comms_s / total,
"lost to optimizer": optimizer_s / total,
}
for k, v in mfu_budget(matmul_s=100, vector_s=45, hbm_s=60,
comms_s=50, optimizer_s=25).items():
print(f"{k:32s} {v:6.1%}")
MFU (matmul busy fraction) 35.7%
lost to vector ops 16.1%
lost to memory traffic 21.4%
lost to collectives 17.9%
lost to optimizer 8.9%
35.7% MFU is not a failure — it is an accounting identity. And notice what the breakdown gives you: an agenda. Comms at 18% says overlap them. Memory at 21% says fuse kernels. Vector at 16% says fuse the norm into the matmul epilogue.
His deeper point connects straight back to co-design: different shapes stress different units. The job is "choosing shapes for this neural net that fully saturate all of those hardware units" — not maximizing one number.
Two traps. (1) MFU is comparable only within a hardware/precision/model class — MFU on FP8 with sparsity is a different denominator. (2) Some teams quote HFU (Hardware FLOPs Utilization) which counts recomputation from activation checkpointing as useful work. HFU is always ≥ MFU. Always ask which one you are being shown.
Claim 7 — Pre-training is a one-shot extrapolation problem
His slides put it starkly. Before: "Maybe 2 stages; toy problem for iteration (CIFAR10) then you apply to Imagenet. LR searches by doing multiple 'final runs'. The last data point is our test set!" Now: "every single time you go up for a pre-training run, you're about to put in more FLOPs into this run than you've ever done before." So "every next run requires extrapolation."
Decode it
In classical ML, "the last data point is our test set" is a joke about a real practice: you try 20 learning rates, and the best one on the held-out set is your answer. You interpolate within a region you have already explored.
Pre-training breaks this in a way that is genuinely new:
- Each flagship run costs $10M–$100M+ and takes 1–3 months.
- You get one shot per generation.
- The run is larger than anything you or anyone has run before — by construction, since the whole point is to push the frontier.
So you are not interpolating. You are extrapolating beyond every data point you own. That is a fundamentally different statistical problem, and it is why "scaling laws" is a research area rather than a spreadsheet.
His slides add the crucial qualifier, which most summaries drop:
"Analysis made in the context of a parameterized LLM training recipe! Must already have architecture scaling, schedule defined for
N,D. Loss forecast implies model/recipe selection capability!"
Unpacked: a scaling law is not a law of nature. It is a property of your recipe. Before
you can fit one, you must have already decided how every hyperparameter scales with N and
D — how depth grows with width, how LR decays with batch size, how warmup scales. The law
then describes that family. Change the family and you must refit.
And the last sentence is the punchline of the whole field: if you can forecast loss, you can select recipes. Forecasting is not a reporting tool. It is the decision procedure.
Under the hood — the ladder, in code
# You cannot run the flagship twice. So you run a LADDER of small models,
# fit a curve, and extrapolate. Here is the shape of that in miniature.
def run_ladder(budgets):
"""Pretend-train at several compute budgets, recording (C, best_loss)."""
results = []
for C in budgets:
# In reality: sweep N at fixed C, take the minimum. See Phase 01.
loss = 1.69 + 406.4 / (compute_optimal_N(C) ** 0.34) \
+ 410.7 / (compute_optimal_D(C) ** 0.28) # Chinchilla-form
results.append((C, loss))
return results
def compute_optimal_N(C): return 0.6 * (C ** 0.5) # Chinchilla-ish: N ∝ C^0.5
def compute_optimal_D(C): return C / (6 * compute_optimal_N(C))
ladder = run_ladder([1e18, 1e19, 1e20, 1e21]) # cheap: hours, not months
flagship = 1e24 # the real run: months, $$$
print(run_ladder([flagship])) # THE FORECAST
The whole game: make the extrapolation from 1e21 to 1e24 — three orders of magnitude —
trustworthy. Phase 01 makes you do it for real, with the fit, the error bars, and the failure
modes.
Tip for interviews. If asked "how would you decide between two pre-training recipes?", the senior answer is never "train both and compare." It is: "fit a scaling law for each over a ladder of small runs, compare the fitted curves at the target FLOP count, and report the crossover point and the confidence interval." That is literally what his slides show: "To make a change, compare baseline vs candidate laws."
Claim 8 — C = 6ND, derived
From his slides: "for a transformer
C = 6 * N * Dis a very good approximation of FLOPs." The footnote: "Excluding self-attention, an N-parameter decoder-only model requires 6N matmul FLOPs per token seen (2N for forward and 4N for backward), because each matmul performs one multiplication and one addition per pair of input values, and the backward pass includes two matmuls for each one in the forward pass."
This is the most important equation in the field. Derive it once and you own it forever.
Step 1 — a single matmul costs 2 × (number of weights) FLOPs per token
Take a linear layer mapping a d_in-dimensional vector to d_out. Its weight matrix has
d_in × d_out entries. For one input vector:
output[j] = Σ_i input[i] * W[i][j] for each of d_out outputs
Each of the d_out × d_in terms is one multiply and one add = 2 FLOPs. Total:
2 × d_in × d_out = 2 × (#weights). Hence, over the whole model: 2N FLOPs per token,
forward.
Step 2 — backward costs twice forward
Backpropagation through the same linear layer needs two matmuls, not one:
forward : Y = X · W (1 matmul)
backward: dX = dY · Wᵀ <- to pass gradient to the previous layer
dW = Xᵀ · dY <- to update this layer's weights
(2 matmuls)
Both are the same size as the forward matmul. So backward = 2 × 2N = 4N FLOPs per token.
Step 3 — add them
6N FLOPs per token × D tokens = C = 6ND
def training_flops(n_params, n_tokens):
return 6 * n_params * n_tokens
def days_on_cluster(flops, n_chips, peak_flops_per_chip, mfu=0.4):
return flops / (n_chips * peak_flops_per_chip * mfu) / 86400
# Llama-3-70B-scale: 70B params, 15T tokens
C = training_flops(70e9, 15e12)
print(f"{C:.2e} FLOPs") # 6.30e+24 FLOPs
print(f"{days_on_cluster(C, 16000, 1e15):.1f} days on 16k H100") # 11.4 days
print(f"{days_on_cluster(C, 1000, 1e15):.1f} days on 1k H100") # 182.3 days
Now you can answer his opening question — "if I give you 1000 H100 for 30 days, what is the
best LLM you can train?" — because you can convert chips × days into C, and Phase 01 turns
C into (N, D).
Where 6ND breaks — know these four
- Attention is excluded. The
QKᵀandattn·Vmatmuls cost roughly12 · L · n_ctx · d_modelFLOPs per token, which does not scale withN. At short context this is a few percent; at 128k context it dominates. His slides show the exact per-step count:18BTDF + 24BTDNH = 6·BT·(3DF + 4DNH), whereB=batch,T=sequence,D=d_model,F=d_ff,N=num heads,H=head dim — the first term is the MLP, the second is attention projections. - MoE.
Nmust be the active parameter count (what each token actually routes through), not the total. A 400B-total / 40B-active MoE costs like a 40B dense model to train, and like a 400B model to store. This is exactly the audience question on his slide: "What about MoEs?" - Embeddings. The input embedding is a lookup (~free); the output unembedding is a real
matmul (
2 · d_model · vocabper token). At smallNwith a 256k vocab this is a large fraction — a common source of wrong small-model FLOP counts. - Activation checkpointing adds a partial extra forward pass, pushing you toward
~8NDof hardware FLOPs while the model FLOPs stay6ND. (This is the MFU/HFU gap from Claim 6.)
Takeaway.
C = 6NDis the currency conversion of the entire field: it turns money and time into model size and data. Memorize the derivation, not the formula.
Claim 9 — Kaplan said scale parameters; Chinchilla said scale both
His slides walk through this in detail, and it is the single most instructive story in scaling research — because it is a story about a methodological bug producing a wrong industry-wide strategy for two years.
Kaplan et al., 2020
Found that loss follows clean power laws in N, D, and C. Their compute-optimal
allocation, in his slides' words: "With a 10x compute budget, parameters should increase by
5.37x and the amount of data by 1.86x." Their own line: data requirements grow "very slowly
as D ∼ C^0.27."
Industry consequence (his slide states it directly): "We should heavily invest in scaling the model size rather than the data size!" This is why 2020–2022 was the era of ever-bigger, relatively under-fed models — GPT-3 at 175B trained on ~300B tokens.
His slide also lists the caveat, which everyone ignored at the time:
- "These 'laws' are only empirical"
- "The fitting of these laws depends a lot on the experimental setup as well as the implicit assumptions being made there."
Chinchilla (Hoffmann et al., GDM, March 2022)
His slide names the bug precisely:
"Kaplan et al. run a single training run per model size and uses intermediate losses to estimate the loss at different token horizon. ... This is a bad approximation as you can get much better losses through proper learning rate decay. Only the final loss value is optimal."
Here is why that is fatal, and it is worth being very concrete because it is subtle.
Learning-rate schedules decay to near zero at the end of training. That final decay phase is where a big chunk of the loss improvement happens — the model stops bouncing around the minimum and settles into it. So:
- A model mid-run at 100B tokens (LR still high, still bouncing) has a much worse loss
- than a model whose entire schedule was designed to end at 100B tokens (LR fully decayed).
Kaplan used the first as a proxy for the second. That systematically overstates how bad it is to train on more data — every data point in his "more tokens" direction was unfairly penalized. Correct for it, and the optimal shifts toward more tokens.
# Why reading loss mid-run is a biased estimator of "loss if I had stopped here".
def loss_at(tokens, horizon):
"""Toy: base curve + a penalty for not having decayed the LR yet."""
base = 3.0 / (tokens ** 0.1)
frac_done = tokens / horizon
lr_penalty = 0.15 * (1 - frac_done) # high LR = still noisy = worse loss
return base + lr_penalty
# Kaplan-style: peek at the 300B-token run when it has seen 100B tokens
print(round(loss_at(100e9, horizon=300e9), 4)) # 0.3383 <- biased HIGH
# Chinchilla-style: a run actually designed to end at 100B tokens
print(round(loss_at(100e9, horizon=100e9), 4)) # 0.2383 <- the truth
# The gap (0.1 nats) is enormous at this scale, and it is pure methodology.
And note the direction of the bias, which is what makes it fatal rather than merely noisy:
the penalty is proportional to (1 - frac_done), so it is largest exactly for the points
with the most tokens relative to their horizon. The error is not random — it systematically
tilts the fitted curve against training on more data.
The IsoFLOPs method — his slides' six steps
Chinchilla's cleanest approach, exactly as his slides enumerate:
1. Fix a target FLOPs budget ── e.g. C = 1e20
2. Train a few models, vary model size ── N = 100M, 300M, 1B, 3B (D = C/6N each)
3. Fit a parabola and find the minimum ── loss vs log(N) is U-shaped; take the vertex
4. Repeat 1–3 for various FLOPs budgets ── C = 1e19, 1e20, 1e21, 1e22
5. Fit a power law: FLOPs budget → optimal N ── N_opt ∝ C^a
6. Fit a power law: FLOPs budget → optimal D ── D_opt ∝ C^b
Why a parabola? Because at fixed C, there is a genuine trade-off with a single interior
minimum:
- Too small
N: you have tons of data but not enough capacity to absorb it → underfit. - Too large
N: enormous capacity but you starve it of data → also bad. - In between: the sweet spot. Plotting loss against
log Ngives a clean U.
# The IsoFLOPs inner loop, in full.
def isoflop_curve(C, sizes):
pts = []
for N in sizes:
D = C / (6 * N) # the budget constraint
L = 1.69 + 406.4 / N**0.34 + 410.7 / D**0.28 # Chinchilla parametric form
pts.append((N, D, L))
return pts
for N, D, L in isoflop_curve(1e21, [1e8, 3e8, 1e9, 3e9, 1e10, 3e10]):
print(f"N={N:8.1e} D={D:8.1e} loss={L:.4f}")
N= 1.0e+08 D= 1.7e+12 loss=2.6198 <- too small: underfit
N= 3.0e+08 D= 5.6e+11 loss=2.4344
N= 1.0e+09 D= 1.7e+11 loss=2.3400
N= 3.0e+09 D= 5.6e+10 loss=2.3363 <- the minimum
N= 1.0e+10 D= 1.7e+10 loss=2.4160
N= 3.0e+10 D= 5.6e+09 loss=2.5687 <- too big: data-starved
Notice how flat the bottom of that U is: 1e9 and 3e9 differ by only 0.004 nats. That
flatness is a gift and a trap. A gift, because you can move off the exact optimum for serving
reasons (Phase 02) at almost no quality cost. A trap, because with noisy measurements the
fitted minimum can wander by a factor of 3 — which is precisely why the parabola fit and its
confidence interval matter more than the single best point.
The result and its consequence
His slide: "the exponent in the power law is ~0.5, meaning model and data size should be scaled at the same rate! This is widely different from Kaplan et al." And he labels the old regime on the plot with one word: UNDERTRAINED!
"Consequences: Given a compute budget, models should be smaller and trained for longer. Kaplan's scaling laws meant that models were undertrained — which is obviously bad given bigger models are more expensive to serve and use downstream!"
That final clause is the bridge to his entire research agenda. Chinchilla didn't just improve loss-per-FLOP — it made models smaller at the same quality, which makes them cheaper to serve, which is the thing his team optimizes for.
Concretely: Chinchilla (70B, 1.4T tokens) beat Gopher (280B, 300B tokens) at the same training compute — with a model 4× smaller to serve.
Takeaway. The most consequential result in scaling laws came from fixing an experimental-design flaw, not from a new idea. Feinberg's own listed research direction — "Least squares vs MLE & formal stats model imply different scaling recommendations! Formalize." — says the field still has this class of bug in it. Methodology is the frontier.
Claim 10 — Scaling laws are brittle and dataset-dependent
From his closing slides: "Scaling laws are brittle, dataset dependent." And: "
L(N, D, etc.)— of course we can add more dims to improve fit. Least squares vs MLE & formal stats model imply different scaling recommendations! Formalize." Plus: "Rather than grid(N, D)where do we get max info gain? Active learn…"
Decode it
Three distinct criticisms hiding in there. Take them one at a time.
(a) The fit depends on your loss function. You have ~30 noisy (N, D, L) points and you
want parameters (A, B, E, α, β) for L = E + A/N^α + B/D^β. How you measure "fit" changes
the answer:
| Method | What it minimizes | Bias |
|---|---|---|
Least squares on L | Σ (L_pred − L_obs)² | dominated by large-loss (small-model) points |
Least squares on log L | relative error | treats all scales equally |
Huber loss on log L | relative error, outlier-robust | what Chinchilla actually used |
| MLE with an explicit noise model | likelihood under stated assumptions | requires you to state the noise model |
These give materially different exponents on the same data — and therefore different recommendations for the flagship run. Feinberg's "Formalize" is a call to stop hand-waving: write down the statistical model, then the estimator follows.
(b) The design points are chosen badly. Everyone runs a grid: N ∈ {100M, 300M, 1B, 3B} ×
C ∈ {1e19, 1e20, 1e21}. But a grid is not an efficient experiment. Optimal experimental
design asks: given my current uncertainty, which next run most reduces the variance of my
extrapolation at C = 1e24? Almost always the answer is "the largest one you can afford,
plus one that breaks a collinearity" — not "fill in the grid." This is the "active learn"
remark, and it is a genuinely open, publishable, cheap research direction.
# The intuition behind active learning for scaling laws, in miniature.
# Fitting a line from points clustered together gives a terrible slope estimate.
def slope_variance(x_points):
n = len(x_points)
xbar = sum(x_points) / n
sxx = sum((x - xbar) ** 2 for x in x_points)
return 1.0 / sxx # Var(slope) ∝ 1 / Σ(x - x̄)²
clustered = [19.0, 19.2, 19.4, 19.6] # four runs, all about the same size
spread = [18.0, 19.0, 20.0, 21.0] # same COUNT of runs, spread out
print(round(slope_variance(clustered), 3)) # 5.0 <- bad extrapolation
print(round(slope_variance(spread), 3)) # 0.2 <- 25x better, same budget
Same number of runs. 25× lower variance on the slope purely from where you placed them. That is what "where do we get max info gain?" means, and it costs nothing to apply.
(c) D is not what you think it is. This leads directly into Claim 16 — his slide says
D "was opaque and recipe-specific. You wouldn't be blamed for assuming iid." It is not iid:
repeated data, deduplication, and mixture weights all change what a "token" is worth.
Takeaway. Three fundable research projects sit in one slide, and none of them needs a supercomputer. This is the most actionable slide in the whole talk.
Claim 11 — MoE: what it is, why it wins, why it hurts
This is the concept the user asked about by name, so we build it completely from zero.
The problem MoE solves
In a dense transformer, every parameter participates in every token. If you want more knowledge in the model, you add parameters — and every token now costs more to process. Cost and capacity are welded together.
Mixture of Experts (MoE) breaks the weld. Replace the feed-forward block with E parallel
copies ("experts") plus a small router. For each token, the router picks the top-k
experts (typically k = 1 or 2) and only those run.
DENSE FFN MoE FFN (E=8, k=2)
token token
│ │
▼ ▼
┌────────┐ ┌─────────┐
│ FFN │ all params │ ROUTER │ tiny: d_model × E
│ (100%) │ run for every └────┬────┘
└────────┘ token │ scores 8 experts, picks best 2
│ ┌───┬───┼───┬───┬───┬───┬───┐
▼ ▼ ▼ ▼ ▼ ▼ ▼ ▼ ▼
output E0 E1 E2 E3 E4 E5 E6 E7
· ✓ · · ✓ · · ·
└───────┬───────┘
▼ weighted sum
output
8× the parameters. 2/8 = 25% of the compute per token.
The trade in one line: parameters (memory) scale with E; FLOPs scale with k.
The router, in full
import math
def softmax(xs):
m = max(xs)
e = [math.exp(x - m) for x in xs]
s = sum(e)
return [v / s for v in e]
def route(token_vec, router_w, k=2):
"""router_w: E x d_model. Returns [(expert_id, gate_weight), ...] of length k."""
logits = [sum(w * t for w, t in zip(row, token_vec)) for row in router_w]
probs = softmax(logits)
top = sorted(range(len(probs)), key=lambda i: -probs[i])[:k]
# Renormalize over the chosen k so the gates sum to 1.
total = sum(probs[i] for i in top)
return [(i, probs[i] / total) for i in top]
token = [0.5, -0.2, 0.9, 0.1]
router_w = [ # 6 experts, d_model = 4
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
[0.5, 0.5, 0.0, 0.0],
[0.0, 0.0, 0.5, 0.5],
]
print(route(token, router_w, k=2))
# [(2, 0.5987...), (0, 0.4013...)] -> experts 2 and 0, with those gate weights
Then the layer output is Σ_over_chosen gate_i × expert_i(token).
Load balancing — the thing that actually breaks
Left alone, routers collapse. Early in training one expert is randomly slightly better,
so it gets more tokens, so it trains more, so it gets better, so it gets more tokens. Within a
few thousand steps you have one expert doing everything and E−1 dead ones — you paid for 8×
the parameters and got a dense model.
The fix is an auxiliary load-balancing loss added to the training objective:
def load_balance_loss(assignments, gate_probs, n_experts):
"""Switch-Transformer style: L_aux = E * Σ_i f_i * P_i.
f_i = fraction of TOKENS routed to expert i (discrete, no gradient)
P_i = mean ROUTER PROBABILITY assigned to expert i (continuous, differentiable)
Multiplying them makes the loss differentiable through P while being driven
by the actual imbalance in f. Minimized when both are uniform (= 1/E each),
giving L_aux = E * E * (1/E) * (1/E) = 1.0.
"""
n_tokens = len(assignments)
f = [0.0] * n_experts
for e in assignments:
f[e] += 1.0 / n_tokens
P = [0.0] * n_experts
for probs in gate_probs:
for i, p in enumerate(probs):
P[i] += p / n_tokens
return n_experts * sum(fi * Pi for fi, Pi in zip(f, P))
E = 4
balanced = [0, 1, 2, 3] * 2
collapsed = [0] * 8
uni = [[0.25] * 4] * 8
skew = [[0.97, 0.01, 0.01, 0.01]] * 8
print(round(load_balance_loss(balanced, uni, E), 4)) # 1.0000 <- ideal
print(round(load_balance_loss(collapsed, skew, E), 4)) # 3.8800 <- heavily penalized
You add α · L_aux (typically α ≈ 0.01) to the main loss. Too small and you collapse; too
large and you damage quality by forcing nonsensical routing. This coefficient is one of the
most finicky hyperparameters in modern pre-training.
Capacity factor — the batching constraint nobody mentions
Hardware wants fixed-size tensors. So each expert gets a fixed buffer:
capacity = capacity_factor × (tokens_in_batch × k / n_experts)
If more tokens route to an expert than fit, the overflow is dropped (passed through by the residual connection, unprocessed). If fewer arrive, the buffer is padded with zeros — wasted compute.
def capacity_stats(assignments, n_experts, capacity_factor=1.25, k=1):
n_tokens = len(assignments)
cap = int(capacity_factor * n_tokens * k / n_experts)
counts = [0] * n_experts
for e in assignments:
counts[e] += 1
dropped = sum(max(0, c - cap) for c in counts)
padded = sum(max(0, cap - c) for c in counts)
return {"capacity_per_expert": cap, "dropped": dropped, "padded_slots": padded,
"drop_rate": dropped / n_tokens}
# 100 tokens, 4 experts, mildly imbalanced routing
assign = [0]*40 + [1]*30 + [2]*20 + [3]*10
print(capacity_stats(assign, 4, capacity_factor=1.25))
# {'capacity_per_expert': 31, 'dropped': 9, 'padded_slots': 33, 'drop_rate': 0.09}
9% of tokens silently skip the FFN. This is a real, live source of quality loss in production
MoE models, and tuning capacity_factor against drop rate is routine pre-training work.
What his slides say about MoE scaling
"MoE scaling laws are better, but have implications for token hunger. We're running out of internet! ... Notice relative data hunger compared to dense! At same active param count and fixed 100B token training, MoE 64E improves on dense."
Decoded: for a fixed compute budget, an MoE reaches lower loss than a dense model — the
scaling law is strictly better. But the compute-optimal D for an MoE is larger. MoE
converts "we have compute" into "we need more unique tokens," and unique high-quality tokens
are the resource that is actually running out. MoE trades a compute problem for a data
problem. (The reference here is Clark et al., Unified Scaling Laws for Routed Language
Models, 2022.)
Takeaway. MoE is not free capacity. You pay in HBM, in routing instability, in dropped tokens, in data hunger, and — the big one, next — in communication. Phase 03 builds all of this; Phase 06 fixes the communication.
Claim 12 — The MoE serving bottleneck, and Geng Yan's pipeline-prefill fix
This is the most technically specific story in the interview, and it is worth full precision.
The problem
An MoE has many more parameters than a dense model of equal compute. Those parameters must live in HBM. A single chip's HBM (16 GB on a v5e; 80 GB on an H100) cannot hold them. So you shard the experts across chips — expert 0–7 on chip 0, experts 8–15 on chip 1, and so on. This is expert parallelism (EP).
Now trace one token through one layer:
Layer ℓ: token lives on chip 0
router says "you need expert 37"
expert 37 lives on chip 4
──> send the token's activation vector to chip 4 [NETWORK]
chip 4 computes
──> send the result back to chip 0 [NETWORK]
Layer ℓ+1: router says "you need expert 12" (on chip 1)
──> send to chip 1 [NETWORK]
──> send back [NETWORK]
... repeat for every one of ~60 layers.
His description of exactly this: "that token might live on the first TPU, but it needs to go to the last TPU." Every layer, for every token. The collective involved is an all-to-all (every chip sends a different slice to every other chip), which is the most expensive collective there is, and its cost, as he notes, "increases dramatically with N."
def moe_alltoall_cost(layers, tokens, d_model, n_chips, bytes_per=2,
link_bw_gbps=100, latency_us=5):
"""Rough per-forward-pass communication cost of naive expert parallelism."""
bytes_per_hop = tokens * d_model * bytes_per
# Each layer: dispatch to experts + combine back = 2 all-to-alls.
total_bytes = layers * 2 * bytes_per_hop
transfer_s = total_bytes / (link_bw_gbps * 1e9 / 8)
latency_s = layers * 2 * latency_us * 1e-6 # fixed cost, paid per collective
return {"GB moved": total_bytes / 1e9,
"transfer_s": transfer_s,
"latency_s": latency_s,
"total_s": transfer_s + latency_s}
print(moe_alltoall_cost(layers=60, tokens=8192, d_model=8192, n_chips=16))
{'GB moved': 16.1, 'transfer_s': 1.288, 'latency_s': 0.0006, 'total_s': 1.289}
1.3 seconds of pure network time, before a single useful FLOP. For a product with a sub-second latency budget, that is fatal. This is the wall the Flash team hit.
The insight
He credits Geng Yan — described as a junior member of the team — with the fix. Restated precisely:
Stop parallelizing across experts. Parallelize across layers instead.
That is pipeline parallelism, applied at prefill time. Chip 0 holds layers 1–10 (with all their experts). Chip 1 holds layers 11–20. And so on.
Now the token does not hop around per layer. It flows forward through the pipeline, once:
NAIVE EXPERT PARALLELISM PIPELINED PREFILL
(shard experts, keep all layers) (shard layers, keep all experts local)
chip0 ⇄ chip1 ⇄ chip2 ⇄ chip3 chip0 ──> chip1 ──> chip2 ──> chip3
↕ ↕ ↕ ↕ L1-15 L16-30 L31-45 L46-60
all-to-all EVERY layer (×60) one hop per STAGE (×3), point-to-point
And now the crucial second half of the idea — his exact framing: "layer one on the first chip is processing the second thousand tokens... while layer two is working on the first thousand tokens." Chunk the prompt and stream the chunks through the pipeline, so every stage is busy on a different chunk simultaneously:
time ──────────────────────────────────────────────────>
chip0 (L1-15) : [chunk1][chunk2][chunk3][chunk4][chunk5]
chip1 (L16-30): [chunk1][chunk2][chunk3][chunk4]
chip2 (L31-45): [chunk1][chunk2][chunk3]
chip3 (L46-60): [chunk1][chunk2]
└ bubble ┘ └──── steady state, all chips busy ────┘
The chip→chip transfer of chunk i overlaps with computation on chunk i+1.
Communication is HIDDEN, not eliminated.
His summary of the effect: communication went "from something that required a lot of token exchange on every single layer to something that actually can be hidden behind other computation."
def pipeline_prefill(n_chunks, n_stages, compute_per_stage_ms, transfer_ms):
"""Steady-state pipeline: total time = fill + (chunks * per-chunk cost)."""
step = max(compute_per_stage_ms, transfer_ms) # overlapped -> the max, not the sum
fill = (n_stages - 1) * step # the 'bubble' at the start
return fill + n_chunks * step
serial = 4 * 8 * (50 + 20) # 4 stages, 8 chunks, no overlap
piped = pipeline_prefill(8, 4, 50, 20)
print(serial, "ms serial ->", piped, "ms pipelined") # 2240 ms -> 550 ms
Bubble efficiency is the thing to reason about: with S stages and M chunks, utilization
is M / (M + S − 1). With 4 stages and 8 chunks: 8/11 = 73%. With 32 chunks: 32/35 = 91%.
More chunks = smaller bubble, but each chunk is a smaller matmul with worse arithmetic
intensity. That trade-off is the design.
Why prefill specifically
This is the part that makes the insight clever rather than obvious, and most summaries miss it.
- Prefill processes the whole prompt at once. It is compute-bound and has thousands of tokens available to chunk. Pipelining has plenty of work to hide communication behind. ✅
- Decode produces one token at a time. It is memory-bandwidth-bound, and there is
no parallel work to overlap with. A pipeline here just adds
S−1serial hops to every single token. ❌
So you use different parallelism strategies for the two phases of the same request — which is exactly the modern prefill/decode disaggregation design. His remark that this "serving-time innovation" is what "made [an MoE Gemini 2.0 series] possible" is the payoff: the serving strategy unlocked the architecture choice.
Takeaway — the transferable lesson. The fix was not a better kernel or a better model. It was changing which axis you shard along, informed by which phase of inference you are in. Note also who found it: a junior engineer. Feinberg's response was to run "a very transparent technical process to get to the bottom of this" — the leadership behavior that lets a junior insight become a flagship decision. Phase 06 builds this scheduler.
Claim 13 — The Flash 2.0 war story: 40 days, five people, two continents
The largest model they had ever trained at the Flash scale. "40 days of grueling work for a really, really small team" — "five people on the rotation," rotating "day by day, handing off all of this SRE-style work of keeping the training job alive." They monitored data iterators and fixed indexing issues to avoid "wasting all of this GPU time." "We did not do a lot of sleeping" — "dual shifts across the Paris office and Mountain View."
Decode it — why does a training run need a pager?
A newcomer's model of training is model.fit(). At this scale the reality is a distributed
system with thousands of nodes running for weeks, and the failure modes are exotic:
| Failure | What it looks like | Why it is brutal |
|---|---|---|
| Loss spike | loss jumps 2→8 in one step | May recover, may permanently damage the model. You have minutes to decide whether to roll back. |
| Hardware failure | one chip of 10,000 dies | Synchronous training means the whole job stalls. At 10k chips with a 3-year MTBF, expect a failure every few hours. |
| Silent data corruption (SDC) | a chip computes wrong numbers without erroring | The worst one. No crash, no alert; the model just gets subtly worse. Detected only by cross-replica checksums. |
| Data iterator bug | wrong shard, an off-by-one, a repeated segment | This is the one he names. Silent. You can burn days of compute training on the wrong data. |
| Stragglers | one slow node | Everyone waits at the all-reduce barrier; throughput collapses to the slowest node. |
| Checkpoint corruption | the save itself fails | Your rollback point is gone. Discovered at the worst moment. |
Under the hood — goodput, the number that matters
def goodput(total_hours, crash_count, restart_minutes, checkpoint_interval_min):
"""Fraction of wall-clock actually spent making forward progress.
Every crash costs: the restart itself, PLUS the work done since the last
checkpoint (lost), PLUS on average half a checkpoint interval of redo.
"""
lost_per_crash_h = (restart_minutes + checkpoint_interval_min / 2) / 60
lost = crash_count * lost_per_crash_h
return max(0.0, (total_hours - lost) / total_hours)
# 40 days = 960 hours. Compare a fragile setup to a hardened one.
print(round(goodput(960, crash_count=200, restart_minutes=45,
checkpoint_interval_min=120), 3)) # 0.635
print(round(goodput(960, crash_count=200, restart_minutes=10,
checkpoint_interval_min=15), 3)) # 0.939
Same hardware, same number of crashes: 64% vs 94% goodput. That 30-point gap is 12 days of a 40-day run — or, in money, a seven-figure swing, and quite possibly the difference between shipping before a competitor and shipping after. The entire delta comes from two boring engineering decisions: fast restarts and frequent (asynchronous) checkpoints. This is why the SRE hat is not beneath the researcher; it is a large fraction of the deliverable.
Note also which lever matters more. Halving restart time saved ~35 minutes per crash; going from 2-hour to 15-minute checkpoints saved ~52 minutes per crash. Checkpoint cadence is the bigger lever, and it is limited by how fast you can write terabytes of optimizer state — which is why asynchronous and sharded checkpointing is a real engineering discipline, not a config flag.
The loss-spike playbook
The single most common 3 a.m. event. A working decision procedure:
def spike_response(loss_history, window=100, z_threshold=6.0):
"""Detect a spike against a rolling baseline, and pick an action."""
if len(loss_history) < window + 1:
return "warmup: insufficient history"
recent = loss_history[-window - 1:-1]
mean = sum(recent) / len(recent)
var = sum((x - mean) ** 2 for x in recent) / len(recent)
std = var ** 0.5 or 1e-9
z = (loss_history[-1] - mean) / std
if z < z_threshold:
return "normal"
if z < 15:
return "WATCH: skip this batch, log it, continue" # often self-heals
return "ROLLBACK: restore last checkpoint, skip N batches, lower LR"
Standard mitigations, in escalation order: (1) skip the offending batch (a bad data shard is a common cause); (2) lower the learning rate and re-warm; (3) roll back to the last checkpoint and skip forward past the data; (4) if it recurs at the same step, it is the data, not luck — go find it. Real interventions from published reports include z-loss regularization, QK-norm, and clipping gradients more aggressively.
The competitive coda
Flash 2.0 landed around the same time as DeepSeek-V3. He notes that a Wall Street Journal piece had "some elided rows" in its leaderboard comparison, making Gemini look badly positioned — while the actual leaderboard showed "Flash 2.0 Thinking up in the top right corner, way far ahead of DeepSeek-V3."
The transferable lesson has nothing to do with either company: leaderboard screenshots are marketing artifacts. Always ask what was filtered, which variant was tested, what the axes are, and whether the comparison controls for cost. His own slides make the same point from the other side — "LMSys is not the end-all-be-all," and Llama 4 Maverick showed ranking can be "volatile and overfit to human preference."
Takeaway. Frontier pre-training is an operations discipline wearing a research hat. If you want a way in that is less crowded than "research scientist," become excellent at training-run reliability. Phase 09 builds the watchdog.
Claim 14 — Real-time products force small models (the napkin math)
His slides do this calculation live, and it is the single best worked example in the talk of how a product requirement becomes an architecture constraint. Reproduced and extended:
The setup
A web-interaction agent with:
- 128k prefill, but only 8k incremental per turn
- 128 decode tokens (enough to emit an action)
- No more than 1 second of latency between actions
- 250ms of that goes to "scaffolding, load balancing, request validation, kv cache retrieval" — and he flags this as "optimistic!"
Experiment: Llama3-70B on v5e chips. Assume fully compute-bound on prefill, HBM-bound on decode.
The arithmetic
# --- Chip and model constants -------------------------------------------------
V5E_FLOPS = 197e12 # bf16 peak FLOP/s for one TPU v5e
V5E_HBM_BW = 819e9 # bytes/s of HBM bandwidth
V5E_HBM_GB = 16 # GB capacity <- note: a 70B model in bf16 needs 140 GB
N = 70e9 # Llama3-70B parameters
BYTES_PER_PARAM = 2 # bf16
def prefill_seconds(n_tokens, n_params, n_chips, mfu=1.0):
"""Prefill is COMPUTE bound: 2N FLOPs per token (forward only).
mfu=1.0 is his stated idealization: 'assume fully compute bound on prefill'."""
flops = 2 * n_params * n_tokens
return flops / (n_chips * V5E_FLOPS * mfu)
def decode_seconds(n_tokens, n_params, n_chips, efficiency=1.0):
"""Decode is MEMORY bound: every generated token re-reads EVERY weight.
With the model sharded over n_chips, each chip reads its 1/n_chips slice."""
bytes_moved = n_params * BYTES_PER_PARAM * n_tokens
return bytes_moved / (n_chips * V5E_HBM_BW * efficiency)
for chips in (1, 4, 16, 64, 128):
p = prefill_seconds(8192, N, chips) # the 8k incremental prefill
d = decode_seconds(128, N, chips) # 128 tokens to emit one action
total = p + d + 0.25 # + the 250ms scaffolding budget
verdict = "OK " if total <= 1.0 else "MISS"
print(f"{chips:4d} chips: prefill {p:7.3f}s decode {d:7.3f}s "
f"total {total:7.3f}s {verdict}")
1 chips: prefill 5.822s decode 21.881s total 27.953s MISS
4 chips: prefill 1.455s decode 5.470s total 7.175s MISS
16 chips: prefill 0.364s decode 1.368s total 1.982s MISS
64 chips: prefill 0.091s decode 0.342s total 0.683s OK
128 chips: prefill 0.045s decode 0.171s total 0.466s OK
The prefill column reproduces his slide exactly: "Uh oh… 5.7 seconds for 1 chip. So to hit 0.5 sec api limit we already need to have a 4x4 prefill station of v5e." One chip gives 5.8s; a 4×4 = 16-chip station brings prefill to 0.36s, under his 0.5s API limit. And the audience question he poses next is the right one: "how would we shard on 4x4?" — 16 chips is not a number, it is a topology (tensor/pipeline split, mesh shape, and the collectives that follow). Phase 04 answers it.
The five conclusions that follow
- The prefill station alone needs 16 chips to serve one conversation inside the latency budget. Not 16 chips for the service — 16 for one user.
- Decode is worse than prefill at batch 1, by ~3.8×. Every generated token re-reads all 140 GB of weights to do a trivial amount of arithmetic. This is the memory-bound regime from Phase 00, and it is why the full budget needs ~64 chips, not 16.
- The two phases want different hardware allocations — which is the entire argument for prefill/decode disaggregation, and the setup for Claim 12.
- Therefore: make the model smaller. Halving
Nhalves both columns linearly. This is the entire economic case for Flash and Flash-Lite, and it is why his job exists. - The 70B doesn't even fit. 140 GB of weights across 16 GB chips means ≥9 chips just to hold it, before any latency consideration. Capacity and latency push the same way.
Caveat, stated honestly: batch-1 decode is the worst case. Real serving batches many requests, which amortizes the weight reads across all of them and dramatically improves the decode column — that is why continuous batching exists. But an interactive agent that must respond in under a second cannot always wait to fill a batch. The tension between latency (small batch) and throughput (large batch) is the central trade-off of serving, and this table is where you feel it.
Takeaway. This is the calculation that converts "we want a real-time agent" into "we need a distilled, quantized, inference-co-designed small model." Do this arithmetic before you pick a model, not after. Phase 05 makes you build the full version with a roofline.
Claim 15 — Chinchilla ignores inference cost
His slide: "Chinchilla-style Scaling Ignores Inference Cost." The direct fix: "Globally optimize FLOPs between training and inference?" — citing Sardana et al., Beyond Chinchilla-Optimal, 2024.
Decode it
Chinchilla minimizes loss for a fixed training budget. But a deployed model's lifetime cost is training plus all inference, forever. If you serve a lot, it is worth overtraining a smaller model — spending more training FLOPs than Chinchilla says, on a model smaller than Chinchilla says — because you amortize that over trillions of served tokens.
def total_lifetime_flops(N, D_train, D_inference):
return 6 * N * D_train + 2 * N * D_inference # 6ND to train, 2N/token to serve
# Two models targeting the SAME quality (illustrative, from published fits).
chinchilla = dict(N=70e9, D_train=1.4e12) # "compute-optimal"
overtrained = dict(N=20e9, D_train=8.0e12) # smaller, trained much longer
for name, m in [("Chinchilla-optimal", chinchilla), ("Overtrained-small", overtrained)]:
for served in (1e12, 1e14, 1e16):
total = total_lifetime_flops(m["N"], m["D_train"], served)
print(f"{name:20s} served={served:.0e} lifetime={total:.3e} FLOPs")
print()
Chinchilla-optimal served=1e+12 lifetime=7.280e+23 FLOPs
Chinchilla-optimal served=1e+14 lifetime=1.459e+25 FLOPs
Chinchilla-optimal served=1e+16 lifetime=1.401e+27 FLOPs
Overtrained-small served=1e+12 lifetime=1.000e+24 FLOPs
Overtrained-small served=1e+14 lifetime=4.960e+24 FLOPs
Overtrained-small served=1e+16 lifetime=4.010e+26 FLOPs
At 10¹² served tokens the Chinchilla model wins. By 10¹⁴ the overtrained small model is ~2.9× cheaper overall, and by 10¹⁶ it is 3.5× cheaper. The crossover is the decision, and it depends entirely on a number the research team does not control: how many tokens the product will serve.
The three problems he raises with this idea
His slides are notably skeptical of the clean version, and the objections are the interesting part:
(1) Non-homogeneity of compute. "Inference-optimized chips. Also global optimization is not how cross-org planning actually works." Training FLOPs and inference FLOPs are not the same currency — they happen on different chips, in different datacenters, on different budgets, owned by different VPs. "But in principle can adjust the formulas for 'business cost'" — i.e. the real objective is dollars, and the FLOP-exchange-rate is an org-chart question as much as a physics one.
(2) Non-forecastability of D_inf. You cannot know how many tokens you will serve. Two
named reasons:
- Jevons paradox — making a resource cheaper increases total consumption. Every efficiency win you deliver gets eaten by more usage.
- Market expansion from quality improvements — a better model unlocks use cases that did not exist, so demand is a function of the very quality you are optimizing.
So D_inf appears in your objective and is caused by your objective. Feedback loop, no fixed
point. This is a genuinely unsolved problem, not a modelling nicety.
(3) Badness of fit. He points at the paper's own Fig 5 / Table 1. The inference-aware laws are extrapolating into the heavily-overtrained regime — far past where the fits were calibrated — and that is exactly where the classical functional forms behave worst. Which leads directly to the next claim.
Takeaway. The right objective is lifetime cost, not training cost. But the honest version of the calculation contains a term nobody can forecast, so in practice you do scenario analysis over
D_infand pick something robust across scenarios. Phase 02 builds this.
Claim 16 — "We're running out of internet"
His slide, on MoE's data hunger: "We're running out of internet!" And the fix he spends most of his time on: "Unsurprisingly, this is where we spend most of our time, even as modelling people. Probably half my focus this year so far." — namely (1) multimodal data (audio, visual, 3D, video) and (2) synthetic data.
Note that line carefully: the pre-training lead spends half his time on data. Not on architecture. Not on optimizers. On data.
The data-constrained law
The reference is Muennighoff et al., Scaling Data-Constrained Language Models (2023). His slide's framing is the sharpest summary you will find of why it matters:
"
Dwas opaque and recipe-specific. You wouldn't be blamed for assuming iid." "New dimension: intentionally unique data,L(N, U, R)." "Upshot: yet smaller models, more resilient to repeats."
Decoded: everyone writes L(N, D) where D = "tokens seen." But 1 trillion tokens seen could
mean 1T unique tokens once, or 200B unique tokens five times. Those are not the same
experiment, and the classical law cannot tell them apart. So you split the axis:
U= unique tokens in your corpusR= number of repeats (epochs)D = U × R
And the empirical finding: repeated tokens are worth almost as much as fresh ones for the first few epochs, then their value decays fast.
import math
def effective_tokens(U, R, half_life=5.0):
"""Value of R epochs over U unique tokens, with exponentially decaying returns.
Matches the paper's qualitative finding: ~4 epochs is nearly free,
~16 epochs adds almost nothing, and after that you may be actively hurting.
"""
return U * sum(math.exp(-(r - 1) / half_life) for r in range(1, R + 1))
U = 100e9
for R in (1, 2, 4, 8, 16, 32):
eff = effective_tokens(U, R)
print(f"R={R:2d} raw D={U*R/1e9:6.0f}B effective={eff/1e9:6.1f}B "
f"efficiency={eff/(U*R):5.1%}")
R= 1 raw D= 100B effective= 100.0B efficiency=100.0%
R= 2 raw D= 200B effective= 181.9B efficiency= 90.9%
R= 4 raw D= 400B effective= 303.8B efficiency= 75.9%
R= 8 raw D= 800B effective= 440.3B efficiency= 55.0%
R=16 raw D= 1600B effective= 529.2B efficiency= 33.1%
R=32 raw D= 3200B effective= 550.7B efficiency= 17.2%
At R = 32 you have spent 32× the compute for 5.5× the effective data — and the last 16
epochs bought you 4% more. His slide's "5 epochs" annotation sits right at the elbow of that
curve, which is not a coincidence.
This is a toy model, deliberately. The exponential-decay form is a stand-in that reproduces the paper's qualitative finding (a few epochs are nearly free, many are nearly worthless). The real fit in Muennighoff et al. has a different functional form with fitted half-life parameters. Use this to build intuition; use the paper to make decisions.
Why this changes the architecture decision
If unique data is the binding constraint, big models are the wrong answer, because big models are precisely the ones that need lots of fresh data to justify their capacity. Hence his slide's conclusion: "yet smaller models, more resilient to repeats." Data scarcity and serving economics push in the same direction — which is very convenient, and is a large part of why the industry converged on Flash-class models.
The two escape hatches
Multimodal data. Audio, images, video, 3D. Text on the internet is finite; video is effectively not. It also carries information text never encodes (physical dynamics, spatial relations). The cost: tokenization and encoder design get much harder, and token counts explode (a second of video can cost hundreds of tokens).
Synthetic data. He makes a subtle point most people get wrong:
"Without filter, it can help in the Stein's paradox sense (Jain et al 2024). Tradeoff: Generation Quality vs. Filtering."
Stein's paradox is the famous statistical result that a biased estimator can have lower total error than an unbiased one — shrinking your estimates toward a common point beats using each observation on its own. Applied here: synthetic data is biased (it reflects the generator's distribution, not the truth) but it is lower variance. Adding it can reduce total error even though it adds bias. That is a much more precise and more defensible argument for synthetic data than "we ran out of text."
The trade-off he names is real: generate more carefully (expensive) or generate cheaply and filter hard (also expensive, and filtering can be its own source of bias).
The counterpoint on his own slide
"Llama3:
D_inf= inf!" — quoting the Llama 3 paper: "Both our 8B and 70B parameter models continued to improve log-linearly after we trained them on up to 15T tokens.""Could be quite valid for open source! Just pick sizes and train on all your data! We could be doing research with those FLOPs! Use this forecast to estimate how much regret we got."
That is an honest and slightly self-critical note. If returns have not saturated, the simplest strategy — pick a size and train on everything — is close to optimal, and all the sophisticated scaling work buys you less than it appears to. His framing of "regret" (how much worse off you were than the best strategy in hindsight) is the right way to hold this. And "job is to push the curves right" is the mission statement.
Takeaway. Data is where a pre-training lead's time actually goes. If you want to be useful to a frontier lab fast, get extremely good at data quality, deduplication, mixture weighting, and synthetic-data filtering. It is less glamorous than architecture and worth more.
Claim 17 — Distillation scaling laws and the capacity gap
He cites Busbridge et al., Distillation Scaling Laws (2025), and the question it poses: "How to spend FLOPs with teacher?" Then he pushes back on one of its findings, and the pushback is a masterclass in reading a paper critically.
The question
You have a compute budget. You can spend it on: (a) making the teacher better, (b) running the teacher over more tokens to generate more supervision, or (c) training the student longer. It is a three-way allocation problem, and it has an optimum.
The "capacity gap" and his three objections
The claimed phenomenon: if the teacher is too much better than the student, distillation gets worse, not better — the student cannot represent the teacher's function, so chasing it hurts.
His response, point by point (from his slide, annotated):
(1) "very weak effect from up-trend; and not typical regime." The effect is small in the data, and the region where it appears is not where anyone actually operates. Lesson: always ask whether a reported effect is in the regime you care about.
(2) The temperature objection — this one is elegant. "Teacher pplx can be arbitrarily weakened by just adding temperature! Take a really good teacher → Eq8 predicts bad distill → but add high temp and it will be good?"
Unpacked: the law expresses the capacity gap in terms of teacher perplexity. But you can change a teacher's effective perplexity for free by raising the softmax temperature — that flattens its distribution without changing what it knows. So the formula predicts you could fix the capacity gap by turning a knob that carries no information. That is a reductio: if a free, information-free transformation moves your predictor, your predictor is parameterized on the wrong variable.
def perplexity(logits, T=1.0):
"""Raising temperature flattens the distribution and RAISES perplexity —
without the teacher knowing anything less."""
m = max(logits)
e = [math.exp((z - m) / T) for z in logits]
s = sum(e)
p = [x / s for x in e]
H = -sum(pi * math.log(pi) for pi in p if pi > 0)
return math.exp(H)
teacher = [8.0, 2.0, 1.0, 0.5, 0.0]
for T in (1.0, 2.0, 4.0, 8.0):
print(f"T={T}: perplexity={perplexity(teacher, T):.3f}")
T=1.0: perplexity=1.033
T=2.0: perplexity=1.626
T=4.0: perplexity=3.429
T=8.0: perplexity=4.586
Same teacher, same knowledge, perplexity moved 4.4× from a knob you set at inference time. Any law keyed on that number inherits the knob.
(3) "In practice, you can James-Stein this away with weight tuning with supervised objective." In production you never distill purely — you mix the distillation loss with the ordinary next-token loss:
def combined_loss(student_logits, teacher_logits, true_token, lam=0.5, T=2.0):
"""lam=1 -> pure distillation; lam=0 -> pure supervised. Reality lives in between."""
distill = kl_distillation_loss(student_logits, teacher_logits, T)
ce = cross_entropy(student_logits, true_token)
return lam * distill + (1 - lam) * ce
Tuning λ is literally a shrinkage estimator — the James-Stein reference is exact, not
metaphorical. You are trading bias (teacher's errors) against variance (single-label noise).
And that gives his cleanest reframe of the entire topic:
"Distill as variance reduction. Better teacher will just reduce bias."
That one sentence is the best mental model for distillation in existence. The teacher's dense distribution is a low-variance estimate of the true next-token distribution — enormously less noisy than a one-hot label. The teacher's imperfection is the bias. Better teacher → less bias. Distillation at all → less variance. Everything else is engineering.
Takeaway. Watch how he read that paper: check the regime, look for a free transformation that breaks the parameterization, and check whether the effect survives normal practice. That is what "mathematical maturity" means in his hiring criteria, made concrete.
Claim 18 — The Gemini tick-tock
His slide: "Gemini tick-tock (Flash goal to match Pro of previous gen)." And: "Scaling Work has two flavors: (1) Adding points to Quality × Model Size plot. (2) Increasing the slope of the plot."
Decode it
Borrowed from Intel's old CPU cadence. The pattern:
Gen N : Pro ──── quality X ────┐
│ next generation must deliver X
Gen N+1 : Flash ──── quality X ───┘ at a fraction of the size and cost
Pro ──── quality X+Δ
Every generation, the small cheap model must reach what the previous generation's flagship reached. That is a compression mandate with a deadline, and it is why his three verticals are what they are: distillation, quantization, and serving-friendly architecture are the only three ways to hit it.
The two "flavors" distinction is worth internalizing:
| Flavor | What it means | How you do it |
|---|---|---|
| Adding points | fill in the quality-vs-size curve at a new size | train another model at a new N |
| Increasing the slope | make the whole curve better — more quality per parameter | better architecture, better data, better distillation, better optimizer |
His closing note ties them: "Inference Efficiency Work: Compression work grows with both scaling aspects" — via (1) better distillation recipes, (2) quantization, (3) serving-friendly model design changes.
Takeaway. Adding points is production. Moving the slope is research. Know which one you are being asked for, and know which one you are doing.
Claim 19 — You can do pre-training research without a supercomputer
His slide is titled "Future Pretrain Research Ideas – Without Big Costs!" and opens with the objection it demolishes: "Common refrain: pretraining is expensive, only can be researched in industry." Then four counterexamples. If you want research to put on a resume, this is the list.
1. Kernels and kernel languages. "Developing hardware-focussed kernels is the hot-loop for research now. Kernel programming languages, compiler tools, developer tools that make this easier are crucial. Or come up with the next flash attention." — Needs one GPU. High impact. Directly hireable. (Phase 11.)
2. Vector quantization. "Quantization entering a new frontier from vector quant." Scalar quantization rounds each weight independently. Vector quantization replaces groups of weights with a codebook index, exploiting correlations between them. Far better rate-distortion in principle; largely unexploited in LLMs. Needs a laptop to prototype. (Phase 08.)
3. FunSearch-style inference-vs-quality trade-offs. "For LLM-in-the-loop for search." How much extra inference compute (samples, search width, verification) buys how much quality? This is test-time-compute scaling, and it is an open, cheap research area.
4. The statistics of scaling laws. Covered in Claim 10 — the estimator question and the optimal-design question. Pure statistics on published data. Zero GPUs. Possibly the highest ratio of impact-to-cost on the entire list.
Takeaway. The barrier to entry is not compute. It is knowing which questions are open. He just told you four of them.
Claim 20 — Hiring: intent, mathematical maturity, grit
He screens for three things: "intent, mathematical maturity, grit." Mathematical maturity means "being able to dive into a paper of that level and then understand it, being able to take a research idea from a paper and implement it yourself." He also stresses "having read and having the skills to effectively traverse the historical citation tree for a particular topic" and knowing "what are the high-value papers."
The concrete asks, and how to actually execute them
(a) Do How To Scale Your Model ("The Scaling Book") exercises — handwritten, on video. He publicly offered interviews for this, and referrals when he lacked headcount. This is an open, standing, verifiable offer, and it is stunning how few people take it.
How to do it well: work the arithmetic by hand, show the units at every step, and narrate why each quantity matters, not just what it equals. The video is evidence you did it yourself; the narration is evidence you understood it.
(b) Implement a transformer from scratch. Demonstrates "willingness to get into the weeds engineering-wise" and the "bread-and-butter math that we use every day to size these LLMs."
How to do it well: not a 200-line nanoGPT copy. Include the FLOP counter, the memory accounting, a KV cache, and a benchmark showing you know where time actually goes. The differentiator is the instrumentation, not the model.
(c) Contribute to vLLM / SGLang / TensorRT. "Actual evidence that you've created something of use to other people" is what he wants — an improvement "for this and that setting."
How to do it well: pick a narrow, measurable win (a kernel for an under-served shape, a scheduler edge case, a quantization format), benchmark it honestly, and write the PR description like a small paper — problem, measurement, fix, measurement.
(d) Learn to traverse a citation tree. This is a teachable procedure, not a vibe:
1. Find the newest strong survey or the most-cited recent paper on the topic.
2. Read its related-work section. Note which papers it treats as LOAD-BEARING
(described in detail) vs merely listed.
3. Follow those backward to the origin papers. Read the origin paper in full.
4. Use a citation index to walk FORWARD from the origin: who cites it, and
which of those are highly cited themselves? That intersection is the spine.
5. Build a timeline: what changed at each step, and WHY the previous answer
was insufficient. That "why" chain IS the field's argument.
6. Find the paper that CONTRADICTS the consensus. Understand why it lost —
or whether it actually did.
Step 6 is the one that separates people. Kaplan→Chinchilla is exactly that shape, and Feinberg's critique of the distillation capacity gap (Claim 17) is him doing step 6 live.
Takeaway. All four asks share a property: they produce an artifact someone else can check. That is the entire signal. Not a credential, not a course certificate — a thing that exists, that works, that someone else can use.
Claim 21 — The internal-transfer play
For someone at a big company but outside the frontier team: don't chase the transfer. Ask "how do I help my product area adopt this technology as effectively as possible?" You become "the partner that we work with on the research side." He cites Nate Lintz, who transferred in and "owns so much of what we do... in terms of inference." And he adds that the transfer may not even be necessary, because integrating new technology into real products people use is itself cutting-edge work.
Why this works — the mechanism
Frontier research teams have a structural problem: they build things and need them adopted, but they don't understand your product's constraints, data, users, or latency budget. If you become the person who does understand both sides, you are not applying for a job — you are already doing half of it, and they already depend on you.
The playbook:
1. Become the LLM person for your product area. Actually deploy something.
2. Hit a real wall (latency, cost, quality, a serving limitation).
3. Bring the research team a PROBLEM WITH DATA — not a request.
"Our p99 is 2.1s; here's the profile; the prefill is 80% of it."
4. Collaborate on the fix. Now you have a joint artifact and a joint author list.
5. Repeat. You are now the default partner for that surface.
6. The transfer, if you still want it, is a formality — they will ask you.
Step 3 is the whole thing. A profile and a number gets a researcher's attention. A request does not.
Takeaway. The shortest path into a frontier team usually runs through the product, not around it. And it is far less competitive than the front door.
Claim 22 — On AI replacing engineers
He calls the discourse "FUD everywhere, especially with some of the approach to marketing that some people have." His counter is accountability: "there's an element of making decisions around how we allocate these resources that will always be something that needs to be attributable to a human making that decision" — "you can't hand off blame to AI." The lawyer example: they remain necessary because "they can't be disbarred." And his prescription: "we all have agency over our future and we can start investing in skills that matter for tomorrow today."
The argument, made precise
It is not "AI is bad at this." It is a claim about institutions. Professions with liability — law, medicine, engineering sign-off, financial audit — are structured around a person who can be sanctioned. Licensure, malpractice, disbarment, fiduciary duty. A model cannot hold a license, cannot be sued, cannot be struck off. So even a perfectly capable model needs a human principal to attach accountability to.
The honest counterargument, which you should hold too: this bounds the floor, not the size. One accountable lawyer with excellent tools may replace ten. Accountability preserves the role, not the headcount.
Where that leaves you, practically:
- The durable skills are judgment under uncertainty (Claim 2), verification (can you tell whether the output is right?), and accountability (will you sign your name to it?).
- The compressible skills are pure translation — spec-to-code, paper-to-implementation.
- His own field is a good example of the floor: nobody is going to let a model unilaterally spend $50M of TPU time. Someone signs.
Takeaway. He is not saying "relax." He is saying stop consuming the discourse and start compounding skill — which is, notably, the same thing he says about everything else.
Claim 23 — The career philosophy
Two pieces of advice, both anti-intuitive, both worth more than the technical content.
(1) Chase real problems, including the menial ones
"Chase the problems that people are facing in the world today. Go after the challenges that people see in everyday life, and don't be afraid to tackle a smaller part of this problem or maybe a more menial sounding part."
He is describing his own path. He was in pure research maximizing first-author papers at NeurIPS/ICML/ICLR. His manager Rohan Anil pushed him toward Bard. The work was hyperparameter tuning and getting a model to run on old TPUs — about as unglamorous as frontier AI gets. Jeff Dean gave him a spot bonus for it, and it put him on the path to leading Gemini pre-training.
The mechanism behind why this works:
- Menial work sits close to reality, so it teaches you constraints nobody has written down.
- It is under-supplied, because status-seekers avoid it, so your marginal value is high.
- It makes you legible to the people who ship, which is where decisions are made.
- And it seeds real problems: "quantizing the Ads DNN for pCTR serving efficiency" is not a glamorous line, and it is the direct intellectual ancestor of leading Flash.
(2) Be someone people want to see succeed
"Be the kind of co-worker that people would want to see succeed." Leverage other people's complementary skills "in ways that help them shine" — and then "people will notice, people will want to contribute to projects that you come up with in the future." He explicitly contrasts this with the "workplace psychopath" / Machiavellian model, and credits mentors like Todd Lipkin, who first got him into computer science, as the kind of person who "genuinely inspire[s] me to want to help them succeed."
The game-theoretic version, since it sounds soft but is not: careers are iterated games with reputation and information asymmetry. Zero-sum play works in one-shot games with observable payoffs. Neither condition holds in a career. What actually determines your opportunities is whether someone who already trusts you brings you in — and that is a pure reputation effect, accumulated over years.
Note also the Geng Yan story from Claim 12 in this light: a junior engineer's idea became the central technical bet of a flagship model, and the lead's contribution was to run "a very transparent technical process" around it. That is this philosophy operating at the level of technical decisions.
Takeaway. The technical content of this document is depreciating — architectures change. The two paragraphs above are not.
Master Takeaways
The whole document in twenty lines. If you internalize nothing else:
C = 6ND. Derive it. It converts money into models.- Pre-training is one-shot extrapolation. Every run is bigger than every run you have data for. This is why forecasting is a research field.
- A scaling law is a property of your recipe, not of nature. Change the recipe, refit.
- Loss forecasting is recipe selection. Compare fitted laws at the target FLOP count, not runs at small scale.
- Chinchilla beat Kaplan by fixing an experimental-design bug, not by having a better idea. Methodology is the frontier.
- Chinchilla-optimal is the wrong objective if you serve at scale. Optimize lifetime cost —
but know that
D_infis unforecastable (Jevons, market expansion). - MoE buys capacity with memory, communication, instability, and data hunger. Nothing is free; it just moves.
- The Flash 2.0 unlock was changing the sharding axis — layers instead of experts, at prefill — so communication hides behind compute.
- Prefill and decode are different machines. Compute-bound vs bandwidth-bound. Never apply one's optimization to the other.
- MFU of 35% is an accounting identity, not a failure. Read the breakdown; it is your agenda.
- Power is the bill. ~99% of TCO. A DRAM read costs ~20,000× an integer add. Quantization is an energy lever first and a memory lever second.
- Distillation is variance reduction; a better teacher reduces bias. The teacher's shape is the signal, not its argmax.
- Distillation's hard part is storage, not the loss. 10T tokens × 256k vocab is 5 EB. Top-k or online.
- Data is where the pre-training lead's time goes — half his year. Not architecture.
Dis not iid.L(N, U, R): unique tokens and repeats are different axes. ~4 epochs is nearly free; 32 is mostly waste.- Synthetic data helps in the Stein's-paradox sense — biased but lower variance.
- Training runs need SREs. Goodput 78% vs 94% on identical hardware is checkpoint cadence and restart speed. That is days of a 40-day run.
- Research is an MDP. Buy information before outcomes. Write kill criteria in advance. Diversify across premises, not implementations.
- You can do real pre-training research on one GPU: kernels, vector quantization, test-time-compute trade-offs, and the statistics of scaling laws.
- The signal that gets you hired is an artifact someone else can check — and the career that compounds is built on doing unglamorous work well and making other people shine.
References
Primary sources for this document
- Vlad Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides (PDF)
- Ryan Peterman (Developing Dev), Google DeepMind Pre-Training Lead: How To Get a Job at a Frontier Lab · video
- vladfeinberg.com / about
Papers named in the talk or interview
- Kaplan et al., Scaling Laws for Neural Language Models, 2020 — https://arxiv.org/abs/2001.08361
- Hoffmann et al., Training Compute-Optimal Large Language Models (Chinchilla), 2022 — https://arxiv.org/abs/2203.15556
- Sardana et al., Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Laws, 2024 — https://arxiv.org/abs/2401.00448
- Muennighoff et al., Scaling Data-Constrained Language Models, 2023 — https://arxiv.org/abs/2305.16264
- Busbridge et al., Distillation Scaling Laws, 2025 — https://arxiv.org/abs/2502.08606
- Clark et al., Unified Scaling Laws for Routed Language Models, 2022 — https://arxiv.org/abs/2202.01169
- Pope et al., Efficiently Scaling Transformer Inference, 2022 — https://arxiv.org/abs/2211.05102
- Xu et al., GSPMD: General and Scalable Parallelization for ML Computation Graphs, 2021 — https://arxiv.org/abs/2105.04663
- DeepSeek-AI et al., DeepSeek-V3 Technical Report, 2024 — https://arxiv.org/abs/2412.19437
- Austin et al., How To Scale Your Model (The Scaling Book) — https://jax-ml.github.io/scaling-book/
Supporting references used in the explanations above
- Hinton, Vinyals & Dean, Distilling the Knowledge in a Neural Network, 2015 — https://arxiv.org/abs/1503.02531
- Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated MoE Layer, 2017 — https://arxiv.org/abs/1701.06538
- Fedus, Zoph & Shazeer, Switch Transformers, 2021 — https://arxiv.org/abs/2101.03961
- Dao et al., FlashAttention, 2022 — https://arxiv.org/abs/2205.14135
- Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models, 2023 — https://arxiv.org/abs/2305.13245
- Horowitz, Computing's Energy Problem (and what we can do about it), ISSCC 2014 — the canonical energy-per-operation table
- Jacob Steinhardt, Research as a Stochastic Decision Process — https://cs.stanford.edu/~jsteinhardt/ResearchasaStochasticDecisionProcess.html
- Grattafiori et al., The Llama 3 Herd of Models, 2024 — https://arxiv.org/abs/2407.21783
Lab Standard — Frontier Pre-Training Lead
Every lab in this track builds a runnable, test-verified miniature of a real pre-training decision procedure — the FLOP counter, the IsoFLOPs fitter, the MoE router with capacity, the collective cost model, the roofline, the pipeline scheduler, the teacher-logit store, the 4-bit quantizer, the goodput model, the PPO step, the tile-DSL interpreter, the research planner.
You do not "call scaling_laws.fit()". You implement the parabola fit, the power-law
regression, and the extrapolation, and then you report the confidence interval. That is what
makes the knowledge defensible when a director asks you to justify a $40M run.
Why miniatures and not real training runs
A real pre-training run costs eight figures, needs a datacenter and clearances, is
non-deterministic, and hides every mechanism behind XLA and a scheduler. A lab that reimplements
the decision algebra — 6ND, the IsoFLOPs parabola, the load-balance loss, the all-to-all
byte count, the arithmetic-intensity ridge point, the pipeline bubble fraction, the
quantization error bound, the goodput integral — is offline, deterministic, free, and teaches
exactly what a frontier interview probes.
Every lab README ends with a "How this maps to the real stack" section connecting the miniature to JAX/XLA, GSPMD, MaxText, Megatron-LM, DeepSpeed, vLLM, or the relevant paper — including where the miniature lies.
Required files (per lab)
| File | Contract |
|---|---|
README.md | the problem, what you build, key-concepts table, file map, run commands, success criteria, "How this maps to the real stack", extensions, interview/resume bullets |
lab.py | learner implementation with focused # TODO markers; signatures and docstrings already in place; never a blank file |
solution.py | complete reference; python solution.py runs a worked example and prints output; deterministic |
test_lab.py | positive, negative, boundary, numerical-stability and determinism tests; runnable against either module via LAB_MODULE |
requirements.txt | usually pytest only — labs are pure stdlib otherwise |
The runnable core is Python (stdlib + pytest), offline, deterministic. No GPU, no CUDA, no
network, no model downloads, no pip install torch, no unseeded randomness. Where a lab needs
vectors, matrices, or a least-squares fit, implement it with stdlib lists/math so the
mechanism stays visible — NumPy only if a lab explicitly declares it in its
requirements.txt, and even then the algorithm must be the thing you write.
Determinism rules
- Any randomness goes through an explicit seeded
random.Random(seed); same seed → same bytes. Tests assert this. - Float comparisons in tests use
pytest.approxor an explicit tolerance — never==. softmax, log-sum-exp, KL, and any loss use the max-subtraction / log-space trick, so tests can include large logits without overflow. That numerical detail is part of the lesson.- Curve fits use closed-form or iterative solvers you write, with a fixed iteration count and a fixed tolerance — no library optimizer whose version could change the answer.
Units discipline (specific to this track)
Half of frontier interview failures are unit errors. So:
- Every function that returns a physical quantity names the unit in the identifier or the
docstring:
flops,bytes_moved,seconds,joules,dollars_per_million_tokens. - Tests include at least one dimensional-analysis assertion — e.g. doubling parameters
doubles FLOPs; halving bandwidth doubles memory-bound time;
2Nforward and4Nbackward sum to6N. - Any constant taken from hardware (peak FLOP/s, HBM bandwidth, pJ per bit) lives in a single named table at the top of the module with a source comment, never inline.
The test contract
import importlib, os
lab = importlib.import_module(os.environ.get("LAB_MODULE", "lab"))
Run both ways; the reference must pass, and your lab.py passes once the TODOs are filled:
pytest test_lab.py -v # against your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v # against the reference (must be green)
python solution.py # the worked example
Test taxonomy every lab includes:
- happy path — the textbook case
- malformed / out-of-range input — raises
ValueError(negative parameter counts, zero bandwidth,top_k > n_experts, capacity factor below 1/k) - boundary cases — the off-by-ones interviewers probe: one chip, one expert, one token,
k = n_experts(MoE degenerates to dense), zero repeats, a single pipeline stage, a quantization group where every value is identical - numerical-stability cases — large logits do not overflow;
log(0)is guarded; a degenerate parabola fit is detected rather than dividing by zero - invariants — softmax sums to 1; gate weights over the chosen top-k sum to 1;
forward + backward = 3 × forwardFLOPs;dequant(quant(x))is within the group's error bound; pipeline utilization isM/(M+S−1); goodput ≤ 1 - deterministic output — same seed → same bytes
The two teaching docs (per phase)
| Document | Voice | What it gives you |
|---|---|---|
README.md | the syllabus | why the phase exists, concept map, lab spec, deliverables checklist, key takeaways |
WARMUP.md | the professor | zero-to-principal primer: every term from first principles → what it is → why it exists → how it works under the hood (mechanism, diagrams, math, code) → production significance → common misconceptions; then Lab Walkthrough, Success Criteria, Interview Q&A, Tips & Takeaways, and References to primary sources |
Doc conventions (mdBook-compatible)
WARMUP.mdopens with a Table of Contents of working anchor links, kept in sync with the headings.- MathJax for math:
\( … \)inline,$$ … $$block. - No bare angle brackets in prose — wrap
<like-this>in backticks so mdBook does not eat them as HTML. - File references use relative links. Code fences are language-tagged.
- Every quantitative claim taken from a paper or talk carries a citation in the phase's References section.
Definition of done
A lab is complete only when the reference suite passes (LAB_MODULE=solution pytest), the
learner lab.py passes once the TODOs are filled, python solution.py prints a sensible worked
example, the README's success criteria are testable, and the "How this maps to the real stack"
section is honest about where the miniature diverges from production.
A phase is complete when both teaching docs exist, the WARMUP's ToC anchors resolve, and every lab in it is done.
Glossary — Frontier Pre-Training
One line each. Terms are grouped by where they first matter, and cross-linked to the phase that builds them properly.
Compute and arithmetic
FLOP — one floating-point operation (a multiply or an add). "FLOPs" counts work; "FLOP/s" measures speed. Phase 00
FMA — fused multiply-add: one instruction, counted as two FLOPs.
C = 6ND — training compute for N parameters over D tokens: 2N forward + 4N
backward, per token. The currency conversion of the entire field.
2N — inference FLOPs per token (forward only). Used for prefill and for the serving term
of lifetime cost. Never mix it up with 6N.
Matmul — matrix multiply; 2·m·k·n FLOPs. ~99% of a transformer's arithmetic.
Systolic array — the TPU's fixed-size (e.g. 128×128) grid of multiply-accumulate cells. Dimensions that are not multiples of the tile get padded, and you pay for the padding.
Tensor cores — NVIDIA's equivalent matmul units.
Arithmetic intensity — FLOPs performed per byte moved. Compare against the ridge point to know whether you are compute- or memory-bound.
Ridge point — peak FLOP/s ÷ HBM bandwidth. H100 ≈ 296 FLOP/byte. Below it → memory-bound.
Roofline — the model that plots achievable throughput against arithmetic intensity, with a bandwidth-limited slope and a compute-limited ceiling.
Precision and memory
FP32 / BF16 / FP16 / FP8 / INT4 — number formats at 32/16/16/8/4 bits. BF16 beat FP16 for training because it keeps FP32's exponent range, and gradients span an enormous dynamic range.
HBM — High Bandwidth Memory, physically attached to the accelerator. The scarcest resource in serving. Two numbers matter: capacity (GB) and bandwidth (TB/s).
Mixed precision — bf16 for matmuls, fp32 for the optimizer's master weights and moments. Costs 16 bytes per parameter with Adam.
ZeRO / FSDP — sharding of optimizer states (stage 1), gradients (stage 2) and parameters (stage 3) across data-parallel replicas. Stage 1 is the cheapest big win because optimizer states are ~75% of training memory.
Activation checkpointing — discard intermediate activations and recompute them in the backward pass. Trades ~30% more FLOPs for a large memory saving, and is the reason MFU and HFU differ.
KV cache — the stored keys and values from earlier tokens, so generation does not recompute
them. 2·L·n_kv·d_head·T·B·bytes. The serving wall.
Model shape
d_model — width of the residual stream. d_ff — the MLP's hidden width (~4× d_model,
or ~2.7× for gated). n_heads / n_kv_heads / d_head — attention head counts and
dimension. L — layers. V — vocabulary size.
MHA / GQA / MQA — Multi-Head, Grouped-Query, and Multi-Query attention: every query head gets its own KV head, a shared one per group, or a single shared one. The single biggest lever on KV-cache size, and it is frozen at pre-training time.
SwiGLU / gated MLP — an MLP with three matrices (up, gate, down) instead of two.
Non-embedding parameters — parameter count excluding the embedding and unembedding tables.
The correct N for scaling-law work: a 256k vocabulary can be 89% of a small model.
Scaling laws
Scaling law — an empirical formula predicting test loss from resources, typically
L = E + A/N^α + B/D^β. A property of your recipe, not of nature.
Phase 01
Irreducible loss (E) — the entropy of the data itself; a floor no model beats. ~87% of the
number at frontier scale.
Capacity term (A/N^α) — error from the model being too small. Data term (B/D^β) —
error from not having seen enough.
Recipe — the full parameterized specification of a run: architecture scaling, data mixture, optimizer, schedule, numerics, parallelism. A law is meaningless without one fixed.
IsoFLOPs — fix C, sweep N, derive D = C/6N, fit a parabola in log N, take the vertex.
Repeat across budgets, then fit N_opt ∝ C^a and D_opt ∝ C^b.
a + b ≈ 1 — a consistency check that falls directly out of C = 6ND. Free bug detector.
Kaplan (2020) — found N_opt ∝ C^0.73; concluded the industry should scale parameters over
data. Chinchilla (2022) — found C^0.5 after fixing a measurement bias; concluded models
were undertrained.
LR-decay / schedule-mismatch bias — reading loss part-way through a run whose learning-rate
schedule targets a longer horizon. A uniform bias is absorbed into E; the non-uniform
one tilts the fitted exponents. The mechanism behind Kaplan → Chinchilla.
Huber loss — quadratic near zero, linear in the tails. Used on log-space residuals so a single diverged run cannot dominate the fit.
Bootstrap — resample the ladder with replacement, refit, repeat, take percentiles. How you get a confidence interval when the model is nonlinear and the noise model is unspecified.
Crossover — the compute budget at which two fitted laws swap places. The deliverable of a recipe comparison; "candidate is better" silently assumes a scale.
Data-constrained law L(N, U, R) — splits D into unique tokens U and repeats
R. A few epochs are nearly free; many are nearly worthless.
Inference-aware scaling — optimize 6N·D_train + 2N·D_inf rather than training cost alone.
Complicated by the fact that D_inf is unforecastable (Jevons paradox, market expansion).
Jevons paradox — making a resource cheaper increases total consumption. Efficiency wins get eaten by more usage.
Mixture of Experts
MoE — replace each block's MLP with E parallel experts plus a router that activates
top_k per token. Parameters scale with E; FLOPs scale with k.
Phase 02
Router — the tiny d_model × E matrix that scores experts. The most fragile component in
the model.
Top-k routing — select the k highest-scoring experts. k=1 is Switch; k=2 is the common
default; k == E is a dense model in disguise.
Gate renormalization — rescale the chosen k gates to sum to 1. Without it, the layer's
output is silently attenuated by a per-token amount.
Router collapse — the rich-get-richer loop in which one expert wins everything and the rest die. The default dynamic, not an exotic failure.
Load-balancing loss — L = E · Σ f_i·P_i, where f is the (discrete) token-slot fraction
and P the (differentiable) mean router probability. Perfect balance gives exactly 1.0.
Router z-loss — a penalty on logsumexp(logits)². Prevents bf16 overflow and softmax
saturation, which are two different failures.
Capacity factor — the safety margin on each expert's fixed buffer. Too low drops tokens; too high pads with zeros. No setting avoids both.
Token dropping — overflow past an expert's capacity. Silent: the token rides the residual and nothing is logged.
Shared expert — one that runs for every token unconditionally. Structurally eliminates fully dropped tokens and lets routed experts specialize harder.
Active vs total parameters — active goes into 6ND and serving cost; total goes into HBM,
checkpoints and sharding. Off by 10–20× if swapped.
Expert parallelism (EP) — sharding experts across chips. Forces two all-to-all collectives per layer.
Parallelism and serving
DP / TP / PP / EP — data, tensor, pipeline and expert parallelism.
All-reduce / all-gather / reduce-scatter / all-to-all — the collective operations. All-to-all (every chip sends a different slice to every other) is the most expensive, and is what naive expert parallelism requires.
GSPMD — Google's compiler-driven sharding system: annotate tensors, let XLA infer the rest.
Prefill — process the entire prompt at once. Compute-bound.
Decode — generate one token at a time. Memory-bandwidth-bound, at every realistic batch size. This asymmetry governs all of serving.
Pipelined prefill — shard layers across chips rather than experts, and stream prompt chunks through, so transfers hide behind computation. The Flash 2.0 unlock, credited to Geng Yan.
Pipeline bubble — the idle time while a pipeline fills and drains. Utilization is
M/(M + S − 1) for M chunks and S stages.
Prefill/decode disaggregation — running the two phases on separately-provisioned hardware, because they are bound by different resources.
Continuous batching / PagedAttention — serving techniques that raise decode throughput by amortizing weight reads across many concurrent requests.
Efficiency and operations
MFU (Model FLOPs Utilization) — useful model FLOPs achieved ÷ peak. 35–55% is normal at scale; it is an accounting identity, not a failure grade.
HFU (Hardware FLOPs Utilization) — the same, but counting recomputation from activation checkpointing as useful. Always ≥ MFU. Always ask which you are being shown.
Goodput — the fraction of wall-clock spent making forward progress. Driven mostly by checkpoint cadence and restart speed; 64% vs 94% is ~12 days of a 40-day run.
Loss spike — a sudden jump in training loss. May self-heal or may permanently damage the model; you have minutes to decide between skipping the batch, lowering LR, and rolling back.
Silent data corruption (SDC) — a chip computing wrong numbers without erroring. The worst failure mode, because only cross-replica checksums catch it.
Straggler — one slow node; at a synchronous all-reduce barrier, everyone waits for it.
PUE — Power Usage Effectiveness: total datacenter power ÷ IT power. 1.1–1.5 typically.
Compression
Distillation — training a student to match a teacher's full output distribution rather than a one-hot label. Best understood as variance reduction; a better teacher reduces bias. TRANSCRIPT-DISSECTED §Claim 3
Temperature — the softmax divisor T. Raising it flattens the distribution and amplifies
the "dark knowledge" in small probabilities. The Hinton loss carries a T² factor to restore
gradient magnitude.
Top-k logit store — keeping only the k largest teacher logits plus the renormalized tail
mass. The only tractable way to store teacher outputs at trillion-token scale (full logits would
be ~5 exabytes).
Quantization — representing weights (and sometimes activations) in fewer bits. An energy lever before a memory lever: a DRAM read costs ~20,000× an integer add, and power is ~99% of hardware TCO.
Affine / asymmetric quantization — map a group of values onto 2^b integer levels with a
scale and zero-point. Max error is half a step, by construction.
Outliers — the handful of weights (and far worse, activation channels) that are ~100× larger than the rest, stretching the scale and destroying resolution for everything else. Every serious method (GPTQ, AWQ, SmoothQuant) is fundamentally an outlier-handling strategy.
Vector quantization — replacing groups of weights with a codebook index, exploiting correlations between them. Feinberg names this as an open frontier requiring little compute.
Research practice
Research as an MDP — planning under uncertainty over a stochastic dependency graph, as opposed to engineering's deterministic DAG where progress is monotone.
Research taste — well-calibrated priors on p(success) plus the habit of buying
information before buying outcomes. Trainable, not mystical.
Value of information — the worth of an experiment that changes what you do next, over and above its direct payoff. Usually the dominant term.
Kill criteria — the stopping rule written before the experiment starts, because afterwards you will be attached to it.
Optimal experimental design — choosing where to place runs to minimize the variance of the extrapolation you actually care about. Spread beats density: ~26× at identical cost.
Regret — how much worse off you were than the best strategy in hindsight. Feinberg's framing for auditing whether sophisticated scaling work actually beat "pick a size and train on everything."
Tick-tock — Gemini's generational cadence: this generation's Flash should match last generation's Pro. A compression mandate with a deadline.
Cheat Sheet — Frontier Pre-Training
Everything worth having memorized, on one page. If you can reproduce this from memory you can hold your own in a pre-training design review.
The core equations
| Quantity | Formula | Notes |
|---|---|---|
| Matmul cost | 2·m·k·n FLOPs | the 2 is multiply + add |
| Linear layer, per token | 2 × params FLOPs | forward only |
| Training compute | C = 6ND | 2N forward + 4N backward |
| Inference compute | 2N per token | prefill and serving |
| Exact per-step (MHA, gated) | 18BTDF + 24BTDNH = 6·BT·(3DF + 4DNH) | MLP term + attention-projection term |
| Attention (seq-dependent) | 12·B·n_h·T²·d_h·L | the part 6ND omits |
| Unembedding | 6·d_model·V per token | large fraction at small N |
| Chinchilla split | N = √(C / 6r), D = rN | r ≈ 20 tokens/param |
| Lifetime cost | 6N·D_train + 2N·D_inf | the objective Chinchilla ignores |
| Scaling law | L = E + A/N^α + B/D^β | irreducible + capacity + data |
Compute-optimal N | N_opt ∝ C^(β/(α+β)) | = 0.5 exactly when α = β |
| Consistency check | a + b ≈ 1 | falls out of C = 6ND — free bug detector |
Memory
| Term | Bytes per parameter | Notes |
|---|---|---|
| bf16 weights | 2 | |
| bf16 gradients | 2 | |
fp32 master + Adam m + v | 12 | |
| Mixed-precision Adam total | 16 | before a single activation |
| Adafactor | 4 (optimizer) | factored second moment |
KV cache = 2 · L · n_kv · d_head · T · B · bytes
ZeRO-1 shards optimizer states (75% of the total — cheapest big win)
ZeRO-2 adds gradients
ZeRO-3 adds parameters (an all-gather per layer)
70B model, mixed-precision Adam: 1.12 TB unsharded. An H100 has 80 GB. You must shard.
Hardware constants
| Chip | bf16 peak | HBM | Bandwidth | Watts |
|---|---|---|---|---|
| H100 SXM | 990 TFLOP/s | 80 GB | 3.35 TB/s | 700 |
| A100-80 | 312 TFLOP/s | 80 GB | 2.03 TB/s | 400 |
| TPU v5e | 197 TFLOP/s | 16 GB | 819 GB/s | 170 |
| TPU v5p | 459 TFLOP/s | 95 GB | 2.77 TB/s | 600 |
Ridge point = peak FLOP/s ÷ bandwidth. H100 ≈ 296 FLOP/byte. Below it → memory-bound.
Energy ratios (Horowitz, ISSCC 2014): a DRAM read costs ~20,000× an 8-bit integer add. This is why quantization is an energy lever first.
Numbers to know
| Typical large-scale MFU | 35–55% (not a failure — an accounting identity) |
| MFU vs HFU | HFU counts recomputation as useful work; always ≥ MFU. Ask which. |
| 0.01 nats of loss | ≈ 33% more compute at frontier scale |
| Irreducible fraction of loss | ~87% at 70B/1.4T — the industry fights over the rest |
| Ladder cost | <2% of flagship buys ~5 decades of log C spread |
| Design placement | spread beats density: ~26× lower extrapolation variance, same cost |
| Goodput swing | 64% vs 94% on identical hardware = 12 days of a 40-day run |
$ per 1e21 FLOPs | ~$1,750 at $2.50/H100-hour, 40% MFU |
C = 6ND — where it breaks
| Break | Error | Fix |
|---|---|---|
| Attention at long context | ~4% at 2k, ~15% at 8k, ~245% at 128k, ~20× at 1M | add 12·L·T·d_model per token |
| MoE | 10–20× | use active parameters, not total |
| Embeddings | up to 89% at small N | report non-embedding N |
| Activation checkpointing | 6ND → ~8ND hardware FLOPs | that gap is MFU vs HFU |
MoE
capacity = capacity_factor × tokens × k / E
L_balance = E · Σ f_i · P_i f = token-slot fraction (discrete)
P = mean router probability (differentiable)
L_z = mean( logsumexp(logits)² )
L_total = L_task + α·L_balance + γ·L_z α ≈ 0.01, γ ≈ 1e-3
| Rule | |
|---|---|
Parameters scale with E; FLOPs scale with k | that is MoE, entire |
L_balance = 1.0 means perfectly balanced | larger is worse — the dashboard number |
top_k == n_experts | sparsity 1.0 — a dense model in an MoE costume |
| Active → FLOPs; total → memory | off by 10–20× if swapped |
| Dropped tokens are silent | they ride the residual; no error is raised |
max_over_mean > 3 | the router is collapsing — intervene now |
| Expert-parallel comms | 2 all-to-alls per layer; ~32 GB / ~2.6 s for an 8k prefill on 60 layers |
Prefill vs decode — never confuse them
| Prefill | Decode | |
|---|---|---|
| Processes | the whole prompt at once | one token at a time |
| Bound by | compute | memory bandwidth |
| Arithmetic intensity | high | 1–43 FLOP/byte (far below the ridge) |
| Optimization | pipelining, chunking | batching, GQA, quantization, speculation |
| Parallelism | pipelined prefill works | pipelining just adds serial hops |
This asymmetry is why prefill/decode disaggregation exists, and why the Flash 2.0 fix was pipelined prefill specifically.
The scaling-law workflow
1. FIX the recipe (how depth/LR/batch scale with N, D) — the law describes THIS
2. Design the ladder: geometric budgets, maximize log-C spread, replicate one point
3. IsoFLOPs: fix C, sweep N, fit a parabola in log N, take the vertex
4. Fit power laws N_opt ∝ C^a, D_opt ∝ C^b → CHECK a + b ≈ 1
5. Or fit L = E + A/N^α + B/D^β — Huber loss, log-space residuals
6. Bootstrap a confidence interval (nonlinear model, unknown noise model)
7. Compare baseline vs candidate laws AT THE TARGET C; report the CROSSOVER
8. State the extrapolation distance out loud: "3 decades beyond our largest ablation"
Decision gate: a delta smaller than your confidence interval is not a result.
Kaplan vs Chinchilla, in one box
| Kaplan (2020) | Chinchilla (2022) | |
|---|---|---|
N_opt ∝ | C^0.73 | C^0.5 |
| 10× compute → | 5.37× params, 1.86× data | ~3.2× each |
| Method | one run per size, intermediate losses | separate runs, properly decayed LR |
| Result | models were UNDERTRAINED | smaller models, trained longer |
The bug: reading loss mid-run is a biased estimator, because much of the improvement comes
from the LR decay at the end. A uniform bias would be absorbed into E and change nothing;
it is the non-uniformity across the ladder that tilts the exponents.
Tokens per parameter, historically: GPT-3 1.7 → Chinchilla 20 → Llama-3-70B 214 → Llama-3-8B 1875. A 1000× swing driven first by a methodology fix, then by serving economics.
Feinberg's three verticals
| Vertical | The problem | The hard part |
|---|---|---|
| Distillation | push teacher statistics into a student | storage: 10T tokens × 256k vocab = 5 EB. Top-k or online. |
| Inference co-design | shapes that saturate every hardware unit | quality vs MFU pull opposite ways; differs per chip generation |
| Quantization | FP32 → 4 bits | outliers; power is ~99% of TCO |
Distillation in one sentence: variance reduction — a better teacher reduces bias.
Interview one-liners
- "Backward is exactly 2× forward" — each forward matmul becomes two backward matmuls,
dX = dY·WᵀanddW = Xᵀ·dY, each the same size. - "35% MFU is an accounting identity" — the matmul unit idles during vector ops, memory traffic, collectives and the optimizer step. Read the breakdown; it is your agenda.
- "Active for FLOPs, total for memory."
- "I wouldn't train both and compare" — fit a law for each, evaluate at the target
C, report the delta with intervals plus the crossover, because curves cross. - "That difference is inside our error bars."
- "Prefill is compute-bound, decode is memory-bound" — never apply one's optimization to the other.
- "Power is the bill" — a DRAM read costs ~20,000× an integer add.
- "A scaling law describes your recipe, not the universe."
The five things that will actually kill your run
- Loss spike — decide in minutes: skip the batch, lower LR, or roll back.
- Silent data corruption — no crash, model quietly degrades. Cross-replica checksums.
- Data-iterator bug — wrong shard or an off-by-one. Days of compute on the wrong data.
- Router collapse (MoE) — watch
max_over_mean; above 3 is an emergency. - Slow checkpoints — cadence and restart speed are worth ~12 days of a 40-day run.
Phase 00 — The FLOPs & Memory Algebra of Pre-Training
The phase where "we should train a bigger model" stops being an opinion and becomes an equation. Feinberg opens his Princeton talk with a question — "If I give you a certain amount of compute C (e.g. 1000 H100 for 30 days), what is the best LLM you can train? What should be its size (=N)? How many tokens (=D) should it be trained on?" — and immediately hands you the tool:
C = 6NDis a very good approximation of FLOPs. This phase makes that equation, its derivation, and its four failure modes yours permanently.
Why this phase exists
Every downstream phase is an application of this algebra:
- Phase 01 fits a scaling law over
C, which you compute with6ND. - Phase 02 optimizes lifetime cost =
6ND(train) +2N·D_inf(serve). - Phase 03 asks what
Neven means for an MoE (answer: active, not total). - Phase 05's roofline is this arithmetic divided by hardware constants.
- Phase 08's power model is this arithmetic multiplied by joules per bit.
If 6ND is a formula you memorized rather than derived, all of that becomes cargo cult. So
this phase does the derivation properly, then builds the calculator.
There is a second reason, and it is the practical one: this is what a frontier interview
actually opens with. "You have 1,000 H100s for 30 days. What do you train?" A candidate who
reaches for a framework has already failed. A candidate who converts chips × days into FLOPs,
FLOPs into (N, D), and then sanity-checks against HBM capacity has demonstrated the entire
job in four minutes.
Concept map
┌──────────────────────┐
│ BUDGET │
│ chips × days × MFU │
└──────────┬───────────┘
│ × peak FLOP/s
▼
┌──────────────────────┐
│ C (total FLOPs) │
└──────────┬───────────┘
│ C = 6ND
┌─────────────┴─────────────┐
▼ ▼
┌──────────────────┐ ┌──────────────────┐
│ N (parameters) │ │ D (tokens) │
└────────┬─────────┘ └──────────────────┘
│
┌──────────┼──────────────┬────────────────────┐
▼ ▼ ▼ ▼
weights optimizer activations KV cache
2N bytes states (checkpointing) at serve time
(bf16) 6-12N bytes 2·L·H_kv·d_h·S·B·2
│
▼
┌────────────────────────────────┐
│ DOES IT FIT? (HBM capacity) │
│ if no → shard (Phase 04) │
└────────────────────────────────┘
What you will be able to do
- Derive
2Nforward /4Nbackward /6Ntotal from the definition of a matmul, in under two minutes, on a whiteboard. - Compute the exact per-step FLOP count of a transformer from its shapes — reproducing
Feinberg's slide identity
18BTDF + 24BTDNH = 6·BT·(3DF + 4DNH)— and say which term is the MLP and which is attention. - State the four regimes where
6NDbreaks (attention at long context, MoE, embeddings, activation checkpointing) and estimate the error in each. - Compute training memory: weights + gradients + optimizer states + activations, and say which term dominates at which scale.
- Compute KV-cache size and show how GQA changes it by the group factor.
- Convert between: chips × days ↔ FLOPs ↔
(N, D)↔ dollars ↔ joules. - Answer "1,000 H100 for 30 days — what do you train?" with a number and a defence.
The lab
| Lab | What you build |
|---|---|
| Lab 01 — Transformer FLOPs, Memory & Budget Calculator | A calculator that goes from a hardware budget to a recommended (N, D) and back, with exact per-shape FLOP accounting, MoE-aware active-parameter handling, full training-memory breakdown, KV-cache sizing with GQA, and a fit-check against HBM capacity |
Success criteria. LAB_MODULE=solution pytest test_lab.py -v is green; your lab.py goes
green after the TODOs; python solution.py prints a budget report for a real cluster
configuration that you can defend line by line.
Deliverables checklist
-
I can derive
6NDwithout notes. - I can name the four places it breaks and bound the error.
-
I can compute per-step FLOPs from
(B, T, d_model, d_ff, n_heads, d_head, n_layers). - I can compute the full training memory footprint for Adam in mixed precision.
- I can compute a KV cache and show the GQA saving.
- I have run the calculator on a cluster I could plausibly be given, and written down what I would train and why.
Key takeaways
- A matmul costs
2 × (number of weights)FLOPs per token. Everything else follows. - Backward is exactly 2× forward because each forward matmul becomes two backward matmuls
(
dXanddW). 6NDexcludes attention. Fine at 2k context, badly wrong at 128k.- For MoE,
Nis active parameters. Total parameters determine memory, not FLOPs. - Memory, not FLOPs, is usually what stops you. Adam in mixed precision costs ~16 bytes per parameter before you have stored a single activation.
- The KV cache is the serving wall, and
kv_headsis the lever that moves it.
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
6NDDerivation - Chapter 6: The Exact Per-Step Count
- Chapter 7: Where
6NDBreaks - 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:
| Term | Meaning | Example |
|---|---|---|
| FLOPs (plural noun) | a quantity of work | "this run needs 10²⁴ FLOPs" |
| FLOP/s (rate) | speed — operations per second | "an H100 does 10¹⁵ FLOP/s" |
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.
| Format | Bits | Range | Precision | Typical use |
|---|---|---|---|---|
| FP32 | 32 | huge | ~7 decimal digits | optimizer states, loss accumulation |
| TF32 | 19 (in 32 slots) | FP32 range | ~3 digits | NVIDIA matmul default |
| BF16 | 16 | FP32 range | ~3 digits | the workhorse for training |
| FP16 | 16 | small | ~3 digits | older training, overflows easily |
| FP8 | 8 | small | ~1 digit | frontier training/inference |
| INT4 | 4 | integer codes | — | weight-only inference |
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 × nentries. - Each entry is a sum of
kproducts. - 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
2in2mknis 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)
| Symbol | Name | Typical value | What it is |
|---|---|---|---|
d or d_model | model dimension | 4096–16384 | width of the residual stream |
d_ff or F | FFN hidden dim | ~4·d (or ~2.7·d for SwiGLU) | width of the MLP's middle |
n_h or N_heads | query heads | 32–128 | parallel attention heads |
n_kv | key/value heads | 1–128 | fewer than n_h under GQA/MQA |
d_h or H | head dimension | 64–128 | usually d / n_h |
L | layers | 32–120 | depth |
V | vocabulary | 32k–256k | number of distinct tokens |
B | batch size | — | sequences processed together |
T | sequence length | 2k–1M | tokens per sequence |
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_hto keep variance stable (without this, larged_hproduces 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 tod.- 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·kFLOPs.dW = Xᵀ · dYis(k × T) @ (T × n)→2·k·T·nFLOPs.
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
2mknis 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
| Quantity | FLOPs per token | When you use it |
|---|---|---|
| Forward only | 2N | inference / prefill; teacher forward in distillation |
| Backward | 4N | — |
| Full training step | 6N | pre-training budgets |
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:
- 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.
- 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 = 6DFFLOPs per token - Backward:
2 ×that= 12DF - Total:
18DFper 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:
24DNHper 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_seqterm 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:
matmul_flops(m, k, n)— the2mknprimitive. Everything else calls this.params_per_layer(...)/total_params(...)— withgated, GQA, and tied-embedding handling, returning the non-embedding count separately.training_flops_exact(...)— reproducing18BTDF + 24BTDNH, with the sequence-dependent attention term as a separate line item so you can see Break 1.moe_parameter_split(...)— total vs active, which is Break 2.training_memory(...)— weights, grads, optimizer states, activations, ZeRO stages.kv_cache_bytes(...)/max_concurrent_requests(...)— the serving wall.budget_to_flops(...)/chinchilla_split(...)— the round trip.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.pypasses after filling the TODOs. -
python solution.pyprints a full budget report. -
You can derive
6NDon 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
6NDand 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
Nfor scaling work. Say so explicitly; it signals you have actually done this. - Carry
$ per 1e21 FLOPsin your head. It turns research proposals into money.
Takeaways
- A matmul costs
2 × paramsper token. Every other number in this phase descends from that. - Backward is exactly 2× forward, for a structural reason.
2Ninference,6Ntraining. Never mix them.6NDexcludes attention: ~3% error at 2k context, 22× at 1M.- For MoE: active for FLOPs, total for memory.
- Memory, not compute, is the binding constraint. Adam mixed-precision is ~16 bytes/param.
- The KV cache is the serving wall;
n_kv_headsis the lever, and it is set before training. - Decode is memory-bound at every realistic batch size. That one fact explains most of serving.
- Chinchilla-optimal is a starting point, not an answer. Serving economics move it.
References
Primary
- Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides (the
C = 6NDslide and its footnote; the18BTDF + 24BTDNHidentity) - 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)
Lab 01 — Transformer FLOPs, Memory & Budget Calculator
Build the napkin. By the end you can answer "I give you 1,000 H100s for 30 days — what do you train?" with a number, a memory check, a serving check, and a dollar figure.
The problem
Feinberg opens his Princeton talk with exactly that question, and hands you one tool:
C = 6ND. But the equation alone is not enough to answer it. You also need to know whether the
answer fits in memory, whether you have that many tokens, whether the resulting model is
servable, and what it costs. That whole chain is one calculator, and this lab is it.
Every downstream phase calls back into this arithmetic: Phase 01 fits scaling laws over C,
Phase 02 optimizes 6ND + 2N·D_inf, Phase 03 needs active-vs-total parameters, Phase 05
divides these numbers by hardware constants to get a roofline.
What you build
| Group | Functions | The idea |
|---|---|---|
| The primitive | matmul_flops, padded_matmul_flops | 2mkn; and what tile padding costs you |
| Parameters | params_per_layer, total_params, moe_parameter_split | GQA/MQA, gated MLPs, the embedding trap, active vs total |
| FLOPs | training_flops, inference_flops, training_flops_exact, unembedding_flops, attention_flop_fraction | 6ND, 2N, the exact per-shape count, and where 6ND breaks |
| Memory | activation_bytes, training_memory, kv_cache_bytes, max_concurrent_requests, decode_arithmetic_intensity, ridge_point | Adam's 16 bytes/param, ZeRO stages, the serving wall |
| Budgets | budget_to_flops, flops_to_days, chinchilla_split, lifetime_flops, cost_report, budget_report | chips × days ↔ FLOPs ↔ (N, D) ↔ dollars ↔ watts |
Key concepts
| Concept | Why it is in this lab |
|---|---|
2mkn | A linear layer costs 2 × params per token. The root of everything. |
6ND | Forward 2N + backward 4N. The currency conversion of the field. |
18BTDF + 24BTDNH | The slide identity — 6ND written out in shapes. A test asserts it. |
| Active vs total parameters | MoE: active for FLOPs, total for memory. Off by 13× if you swap them. |
Non-embedding N | A 256k vocab is 89% of a small model. Scaling ladders must exclude it. |
| 16 bytes/param | Mixed-precision Adam, before a single activation. |
| ZeRO stages | Optimizer states are 75% of training memory — shard those first. |
| KV cache | 2·L·n_kv·d_h·T·B·b. The serving wall; n_kv is the lever. |
| Ridge point | peak FLOP/s ÷ HBM bandwidth. Below it you are memory-bound. |
| Lifetime FLOPs | 6N·D_train + 2N·D_inf. Chinchilla's blind spot. |
Files
| File | What it is |
|---|---|
| lab.py | Your implementation. Signatures, docstrings and validation contracts are given; the arithmetic is yours. |
| solution.py | Reference. python solution.py prints a twelve-part worked example. |
| test_lab.py | 75 tests: happy path, validation, boundaries, invariants, determinism. |
| requirements.txt | pytest only. Pure stdlib otherwise. |
Run
pytest test_lab.py -v # your lab.py — red until you implement
LAB_MODULE=solution pytest test_lab.py -v # the reference — must be green (75 passed)
python solution.py # the worked example
Where to start
matmul_flopsand gettest_matmul_flops_is_2x_weight_countpassing. That single identity is the root of6ND; if it is not obvious to you yet, re-read WARMUP Chapter 2.params_per_layer, thentest_gqa_shrinks_only_kv_projections. It forces you to notice that GQA touchesW_k/W_vand leavesW_q/W_oalone.training_flops_exact, thentest_exact_step_reproduces_the_slide_identity. This is the money test — it checks your per-shape accounting against Feinberg's slide.- Everything else is bookkeeping.
The trap is units. Bytes vs GB, FLOPs vs FLOP/s, per-token vs per-step vs per-run, active vs total. Every signature names its unit; read them.
Success criteria
-
LAB_MODULE=solution pytest test_lab.py -v→ 75 passed. -
Your
lab.pyreaches 75 passed. -
python solution.pyruns and you can explain every one of its twelve sections. -
Without the calculator, you can derive
6NDon paper in under two minutes. -
You can state the four places
6NDbreaks and the error in each. -
You have run
budget_reporton a cluster you might plausibly be given, and written a paragraph on what you would train — including where you would deliberately deviate from Chinchilla and why.
How this maps to the real stack
| This lab | The real thing | Where the miniature lies |
|---|---|---|
training_flops_exact | torch.utils.flop_counter.FlopCounterMode; JAX cost_analysis(); NVIDIA Nsight | Real counters trace the actual graph, so they catch fused ops, custom kernels, and recomputation. Ours assumes a textbook architecture. |
training_memory | torch.cuda.max_memory_allocated(); DeepSpeed's memory estimator; torch.distributed.fsdp | Ignores allocator fragmentation, NCCL buffers, CUDA context (~1 GB), and per-framework overheads. Real usage runs 10–30% above this. |
kv_cache_bytes | vLLM's PagedAttention block manager; TensorRT-LLM's KVCacheManager | Real engines page the cache in fixed blocks, so the true footprint rounds up to block granularity and can share prefixes across requests. |
chinchilla_split | Nobody ships a library for this — every lab has an internal version | Uses a fixed tokens_per_param ratio. Real practice fits the ratio from your own ladder (Phase 01) because it is recipe-dependent. |
ridge_point / decode_arithmetic_intensity | The roofline analysis in Nsight Compute; llm-analysis | Assumes perfect bandwidth utilization. Real kernels achieve 60–90% of peak HBM bandwidth. |
HARDWARE table | Vendor spec sheets | Peak numbers are marketing maxima at ideal clocks. Sustained throughput under thermal load is 5–15% lower. |
What is not a lie: the FLOP arithmetic. 2mkn, 6ND, 18BTDF + 24BTDNH, and the
KV-cache formula are exact. That is why they are what interviewers ask about.
Extensions
For your own hardware, or for a portfolio piece:
- Validate against a real model. Load a small HF checkpoint, count its parameters with
sum(p.numel() for p in model.parameters()), and checktotal_paramsmatches. Then runFlopCounterModeon a forward pass and checkinference_flops. Any mismatch is a lesson. - Add pipeline/tensor-parallel memory. Extend
training_memorywithtp_degreeandpp_degree, and model the activation memory that pipeline stages must hold in flight. Phase 04 will need it. - Add an MoE memory model. Expert-parallel sharding puts different experts on different
chips — extend
moe_parameter_splitto report per-device memory givenep_degree. - Fit the tokens-per-param ratio. Instead of hardcoding 20, take a list of
(N, D, loss)points and solve for the ratio that minimizes loss at fixedC. That is Phase 01, and doing it here first makes Phase 01 trivial. - Plot the roofline. Sweep batch size and context length, and plot arithmetic intensity against the ridge point to find exactly where decode becomes compute-bound.
Interview / resume bullets
- "Derived and implemented the full FLOP and memory accounting for transformer pre-training —
C = 6NDwith exact per-shape decomposition (18BTDF + 24BTDNH), MoE active-vs-total parameter handling, mixed-precision Adam memory with ZeRO stages 0–3, and KV-cache sizing under GQA/MQA — and used it to size training runs against real accelerator budgets." - "Built a compute-budget planner converting chips × days → FLOPs → Chinchilla-optimal
(N, D)→ dollars and watts, with automated feasibility checks for cluster memory, corpus size and single-chip serving; quantified the lifetime-cost crossover where an overtrained smaller model beats the compute-optimal one." - Interview-ready: "backward is exactly 2× forward because each forward matmul becomes two
backward matmuls,
dX = dY·WᵀanddW = Xᵀ·dY, each the same size as the forward one."
Phase 01 — Scaling Laws I: Kaplan → Chinchilla → IsoFLOPs
The phase where you learn to predict the result of an experiment you cannot afford to run twice. Feinberg's framing of why this discipline exists is the sharpest available: old ML iterated on CIFAR-10 before ImageNet and did LR searches "by doing multiple final runs" — "the last data point is our test set!" — whereas now "every single time you go up for a pre-training run, you're about to put in more FLOPs into this run than you've ever done before." Every next run requires extrapolation.
Why this phase exists
This is the intellectual centre of a pre-training lead's job. Everything else — MoE, sharding,
distillation, quantization — is machinery in service of one decision: given C, what should
we build? And that decision is made before the money is spent, on the strength of a fitted
curve.
Three reasons this is genuinely hard, and not just curve-fitting:
- You are extrapolating three orders of magnitude beyond your largest data point. Any functional form fits your ladder; they diverge wildly at the target.
- The law describes your recipe, not nature. His slide is explicit: "Analysis made in the
context of a parameterized LLM training recipe! Must already have architecture scaling,
schedule defined for
N,D." Change how depth scales with width, or how LR decays, and you must refit. - The most consequential result in the field came from fixing an experimental-design bug. Chinchilla beat Kaplan not with a better idea but by noticing that reading loss mid-run is a biased estimator. That is the lesson: methodology is the frontier.
And the payoff line from his slides, which is the whole reason to care: "Loss forecast implies model/recipe selection capability!" Forecasting is not reporting. It is the decision procedure.
Concept map
┌───────────────────────────────────────────────────────────────────┐
│ A PARAMETERIZED RECIPE │
│ how depth/width/LR/batch/warmup all scale with N and D │
│ (without this fixed, a "scaling law" means nothing) │
└────────────────────────────┬──────────────────────────────────────┘
│
▼
┌───────────────────────────────────────────────────────────────────┐
│ THE LADDER — many cheap runs │
│ │
│ IsoFLOPs or Parametric fit │
│ ───────── ────────────── │
│ fix C, sweep N collect (N, D, L) │
│ loss vs log N is a U fit L = E + A/N^a + B/D^b│
│ fit a parabola, take the vertex (Huber on log L) │
│ repeat for several C │
│ fit power laws N*(C), D*(C) │
└────────────────────────────┬──────────────────────────────────────┘
│
▼
┌───────────────────────────────────────────────────────────────────┐
│ EXTRAPOLATE to the flagship C, with a confidence interval │
└────────────────────────────┬──────────────────────────────────────┘
│
▼
┌───────────────────────────────────────────────────────────────────┐
│ DECIDE: compare baseline law vs candidate law AT THE TARGET C │
│ report the crossover point, not just the winner │
└───────────────────────────────────────────────────────────────────┘
What you will be able to do
- State the power-law form
L = E + A/N^α + B/D^βand say what each of the five parameters means physically (E= irreducible entropy of the data;A/N^α= capacity limit;B/D^β= data limit). - Run an IsoFLOPs sweep: fix
C, varyN, fit a parabola inlog N, extract the vertex. - Fit power laws
N_opt ∝ C^a,D_opt ∝ C^band checka + b ≈ 1(which it must, sinceC = 6ND). - Explain precisely why Kaplan got
a ≈ 0.73and Chinchilla gota ≈ 0.5— the LR-decay bias — and demonstrate the bias numerically. - Fit the parametric form with a robust loss on
log L, and show that least-squares on rawLgives different (worse) exponents. - Produce a forecast with a confidence interval, via bootstrap over the ladder points.
- Compare a baseline recipe against a candidate recipe as fitted laws at the target
C, and report the crossover. - Choose ladder design points to minimize extrapolation variance instead of filling a grid.
The lab
| Lab | What you build |
|---|---|
| Lab 01 — IsoFLOPs Ladder, Law Fitting & the Forecast | A complete scaling-law workbench: a synthetic ladder generator with configurable noise and an injectable LR-decay bias, a parabola fitter, power-law and parametric fitters (least-squares vs Huber-on-log), bootstrap confidence intervals, a baseline-vs-candidate decision rule, and an optimal-design scorer for choosing where to spend your next ablation |
Success criteria. LAB_MODULE=solution pytest test_lab.py -v green; python solution.py
reproduces the Kaplan→Chinchilla story end to end from synthetic data, including the moment
where correcting the bias moves the recommended N by ~2×.
Deliverables checklist
-
I can write
L = E + A/N^α + B/D^βand explain all five parameters. - I have run an IsoFLOPs sweep and extracted a vertex from a parabola fit.
- I can explain the Kaplan bug in one paragraph, and I have reproduced it numerically.
-
I know why
a + b ≈ 1is a consistency check I should always run. - I can produce a forecast with error bars and say what drives the width.
-
I have compared two recipes at a target
Cand reported the crossover. - I can justify where to place my next ladder run in terms of information gain.
Key takeaways
- A scaling law is a property of a recipe. Fix the recipe first; the law describes that.
- Loss forecasting is recipe selection. Compare fitted curves at the target, not runs at small scale.
- The U-curve at fixed
Cis real and its bottom is flat. Flatness is a gift (you can move off the optimum for serving reasons cheaply) and a trap (noise moves the fitted vertex a lot). a + b ≈ 1falls out ofC = 6ND. If your fitted exponents violate it, you have a bug.- Reading loss mid-run is biased, systematically against training on more data. This one methodological error set the industry's strategy for two years.
- The estimator matters. Least-squares on
L, least-squares onlog L, Huber onlog L, and MLE give different exponents on identical data — and therefore different flagship recommendations. Feinberg's own listed open problem: "Formalize." - Where you place ladder points matters more than how many you run. Spread beats density.
Warmup Guide — Scaling Laws I: Kaplan → Chinchilla → IsoFLOPs
How to read this. No statistics background is assumed. Power laws, log-log plots, curve fitting, robust losses, bootstrap confidence intervals and experimental design are all built from nothing. If you already know regression, start at Chapter 4 and do not skip Chapter 5 or Chapter 9.
Table of Contents
- Chapter 1: What a Scaling Law Is
- Chapter 2: Power Laws From Nothing
- Chapter 3: What "Loss" Actually Measures
- Chapter 4: The Kaplan Result, and What It Did to the Industry
- Chapter 5: The Chinchilla Correction — a Bug in Experimental Design
- Chapter 6: The IsoFLOPs Method, Step by Step
- Chapter 7: The Parametric Form
L(N, D) - Chapter 8: Fitting It — Least Squares, Logs, and Huber
- Chapter 9: The Estimator Problem — Feinberg's Open Question
- Chapter 10: Confidence Intervals via Bootstrap
- Chapter 11: Designing the Ladder — Where to Spend Your Ablations
- Chapter 12: Making the Decision — Baseline vs Candidate
- Chapter 13: When Scaling Laws Lie
- Lab Walkthrough
- Success Criteria
- Interview Q&A
- Tips & Takeaways
- References
Chapter 1: What a Scaling Law Is
What it is
A scaling law is an empirical formula that predicts model quality from the resources you spend. In its most useful form:
$$ L(N, D) = E + \frac{A}{N^{\alpha}} + \frac{B}{D^{\beta}} $$
L— the test loss (Chapter 3 explains what that number means)N— parametersD— training tokensE,A,B,α,β— five constants you fit from a ladder of small experiments
Why it exists
Because of the one-shot problem. Feinberg's slides state the contrast exactly:
ML training before: "Maybe 2 stages; toy problem for iteration (CIFAR10) then you apply to Imagenet. LR searches by doing multiple 'final runs'. The last data point is our test set!"
Now: "every next run requires extrapolation."
You cannot try five flagship configurations and pick the winner. Each costs eight figures and months. So you build a predictor from cheap experiments and bet on its extrapolation.
How it works underneath — why this functional form?
The three terms are not arbitrary. Each corresponds to a distinct source of error:
| Term | Name | Meaning | What drives it to zero |
|---|---|---|---|
E | irreducible loss | The entropy of language itself. Even a perfect model cannot predict the next token with certainty. | Nothing. It is a floor. |
A/N^α | capacity term | Your model is too small to represent the true function. | More parameters. |
B/D^β | data term | You have not seen enough examples to find the right function. | More tokens. |
This decomposition is the classic approximation / estimation split from learning theory,
written for transformers. And it makes a strong, checkable prediction: at fixed C, there is a
single interior optimum, because pushing N up shrinks the capacity term but (via C = 6ND)
shrinks D and grows the data term. Chapter 6 is that trade-off made visual.
Production significance
Everything. The flagship (N, D), the go/no-go, the choice between two architectures, the
decision to buy more chips — all of it runs through a fitted law. Which is why Feinberg's line
lands so hard: "Loss forecast implies model/recipe selection capability!"
The misconception
"Scaling laws are laws of nature, like
F = ma."
They are not. They are fitted regressions over a specific recipe on a specific dataset. His slide says it twice: "These 'laws' are only empirical" and "The fitting of these laws depends a lot on the experimental setup as well as the implicit assumptions being made there." Change your data mixture, your optimizer, or how your architecture scales, and the constants move — and sometimes the exponents do too.
Chapter 2: Power Laws From Nothing
What a power law is
A relationship of the form y = c · x^k. That is it. The magic is what it looks like in logs:
$$ \log y = \log c + k \log x $$
A power law is a straight line on a log-log plot. The slope of that line is the exponent. That single fact is why every scaling-law paper's figures are log-log, and why you can eyeball an exponent from a plot.
import math
def power_law(x, c, k):
return c * (x ** k)
# Confirm the straight-line property.
for x in (1e18, 1e19, 1e20, 1e21):
y = power_law(x, c=0.6, k=0.5)
print(f"log10(x)={math.log10(x):5.1f} log10(y)={math.log10(y):7.3f}")
log10(x)= 18.0 log10(y)= 8.778
log10(x)= 19.0 log10(y)= 9.278
log10(x)= 20.0 log10(y)= 9.778
log10(x)= 21.0 log10(y)= 10.278
Every step of 1 in log x moves log y by exactly 0.5 — the exponent. Perfectly straight.
Fitting one: linear regression in log space
If your data follows a power law, take logs and fit a line. Ordinary least squares, in eight lines, no libraries:
def fit_line(xs, ys):
"""Least-squares slope and intercept. Returns (slope, intercept)."""
n = len(xs)
if n < 2:
raise ValueError("need at least 2 points")
mx = sum(xs) / n
my = sum(ys) / n
sxx = sum((x - mx) ** 2 for x in xs)
if sxx == 0:
raise ValueError("all x values identical — slope is undefined")
sxy = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
slope = sxy / sxx
return slope, my - slope * mx
def fit_power_law(xs, ys):
"""Fit y = c * x^k by regressing log y on log x. Returns (c, k)."""
if any(x <= 0 or y <= 0 for x, y in zip(xs, ys)):
raise ValueError("power-law fit needs strictly positive data")
lx = [math.log(x) for x in xs]
ly = [math.log(y) for y in ys]
k, log_c = fit_line(lx, ly)
return math.exp(log_c), k
xs = [1e18, 1e19, 1e20, 1e21, 1e22]
ys = [power_law(x, 0.6, 0.5) for x in xs]
c, k = fit_power_law(xs, ys)
print(f"recovered c={c:.4f} k={k:.4f}") # recovered c=0.6000 k=0.5000
Why power laws are everywhere in deep learning
Nobody has a fully satisfying answer, and it is an active research question. The best current intuitions:
- Data manifold dimension. If natural language lies near a
d-dimensional manifold, covering it to resolutionεneeds~ε^{-d}samples — which produces a power law with exponent related to1/d. - Heavy-tailed feature frequency. Word and concept frequencies follow Zipf's law (itself a
power law). Learning the
n-th most common concept requires seeing it, and the marginal return of more data falls off as a power. - Random-feature / kernel arguments give power-law generalization curves under mild assumptions.
The honest position: the power-law form is an excellent empirical description over the ranges we have measured, with plausible but unproven theoretical motivation. Treat it as a very good interpolant, not a physical law. Chapter 13 is about what happens when it breaks.
The misconception
"A straight line on a log-log plot proves a power law."
Over a small range, almost anything looks straight in logs. You need several decades of x
to distinguish a power law from a log, a stretched exponential, or a saturating curve. This
matters enormously, because those alternatives disagree violently when extrapolated — which is
exactly what you are about to do.
Chapter 3: What "Loss" Actually Measures
What it is
Language models are trained with cross-entropy loss: the negative log-probability the model assigned to the token that actually came next, averaged over tokens.
$$ L = -\frac{1}{T}\sum_{t=1}^{T} \log p_\theta(x_t \mid x_{<t}) $$
def cross_entropy(probs_assigned_to_true_token):
"""Mean negative log-probability. Units: nats (natural log)."""
if any(p <= 0 for p in probs_assigned_to_true_token):
raise ValueError("zero probability on a true token gives infinite loss")
return -sum(math.log(p) for p in probs_assigned_to_true_token) / len(
probs_assigned_to_true_token)
print(round(cross_entropy([0.9, 0.8, 0.95, 0.7]), 4)) # 0.2113 (confident, correct)
print(round(cross_entropy([0.1, 0.2, 0.05, 0.3]), 4)) # 2.0369 (bad)
Why this and not accuracy
Three reasons, all of which matter for scaling work:
- It is smooth. Accuracy is a step function — the gradient is zero almost everywhere. Cross-entropy has a gradient everywhere.
- It is graded. A model that puts 0.49 on the right token is genuinely better than one that puts 0.01, even though both are "wrong." Accuracy cannot see that.
- It scales predictably. This is the empirical gift: cross-entropy follows clean power laws. Downstream benchmark accuracy does not — it plateaus, jumps, and saturates.
The units, and how to read them
Loss in nats (natural log). Two conversions worth knowing:
def nats_to_bits(nats): return nats / math.log(2)
def loss_to_perplexity(nats): return math.exp(nats)
for L in (2.0, 1.8, 1.7, 1.6):
print(f"loss {L:.2f} nats = {nats_to_bits(L):.3f} bits/token, "
f"perplexity {loss_to_perplexity(L):.2f}")
loss 2.00 nats = 2.885 bits/token, perplexity 7.39
loss 1.80 nats = 2.597 bits/token, perplexity 6.05
loss 1.70 nats = 2.453 bits/token, perplexity 5.47
loss 1.60 nats = 2.308 bits/token, perplexity 4.95
Perplexity is e^L — read it as "the model is as confused as if it were choosing uniformly
among this many options."
The scale intuition you must build
Loss differences look tiny and are not. 0.01 nats is a meaningful, fought-over improvement at frontier scale. Why:
def flops_for_loss_delta(delta_nats, alpha=0.5):
"""Under L ~ C^-alpha near loss ~1.7, roughly how much more compute buys delta?"""
# dL/dC = -alpha * (L - E) / C => fractional compute increase ~ delta / (alpha*(L-E))
L, E = 1.75, 1.69
return delta_nats / (alpha * (L - E))
for d in (0.001, 0.01, 0.05):
print(f"{d:5.3f} nats needs ~{flops_for_loss_delta(d):5.1%} more compute")
0.001 nats needs ~ 3.3% more compute
0.010 nats needs ~33.3% more compute
0.050 nats needs ~166.7% more compute
0.01 nats is roughly a third more compute. On a $30M run that is $10M. That is why teams argue about the third decimal place, and why your fit's confidence interval had better be narrower than the effect you are claiming.
The misconception
"Lower loss always means a better product."
Loss is measured on your held-out distribution. A model can improve loss by getting better at boilerplate that dominates the corpus while getting worse at the reasoning your users care about. This is why every serious team pairs the scaling law with downstream evals — and why Feinberg's slide on refinements mentions "Joint loss, eval fit." The law forecasts loss; loss is a proxy; the proxy needs auditing.
Chapter 4: The Kaplan Result, and What It Did to the Industry
The finding
Kaplan et al. (2020) established that transformer loss follows clean power laws in N, D, and
C, and — the consequential part — derived the compute-optimal allocation. Feinberg's slide
quotes their conclusion directly: data requirements grow "very slowly as D ∼ C^0.27 with
training compute," and states their headline result:
"With a 10x compute budget, parameters should increase by 5.37x and the amount of data by 1.86x."
Decoding those numbers
def kaplan_allocation(compute_multiplier):
"""Kaplan: N grows as C^0.73, D as C^0.27."""
return compute_multiplier ** 0.73, compute_multiplier ** 0.27
n_mult, d_mult = kaplan_allocation(10)
print(f"10x compute -> {n_mult:.2f}x params, {d_mult:.2f}x data")
# 10x compute -> 5.37x params, 1.86x data
print(f"consistency check (must be ~10): {n_mult * d_mult:.2f}")
# consistency check (must be ~10): 10.00
Note that consistency check — N_mult × D_mult must equal the compute multiplier, because
C = 6ND. Equivalently, the exponents must sum to 1: 0.73 + 0.27 = 1.00. Always run this
check on any pair of fitted exponents. If they do not sum to ~1, something is wrong with your
fit or your FLOP accounting.
The industry consequence
His slide states it flatly:
"Consequences for the industry: We should heavily invest in scaling the model size rather than the data size!"
And that is exactly what happened. GPT-3: 175B parameters on ~300B tokens — a ratio of 1.7 tokens per parameter. Compare Chinchilla's later recommendation of ~20, or Llama 3's ~190.
for name, N, D in [("GPT-3 (2020)", 175e9, 300e9),
("Gopher (2021)", 280e9, 300e9),
("Chinchilla (2022)", 70e9, 1.4e12),
("Llama-3-70B (2024)", 70e9, 15e12),
("Llama-3-8B (2024)", 8e9, 15e12)]:
print(f"{name:20s} {D/N:7.1f} tokens per parameter")
GPT-3 (2020) 1.7 tokens per parameter
Gopher (2021) 1.1 tokens per parameter
Chinchilla (2022) 20.0 tokens per parameter
Llama-3-70B (2024) 214.3 tokens per parameter
Llama-3-8B (2024) 1875.0 tokens per parameter
That table is the history of the field in one column. A 1000× swing in a single design parameter, driven first by a methodological error and then by serving economics (Phase 02).
The caveats Kaplan themselves flagged
Feinberg's slide lists them, and the point of listing them is that everyone ignored them:
- "These 'laws' are only empirical"
- "The fitting of these laws depends a lot on the experimental setup as well as the implicit assumptions being made there"
Tip. When a paper flags its own assumptions, that is where the next paper comes from. Chinchilla is literally the second bullet, taken seriously.
Chapter 5: The Chinchilla Correction — a Bug in Experimental Design
This is the most instructive story in the field. Feinberg's slide names the bug precisely:
"Kaplan et al. run a single training run per model size and uses intermediate losses to estimate the loss at different token horizon. ... This is a bad approximation as you can get much better losses through proper learning rate decay. Only the final loss value is optimal."
What a learning-rate schedule is (from zero)
The learning rate controls how big a step the optimizer takes. Too big and you bounce around the minimum forever; too small and you crawl. The standard solution is a schedule: warm up, hold high, then decay toward zero.
LR
│ ╭──────────╮
│ ╱ ╲
│ ╱ ╲___
│ ╱ ╲___
│ ╱ ╲____
│ ╱ ╲___
└─────────────────────────────────────────► training steps
warmup high (explore) decay (settle)
The decay phase is where a large chunk of the final loss improvement happens — the model stops oscillating and settles into the basin. A model mid-run has not had its decay yet.
The bug, made precise
Kaplan wanted L(N, D) for many D. Running a separate full experiment for each D is
expensive. So: train once to a large D, and read the loss curve at intermediate points.
That is a biased estimator. At token count D' mid-run, the LR is still high, so the loss
is worse than it would be for a run scheduled to end at D'.
def loss_curve(tokens_seen, horizon, base_a=3.0, base_b=0.1, decay_penalty=0.15):
"""Toy loss: an intrinsic power-law improvement plus a penalty for undecayed LR.
The penalty vanishes only at the end of the SCHEDULE, not at a fixed token count.
"""
if tokens_seen > horizon:
raise ValueError("cannot read a curve past its horizon")
intrinsic = base_a / (tokens_seen ** base_b)
frac_remaining = 1.0 - tokens_seen / horizon
return intrinsic + decay_penalty * frac_remaining
# The comparison that broke the field.
D_target = 100e9
kaplan_style = loss_curve(D_target, horizon=1000e9) # peek at a long run mid-flight
chinchilla_style = loss_curve(D_target, horizon=D_target) # a run scheduled to end here
print(f"Kaplan-style estimate at D=100B : {kaplan_style:.4f}")
print(f"Chinchilla-style (true) : {chinchilla_style:.4f}")
print(f"bias : {kaplan_style - chinchilla_style:+.4f} nats")
Kaplan-style estimate at D=100B : 0.3733
Chinchilla-style (true) : 0.2383
bias : +0.1350 nats
Why the bias is fatal and not merely noisy
Here is the structural point that makes this a real bug rather than extra noise. The bias is
penalty = decay_penalty × (1 − D_measured / D_schedule)
so it is not the same for every point:
- Points read early in their schedule get a large penalty.
- Points read at the end of their schedule get none.
def bias_across_horizons(decay_penalty):
"""Measure the SAME token count from runs with different schedule lengths."""
D = 100e9
for horizon_mult in (1, 2, 5, 10):
biased = loss_curve(D, D * horizon_mult, decay_penalty=decay_penalty)
truth = loss_curve(D, D, decay_penalty=decay_penalty)
print(f" schedule={horizon_mult:2d}x D -> measured {biased:.4f} "
f"(bias {biased - truth:+.4f})")
print("With the schedule-mismatch bias present:")
bias_across_horizons(0.15)
print("With it removed (each run scheduled to its own horizon):")
bias_across_horizons(0.0)
With the schedule-mismatch bias present:
schedule= 1x D -> measured 0.2383 (bias +0.0000)
schedule= 2x D -> measured 0.3133 (bias +0.0750)
schedule= 5x D -> measured 0.3583 (bias +0.1200)
schedule=10x D -> measured 0.3733 (bias +0.1350)
With it removed (each run scheduled to its own horizon):
schedule= 1x D -> measured 0.2383 (bias +0.0000)
schedule= 2x D -> measured 0.2383 (bias +0.0000)
schedule= 5x D -> measured 0.2383 (bias +0.0000)
schedule=10x D -> measured 0.2383 (bias +0.0000)
Non-uniformity is the whole mechanism, and it is worth being precise about why. A
uniform offset — "all our losses were 0.05 too high" — would simply be absorbed into the
fitted E and would change no exponent and no recommendation. The lab has a dedicated test
for exactly that control (test_uniform_offset_is_absorbed_into_E_and_changes_nothing). It is
because the penalty varies systematically across the (N, D) grid that it tilts the fitted
exponents, and the exponents are what determine the flagship recommendation.
On the direction of the tilt — a note on honesty. Which way the exponent moves depends on which ladder points the mismatch hits hardest, and that is a property of the specific experimental design. Reconstructing Kaplan's design precisely enough to derive the sign from first principles is beyond what the published record settles cleanly; what is documented is the empirical outcome — Chinchilla re-ran with properly matched schedules and found the optimum sits at substantially more tokens per parameter. The lab lets you inject the bias under two different designs (
bias_mode="fixed_schedule"and"per_budget") and measure the tilt each one produces. That exercise — "I do not know the sign, so I will simulate the design and find out" — is the actual skill this chapter is teaching.
The result
Chinchilla's exponent came out at ~0.5 rather than ~0.73. Feinberg's slide:
"Chinchilla findings: the exponent in the power law is ~0.5, meaning model and data size should be scaled at the same rate! This is widely different from Kaplan et al."
And he labels the old regime on the plot with a single word: UNDERTRAINED!
"Consequences: Given a compute budget, models should be smaller and trained for longer. Kaplan's scaling laws meant that models were undertrained — which is obviously bad given bigger models are more expensive to serve and use downstream!"
That last clause is the bridge to the rest of this track. Chinchilla did not just lower the loss — it made the same quality available in a smaller model, which is cheaper to serve. That is the objective Feinberg's team optimizes, and Phase 02 pushes it further.
The empirical demonstration: Chinchilla (70B on 1.4T) beat Gopher (280B on 300B) at the same training compute — with a model 4× cheaper to serve.
Takeaway. The single most consequential result in scaling research was an experimental-design fix. Not an architecture, not an optimizer. Methodology is the frontier, and Feinberg's closing slide says the field still has open problems of exactly this kind (Chapter 9).
Chapter 6: The IsoFLOPs Method, Step by Step
Feinberg's slides walk through this as a build-up of six steps. Here is each one with the code.
The six steps
1. Fix a target FLOPs budget C = 1e20
2. Train a few models, vary model size N = 100M, 300M, 1B, 3B (D = C/6N)
3. Fit a parabola and find the minimum loss vs log N is U-shaped
4. Repeat 1-3 for various FLOPs budgets C = 1e19, 1e20, 1e21, 1e22
5. Fit a power law: FLOPs -> optimal N N_opt ∝ C^a
6. Fit a power law: FLOPs -> optimal D D_opt ∝ C^b
"IsoFLOPs" means "equal FLOPs" — every model on one curve costs the same to train. You are asking: given this exact budget, what shape spends it best?
Step 1–2: the sweep
def true_loss(N, D, E=1.69, A=406.4, alpha=0.34, B=410.7, beta=0.28):
"""The Chinchilla parametric form, used here as a stand-in for reality."""
return E + A / (N ** alpha) + B / (D ** beta)
def isoflop_sweep(C, param_counts):
"""Fix C, vary N, derive D from the budget. Returns [(N, D, loss), ...]."""
out = []
for N in param_counts:
D = C / (6 * N) # the budget constraint
out.append((N, D, true_loss(N, D)))
return out
for N, D, L in isoflop_sweep(1e21, [1e8, 3e8, 1e9, 3e9, 1e10, 3e10]):
print(f"N={N:8.1e} D={D:8.1e} D/N={D/N:8.1f} loss={L:.4f}")
N= 1.0e+08 D= 1.7e+12 D/N= 16666.7 loss=2.6198
N= 3.0e+08 D= 5.6e+11 D/N= 1851.9 loss=2.4344
N= 1.0e+09 D= 1.7e+11 D/N= 166.7 loss=2.3400
N= 3.0e+09 D= 5.6e+10 D/N= 18.5 loss=2.3363
N= 1.0e+10 D= 1.7e+10 D/N= 1.7 loss=2.4160
N= 3.0e+10 D= 5.6e+09 D/N= 0.2 loss=2.5687
Why it is a U
- Small
N, hugeD: oceans of data, no capacity to absorb it. TheA/N^αterm dominates. - Huge
N, tinyD: enormous capacity, starved of examples. TheB/D^βterm dominates. - In between: balanced. A single interior minimum.
loss
2.62 │●
2.57 │ ●
2.43 │ ●
2.42 │ ●
2.34 │ ● ●
└──────────────────────────────────────────► log10(N)
8.0 8.5 9.0 9.5 10.0 10.5
underfit OPT data-starved
Step 3: fit the parabola
You have noisy points. You want the vertex. Fit a quadratic in log N and solve for its
minimum analytically.
def fit_quadratic(xs, ys):
"""Least-squares fit of y = a*x^2 + b*x + c. Solves the 3x3 normal equations
with Cramer's rule — no libraries, fully deterministic."""
if len(xs) < 3:
raise ValueError("need at least 3 points for a parabola")
n = len(xs)
s = [sum(x ** p for x in xs) for p in range(5)] # s[0..4]
t = [sum(y * x ** p for x, y in zip(xs, ys)) for p in range(3)]
M = [[s[4], s[3], s[2]],
[s[3], s[2], s[1]],
[s[2], s[1], s[0]]]
rhs = [t[2], t[1], t[0]]
def det3(m):
return (m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
- m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
+ m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]))
D0 = det3(M)
if abs(D0) < 1e-30:
raise ValueError("degenerate design — points are collinear in x")
def replace_col(m, col, v):
return [[v[r] if c == col else m[r][c] for c in range(3)] for r in range(3)]
a = det3(replace_col(M, 0, rhs)) / D0
b = det3(replace_col(M, 1, rhs)) / D0
c = det3(replace_col(M, 2, rhs)) / D0
return a, b, c
def parabola_vertex(a, b):
"""x at the minimum of a*x^2 + b*x + c. Requires a > 0 (a real minimum)."""
if a <= 0:
raise ValueError("parabola opens downward or is flat — no interior minimum")
return -b / (2 * a)
pts = isoflop_sweep(1e21, [1e8, 3e8, 1e9, 3e9, 1e10, 3e10])
log_n = [math.log10(N) for N, _, _ in pts]
losses = [L for _, _, L in pts]
a, b, c = fit_quadratic(log_n, losses)
n_opt = 10 ** parabola_vertex(a, b)
print(f"fitted optimal N = {n_opt:.3e} (D = {1e21/(6*n_opt):.3e}, "
f"ratio {1e21/(6*n_opt)/n_opt:.1f} tokens/param)")
fitted optimal N = 1.951e+09 (D = 8.542e+10, ratio 43.8 tokens/param)
Notice the flat bottom. In the sweep above,
N=1e9andN=3e9differed by 0.004 nats — a factor of 3 in model size for essentially no quality difference. This is a gift: you can move off the exact optimum toward a smaller, cheaper-to-serve model almost for free (Phase 02's whole thesis). It is also a trap: with realistic measurement noise, the fitted vertex wanders. Chapter 10 quantifies how much.
Steps 4–6: the power laws
def run_isoflops_ladder(budgets, sizes_per_budget):
"""Steps 1-4: get (C, N_opt, D_opt) for several budgets."""
results = []
for C, sizes in zip(budgets, sizes_per_budget):
pts = isoflop_sweep(C, sizes)
lx = [math.log10(N) for N, _, _ in pts]
ly = [L for _, _, L in pts]
a, b, _ = fit_quadratic(lx, ly)
N_opt = 10 ** parabola_vertex(a, b)
results.append((C, N_opt, C / (6 * N_opt)))
return results
budgets = [1e19, 1e20, 1e21, 1e22]
sizes = [[N * m for N in (1e7, 3e7, 1e8, 3e8, 1e9)] for m in (1, 3, 10, 30)]
ladder = run_isoflops_ladder(budgets, sizes)
# Steps 5-6: fit the power laws.
Cs = [C for C, _, _ in ladder]
Ns = [N for _, N, _ in ladder]
Ds = [D for _, _, D in ladder]
_, a_exp = fit_power_law(Cs, Ns)
_, b_exp = fit_power_law(Cs, Ds)
print(f"N_opt ∝ C^{a_exp:.3f}")
print(f"D_opt ∝ C^{b_exp:.3f}")
print(f"consistency: a + b = {a_exp + b_exp:.3f} (must be 1.000)")
N_opt ∝ C^0.454
D_opt ∝ C^0.546
consistency: a + b = 1.000
a + b = 1 exactly, and it must — it falls straight out of C = 6ND. Run this check every
single time. If your fitted exponents do not sum to ~1, you have a bug in your FLOP accounting
or your fit, and you should find it before you spend $30M.
And the exponent, ~0.45, is close to Chinchilla's ~0.5: scale N and D at roughly the same
rate. (It is not exactly 0.5 because the ground-truth α and β used here are not equal —
optimum_exponent(α, β) = β/(α+β), which is 0.5 precisely when α = β.)
Chapter 7: The Parametric Form L(N, D)
IsoFLOPs finds the optimum without ever writing down a formula. The alternative — and what you need if you want to answer "what if I have 3× the data but the same compute?" — is to fit the full surface.
$$ L(N, D) = E + \frac{A}{N^{\alpha}} + \frac{B}{D^{\beta}} $$
Reading the five parameters
CHINCHILLA = dict(E=1.69, A=406.4, alpha=0.34, B=410.7, beta=0.28)
def parametric_loss(N, D, E, A, alpha, B, beta):
if N <= 0 or D <= 0:
raise ValueError("N and D must be positive")
return E + A / (N ** alpha) + B / (D ** beta)
# Decompose the loss at a realistic operating point.
N, D = 70e9, 1.4e12
p = CHINCHILLA
cap = p["A"] / N ** p["alpha"]
dat = p["B"] / D ** p["beta"]
print(f"irreducible E : {p['E']:.4f} ({p['E']/(p['E']+cap+dat):5.1%})")
print(f"capacity A/N^a : {cap:.4f} ({cap/(p['E']+cap+dat):5.1%})")
print(f"data B/D^b : {dat:.4f} ({dat/(p['E']+cap+dat):5.1%})")
print(f"total : {parametric_loss(N, D, **p):.4f}")
irreducible E : 1.6900 (87.3%)
capacity A/N^a : 0.0835 ( 4.3%)
data B/D^b : 0.1632 ( 8.4%)
total : 1.9366
Look at that: 87% of the loss is irreducible. All the money in the industry is being spent on the remaining 13%. That is worth internalizing — it explains why improvements look so small in nats, and why compressing a 0.01-nat gap is worth a nine-figure budget.
It also tells you which lever to pull: here the data term (0.163) is nearly twice the capacity term (0.084), so at this operating point more data helps more than more parameters — which is precisely the Chinchilla recommendation, visible directly in the decomposition.
Deriving the optimum analytically
With the parametric form you can solve for the optimum instead of sweeping. Minimize
L(N, C/(6N)) over N:
$$ \frac{\partial}{\partial N}\left[\frac{A}{N^\alpha} + \frac{B}{(C/6N)^\beta}\right] = 0 $$
$$ \Rightarrow \quad -\alpha A N^{-\alpha-1} + \beta B (6/C)^{\beta} N^{\beta - 1} = 0 $$
$$ \Rightarrow \quad N_{\text{opt}} = \left[\frac{\alpha A}{\beta B}\left(\frac{C}{6}\right)^{\beta}\right]^{\frac{1}{\alpha+\beta}} $$
which is a power law in C with exponent β/(α+β).
def analytic_optimum(C, E, A, alpha, B, beta):
"""Closed-form compute-optimal (N, D) under the parametric law."""
N = ((alpha * A) / (beta * B) * (C / 6.0) ** beta) ** (1.0 / (alpha + beta))
return N, C / (6 * N)
for C in (1e19, 1e21, 1e23, 1e25):
N, D = analytic_optimum(C, **CHINCHILLA)
print(f"C={C:.0e} N={N:9.3e} D={D:9.3e} ratio={D/N:7.1f} tok/param")
exponent = CHINCHILLA["beta"] / (CHINCHILLA["alpha"] + CHINCHILLA["beta"])
print(f"\nN_opt ∝ C^{exponent:.4f}")
C=1e+19 N=2.280e+08 D=7.311e+09 ratio= 32.1 tok/param
C=1e+21 N=1.824e+09 D=9.136e+10 ratio= 50.1 tok/param
C=1e+23 N=1.460e+10 D=1.142e+12 ratio= 78.2 tok/param
C=1e+25 N=1.168e+11 D=1.427e+13 ratio= 122.1 tok/param
N_opt ∝ C^0.4516
A caution worth stating. Notice that the token-per-parameter ratio is not constant — it climbs from 32 at
1e19to 122 at1e25. The familiar "20 tokens per parameter" is only a rule of thumb at one scale; the parametric law says the ratio drifts wheneverα ≠ β. Note also that these particularE/A/Bvalues have been debated in follow-up work (see Besiroglu et al. in the References), and the ratio is extremely sensitive toαandβ. This is not a flaw in the method — it is the lesson. Small changes in fitted exponents produce large changes in the recommendation, which is exactly why Chapter 9 (the estimator problem) and Chapter 10 (confidence intervals) exist. When you fit your own law, always report the ratio and its uncertainty.
Chapter 8: Fitting It — Least Squares, Logs, and Huber
You have (N, D, L) points and want (E, A, α, B, β). This is a nonlinear fit, and how you
measure "fit" changes the answer.
The objective choices
def residuals(points, params, space="log"):
"""points: [(N, D, L_observed), ...]; params: (E, A, alpha, B, beta)."""
E, A, alpha, B, beta = params
out = []
for N, D, L_obs in points:
L_pred = E + A / N ** alpha + B / D ** beta
if space == "linear":
out.append(L_pred - L_obs)
elif space == "log":
if L_pred <= 0 or L_obs <= 0:
raise ValueError("log-space residuals need positive losses")
out.append(math.log(L_pred) - math.log(L_obs))
else:
raise ValueError(f"unknown space {space!r}")
return out
def squared_loss(rs):
return sum(r * r for r in rs)
def huber_loss(rs, delta=1e-3):
"""Quadratic near zero, LINEAR in the tails — so one bad run cannot dominate.
Chinchilla used Huber on log-space residuals. This is not incidental: scaling
ladders contain genuine outliers (a run that diverged, a bad data shard), and
least squares would let one of them set your flagship recommendation.
"""
total = 0.0
for r in rs:
if abs(r) <= delta:
total += 0.5 * r * r
else:
total += delta * (abs(r) - 0.5 * delta)
return total
Why the space matters
Losses in a ladder span a wide range — a 10M-parameter model might sit at 4.0 nats, a 10B model at 1.9.
- Least squares on raw
Lweights absolute error. A 0.1-nat miss on the 4.0 point counts the same as a 0.1-nat miss on the 1.9 point. Since the big-loss points are the small models, the fit gets dragged toward matching small models well — exactly the points you care least about, since you are extrapolating upward. - Least squares on
log Lweights relative error, treating all scales evenly. - Huber on
log Ldoes that and refuses to let one outlier dominate.
def compare_objectives(points, params):
lin = residuals(points, params, "linear")
log = residuals(points, params, "log")
print(f" sum-sq on L : {squared_loss(lin):.6f}")
print(f" sum-sq on log L : {squared_loss(log):.6f}")
print(f" Huber on log L : {huber_loss(log):.6f}")
pts = [(N, D, true_loss(N, D)) for N, D in
[(1e7, 1e10), (1e8, 1e11), (1e9, 1e11), (1e10, 1e12)]]
# Inject one outlier: a run that diverged and reported a bad loss.
pts_outlier = pts[:-1] + [(1e10, 1e12, pts[-1][2] + 1.5)]
print("clean ladder, true params:")
compare_objectives(pts, (1.69, 406.4, 0.34, 410.7, 0.28))
print("ladder with ONE diverged run, same params:")
compare_objectives(pts_outlier, (1.69, 406.4, 0.34, 410.7, 0.28))
clean ladder, true params:
sum-sq on L : 0.000000
sum-sq on log L : 0.000000
Huber on log L : 0.000000
ladder with ONE diverged run, same params:
sum-sq on L : 2.250000
sum-sq on log L : 0.305852
Huber on log L : 0.000553
One bad run contributes 2.25 to the least-squares objective and 0.00055 to Huber — a factor of about 4,000. Under least squares the fitter will distort all five parameters to chase that outlier. Under Huber it barely notices. On a real ladder, where a diverged run is a routine occurrence, this is the difference between a usable law and a garbage one.
Fitting without a library
You need a nonlinear optimizer. Coordinate descent with a shrinking step is enough, fully deterministic, and about twenty lines:
def fit_parametric(points, init=(1.5, 400.0, 0.35, 400.0, 0.30),
objective="huber", space="log",
iters=200, seed_step=0.5, shrink=0.97):
"""Deterministic coordinate descent. No randomness, no library optimizer,
so the same input always gives the same fit — which is a testable property."""
params = list(init)
step = [seed_step * abs(p) if p != 0 else seed_step for p in params]
def score(p):
rs = residuals(points, tuple(p), space)
return huber_loss(rs) if objective == "huber" else squared_loss(rs)
best = score(params)
for _ in range(iters):
improved = False
for i in range(len(params)):
for direction in (+1, -1):
trial = list(params)
trial[i] = params[i] + direction * step[i]
# Keep exponents and scales in a sane region.
if trial[i] <= 0:
continue
s = score(trial)
if s < best:
best, params, improved = s, trial, True
break
if not improved:
step = [s * shrink for s in step]
return tuple(params), best
Why hand-rolled? Because a library optimizer's version, tolerance, and default method all change the fitted exponents, and therefore your flagship recommendation. If a
scipyupgrade can move yourN_optby 20%, your process is not reproducible. Frontier teams pin this ruthlessly.
Chapter 9: The Estimator Problem — Feinberg's Open Question
His closing slide lists this as a research direction, and it is worth quoting in full because it is a rare, concrete, fundable, GPU-free open problem stated by someone who would know:
"Scaling laws are brittle, dataset dependent.
L(N, D, etc.)– of course we can add more dims to improve fit. Least squares vs MLE & formal stats model imply different scaling recommendations! Formalize. Rather than grid(N, D)where do we get max info gain? Active learn…"
The two problems
Problem A — the estimator is unspecified. Chinchilla used Huber on log-residuals. That is a choice, not a derivation. A proper statistical treatment would:
- Write down a generative model: what is the noise on a measured loss? Is it additive? Multiplicative? Heteroscedastic (bigger for small models, which are noisier)? Correlated across points from the same run?
- Derive the maximum-likelihood estimator under that model.
- Get calibrated uncertainty for free, instead of bootstrapping.
Nobody has fully done this, and the choice demonstrably changes the recommendation:
def recommendation_under_estimator(points, objective, space):
params, _ = fit_parametric(points, objective=objective, space=space)
N_opt, D_opt = analytic_optimum(1e24, *params)
return params, N_opt, D_opt
ladder = [(N, D, true_loss(N, D)) for N, D in
[(1e7, 2e10), (3e7, 6e10), (1e8, 2e11), (3e8, 6e11), (1e9, 2e12)]]
for obj, sp in (("squared", "linear"), ("squared", "log"), ("huber", "log")):
params, N_opt, D_opt = recommendation_under_estimator(ladder, obj, sp)
print(f"{obj:8s}/{sp:6s} -> flagship N = {N_opt:.3e}, D = {D_opt:.3e}")
Run it and watch the recommendations diverge on identical data. That divergence is the open problem.
Problem B — the design is a grid. Everybody runs N ∈ {…} × C ∈ {…}. But a grid is not an
efficient experiment. The right question is: given my current posterior, which next run most
reduces the variance of my prediction at C = 1e25?
def extrapolation_variance(log_c_points, target_log_c):
"""Variance of a linear extrapolation to `target_log_c`, up to a noise constant.
Standard OLS result: Var(y_hat) ∝ 1/n + (x* - x̄)^2 / Σ(x - x̄)^2
"""
n = len(log_c_points)
if n < 2:
raise ValueError("need at least 2 design points")
xbar = sum(log_c_points) / n
sxx = sum((x - xbar) ** 2 for x in log_c_points)
if sxx == 0:
raise ValueError("all design points identical")
return 1.0 / n + (target_log_c - xbar) ** 2 / sxx
TARGET = 25.0 # forecasting a 1e25 FLOP run
designs = {
"clustered (4 runs, all ~1e19)": [19.0, 19.2, 19.4, 19.6],
"spread (4 runs, 1e18-1e21)": [18.0, 19.0, 20.0, 21.0],
"budget-skewed (3 small, 1 big)": [18.0, 18.5, 19.0, 21.5],
"two-point extremes": [18.0, 21.0],
}
for name, d in designs.items():
print(f"{name:32s} Var ∝ {extrapolation_variance(d, TARGET):8.2f}")
clustered (4 runs, all ~1e19) Var ∝ 162.70
spread (4 runs, 1e18-1e21) Var ∝ 6.30
budget-skewed (3 small, 1 big) Var ∝ 4.81
two-point extremes Var ∝ 7.22
The clustered design is 26× worse than the spread one, using the same number of runs. And note the last row: two-point extremes, with half the runs, beats four clustered runs by 23×. Note also that the budget-skewed design — three cheap points plus one expensive one far out — edges out the evenly spread one, because that distant anchor is doing most of the work.
The practical rule this gives you: when planning a ladder, maximize the spread of
log C, subject to your budget and to the smallest model still being large enough to be in the scaling regime. Do not add a fifth point near your existing four; add one an order of magnitude out.
Chapter 10: Confidence Intervals via Bootstrap
A forecast without error bars is not a forecast. And you cannot derive analytic intervals here because the model is nonlinear and the noise model is unknown (Chapter 9). So: bootstrap.
What the bootstrap is
You have n ladder points. Resample n of them with replacement, refit, record the
prediction. Do it 1,000 times. The spread of those 1,000 predictions estimates the spread of
your prediction.
The logic: your sample is your best available picture of the population, so resampling from it mimics drawing fresh samples from the population.
import random
def bootstrap_forecast(points, target_C, n_boot=200, seed=0):
"""Bootstrap CI for the loss forecast at target_C. Seeded -> reproducible."""
rng = random.Random(seed) # SEEDED. Same seed -> same bytes.
forecasts = []
for _ in range(n_boot):
sample = [points[rng.randrange(len(points))] for _ in range(len(points))]
try:
params, _ = fit_parametric(sample, iters=60)
N, D = analytic_optimum(target_C, *params)
forecasts.append(parametric_loss(N, D, *params))
except (ValueError, OverflowError, ZeroDivisionError):
continue # a degenerate resample; skip it
if len(forecasts) < 10:
raise ValueError("too few successful bootstrap fits — ladder is too small")
forecasts.sort()
lo = forecasts[int(0.025 * len(forecasts))]
hi = forecasts[int(0.975 * len(forecasts))]
mid = forecasts[len(forecasts) // 2]
return {"median": mid, "ci_low": lo, "ci_high": hi, "width": hi - lo,
"n_successful": len(forecasts)}
How to read the width
Remember Chapter 3: 0.01 nats is roughly a third more compute. So:
| CI width | Verdict |
|---|---|
| < 0.01 nats | Tight. You can make a confident recipe call. |
| 0.01–0.05 | Usable for go/no-go, not for choosing between similar recipes. |
| > 0.05 | Your ladder cannot support this decision. Run more/wider points before the meeting. |
The senior move. When someone presents a scaling forecast without an interval, ask for one. When they present an interval wider than the effect they are claiming, the honest conclusion is "we do not know yet" — and saying so is far more valuable than a confident wrong number that gets baked into a $40M plan.
Chapter 11: Designing the Ladder — Where to Spend Your Ablations
Pulling Chapters 9 and 10 into a procedure.
The constraints
- Budget. The ladder should cost a small fraction of the flagship — a few percent is
typical, and Phase 00's
$ per 1e21 FLOPsmakes that concrete. - The scaling regime. Below ~10M non-embedding parameters, models behave differently (embeddings dominate, optimization is qualitatively different). Points below the regime are worse than useless — they actively bias the fit.
- Spread. Chapter 9: variance is driven by the spread of
log C. - Replication. At least one budget should be run twice to measure your noise level. Without that, you are guessing at the thing your confidence interval depends on.
A workable recipe
def design_ladder(flagship_C, ladder_budget_fraction=0.03, n_budgets=5,
min_C=1e18):
"""Geometrically spaced budgets consuming a fixed fraction of the flagship."""
total = flagship_C * ladder_budget_fraction
# Geometric spacing means each budget is the same multiplicative step apart,
# which maximizes spread in log space for a given range.
top = total / 2.0 # the largest ladder run gets half the ladder budget
ratio = (top / min_C) ** (1.0 / (n_budgets - 1))
budgets = [min_C * ratio ** i for i in range(n_budgets)]
return budgets
flagship = 1e25
budgets = design_ladder(flagship)
print(f"flagship: {flagship:.0e} FLOPs")
print(f"ladder total: {sum(budgets):.3e} FLOPs "
f"({sum(budgets)/flagship:.2%} of flagship)")
for C in budgets:
print(f" C={C:.3e} ({C/flagship:.2e} of flagship)")
print(f"log-C spread: {math.log10(budgets[-1]) - math.log10(budgets[0]):.1f} decades")
flagship: 1e+25 FLOPs
ladder total: 1.580e+23 FLOPs (1.58% of flagship)
C=1.000e+18 (1.00e-07 of flagship)
C=1.968e+19 (1.97e-06 of flagship)
C=3.873e+20 (3.87e-05 of flagship)
C=7.622e+21 (7.62e-04 of flagship)
C=1.500e+23 (1.50e-02 of flagship)
log-C spread: 5.2 decades
Under 2% of the flagship budget buys you 5.2 decades of spread. In money (Phase 00's $1,753 per 1e21 FLOPs): the flagship is ~$17.5M and the entire ladder is ~$277k. That is the single best-value spend in the whole project, and it is the argument you make when someone suggests skipping ablations to save time.
Chapter 12: Making the Decision — Baseline vs Candidate
Feinberg's slide gives the procedure in three lines:
"How changes get adopted in classical setting:
- Derive
L*(flops)baselineL*(flops)candidate"and, on the refinements slide: "To make a change, compare baseline vs candidate laws."
Why you compare laws, not runs
The naive approach — "train both at 1B parameters, pick the winner" — fails because the curves can cross. An architecture change that helps small models can hurt large ones, and vice versa. Comparing at one scale tells you about that scale only.
def compare_recipes(baseline_params, candidate_params, target_C):
"""Compare two fitted laws AT THE TARGET, and find the crossover."""
def loss_at(params, C):
N, D = analytic_optimum(C, *params)
return parametric_loss(N, D, *params)
base = loss_at(baseline_params, target_C)
cand = loss_at(candidate_params, target_C)
# Scan for a crossover on a log grid.
crossover = None
prev = None
for i in range(100):
C = 10 ** (16 + i * 0.12)
sign = loss_at(candidate_params, C) < loss_at(baseline_params, C)
if prev is not None and sign != prev:
crossover = C
break
prev = sign
return {"baseline_loss": base, "candidate_loss": cand,
"delta": cand - base, "candidate_wins": cand < base,
"crossover_C": crossover}
baseline = (1.69, 406.4, 0.34, 410.7, 0.28)
# A candidate that is better per-parameter (bigger A-exponent) but worse per-token.
candidate = (1.69, 500.0, 0.38, 380.0, 0.26)
for C in (1e20, 1e22, 1e24, 1e26):
r = compare_recipes(baseline, candidate, C)
verdict = "CANDIDATE" if r["candidate_wins"] else "baseline "
print(f"C={C:.0e}: base {r['baseline_loss']:.4f} "
f"cand {r['candidate_loss']:.4f} delta {r['delta']:+.4f} -> {verdict}")
r = compare_recipes(baseline, candidate, 1e24)
print(f"\ncrossover at C ≈ {r['crossover_C']:.2e}" if r["crossover_C"]
else "\nno crossover in range")
The crossover is the deliverable, not the winner. "Candidate wins above 1e23 FLOPs" is an actionable statement; "candidate is better" is not, because it silently assumes a scale.
The decision checklist
Before you recommend a recipe change, you should be able to answer all six:
- Have both laws been fitted on the same ladder design and the same estimator?
- Is the delta at the target larger than the confidence interval on either forecast?
- Where is the crossover, and is the flagship comfortably on one side of it?
- Does the candidate change the serving cost (Phase 02, Phase 05)? A 0.01-nat win that doubles inference cost is a loss.
- Does the candidate change the failure modes (Phase 09)? A recipe that is 0.005 nats better and spikes twice as often will cost you more in goodput than it gains.
- Do the downstream evals agree with the loss delta, or is the loss improvement coming from something users do not care about?
Chapter 13: When Scaling Laws Lie
Feinberg's slide "The End of Scaling?" is measured about this, and worth taking seriously in both directions.
The failure modes
1. Out-of-regime extrapolation. Every law is fitted over a range. Three decades beyond it, the functional form itself may be wrong — a power law and a saturating curve can be indistinguishable on your ladder and disagree by a lot at the target.
2. Recipe drift. The law describes the recipe you fitted. If between the ladder and the flagship you change the data mixture, the tokenizer, the optimizer, or how depth scales — the law no longer applies. This is a shockingly common real-world failure: the ladder was run in March, the flagship in July, and the data team shipped three improvements in between.
3. Data exhaustion. The law says "more tokens." If you do not have them, the law is
answering a question you cannot act on. This is Phase 02's L(N, U, R).
4. Benchmark saturation, not model saturation. His slide: "LMSys is not the end-all-be-all. Llama 4 Maverick demonstrated that ranking can be volatile and overfit to human preference." A flat benchmark can mean a flat model or a dead benchmark. They look identical.
His counterargument to the doomers
He does not conclude scaling is over. The slide lists two reasons to expect continued progress:
"1. Better NN design still coming 2. Data from new sources being added"
and the whole "More Data Sources" slide is about multimodal and synthetic data (Phase 02). His framing of the mission is worth keeping: "job is to push the curves right."
The practical posture
def extrapolation_risk(ladder_max_C, target_C, ci_width):
"""A crude, honest risk score for a forecast."""
decades = math.log10(target_C / ladder_max_C)
if decades <= 0:
return "INTERPOLATION — low risk"
risk = decades * (1 + 20 * ci_width)
band = ("LOW" if risk < 2 else "MODERATE" if risk < 4 else "HIGH")
return (f"{decades:.1f} decades of extrapolation, CI {ci_width:.3f} nats "
f"-> {band} risk (score {risk:.1f})")
print(extrapolation_risk(1e22, 1e25, 0.008))
print(extrapolation_risk(1e20, 1e25, 0.040))
3.0 decades of extrapolation, CI 0.008 nats -> MODERATE risk (score 3.5)
5.0 decades of extrapolation, CI 0.040 nats -> HIGH risk (score 9.0)
Present it this way. "We forecast 1.712 nats, 95% CI [1.704, 1.721], extrapolating 3 decades beyond our largest ablation, assuming the data mixture is frozen from today." That sentence is what a principal engineer says. Everything before the comma is arithmetic; everything after it is judgment, and the judgment is what you are paid for.
Lab Walkthrough
Lab 01 — IsoFLOPs Ladder, Law Fitting & the Forecast
Implement in this order:
fit_line,fit_power_law,fit_quadratic,parabola_vertex— the numerical spine.fit_quadraticuses Cramer's rule on the 3×3 normal equations; it must raise on a degenerate (collinear) design rather than divide by zero.synthetic_ladder— generate(N, D, L)with seeded noise and an injectablelr_decay_bias. This is what lets you reproduce the Kaplan bug.isoflop_optimum— sweep, fit the parabola, return the vertex.fit_scaling_exponents— steps 5–6, plus thea + b ≈ 1consistency check.residuals/huber_loss/fit_parametric— the surface fit. Deterministic coordinate descent; same input must give byte-identical output.analytic_optimum— the closed form. A test checks it agrees with the numerical sweep.bootstrap_forecast— seeded resampling, percentile interval.extrapolation_variance/design_ladder— the optimal-design tooling.compare_recipes— the decision rule with the crossover scan.
The money test is test_lr_decay_bias_shifts_recommendation_toward_smaller_models: generate
a ladder with the bias, fit, and confirm the recommended N is meaningfully larger than the
unbiased fit's. That is Kaplan → Chinchilla, reproduced from scratch.
The traps:
fit_quadraticon 2 points, or on collinear points → must raise, notZeroDivisionError.- A parabola that opens downward has no minimum → must raise.
- Bootstrap resamples can be degenerate (all identical points) → catch and skip, and fail loudly if too few succeed.
- Everything random must go through a seeded
random.Random(seed).
Success Criteria
-
LAB_MODULE=solution pytest test_lab.py -vall green. -
Your
lab.pypasses after the TODOs. -
python solution.pyreproduces the Kaplan → Chinchilla story from synthetic data. - You can explain, without notes, why reading loss mid-run is a biased estimator and why the bias has a direction.
-
You can state why
a + b ≈ 1must hold and what it catches. - You can produce a forecast with a CI and say whether it is tight enough to decide on.
- You can justify a ladder design in terms of extrapolation variance and cost.
Interview Q&A
Q: What is a scaling law and why does the field need them?
An empirical formula, typically L = E + A/N^α + B/D^β, predicting test loss from parameters
and tokens. The field needs them because pre-training is a one-shot extrapolation problem: each
flagship run costs eight figures, takes months, and is by construction larger than anything
previously run, so you cannot validate by trying alternatives. You fit the law on a cheap ladder
and extrapolate. Critically, the law is a property of a fixed, parameterized recipe, not of
nature.
Q: Kaplan said N ∝ C^0.73; Chinchilla said C^0.5. What happened?
A methodological bug. Kaplan ran one training run per model size and read intermediate losses
to estimate loss at smaller token horizons. But learning-rate schedules decay to near zero at
the end, and much of the final improvement comes from that decay. A mid-run reading is therefore
systematically worse than a run actually scheduled to stop there — and the bias is largest for
the points furthest from their horizon, which tilts the fit against training on more data.
Chinchilla ran separate, properly-scheduled runs, found the exponent was ~0.5, and concluded
models had been badly undertrained.
Q: Walk me through IsoFLOPs.
Fix a FLOP budget. Train several models at different N, deriving D = C/6N so every run costs
the same. Loss versus log N is U-shaped — too small underfits, too large is data-starved — so
fit a parabola and take the vertex as N_opt for that budget. Repeat across budgets, then fit
power laws N_opt ∝ C^a and D_opt ∝ C^b. Sanity check: a + b must be ~1, since C = 6ND.
Q: How do you fit the parametric form, and why does it matter how? Minimize a robust loss — Huber — on log-space residuals. Log space because losses span a wide range and you care about relative not absolute error, and because absolute error would let the small, high-loss models dominate the fit exactly when you are extrapolating upward. Huber because ladders contain genuine outliers — a diverged run, a bad shard — and under least squares one of those can move all five parameters. It matters because different estimators give different exponents on identical data, and therefore different flagship recommendations. That is an open problem Feinberg explicitly calls out: "least squares vs MLE... Formalize."
Q: How do you put error bars on a forecast? Bootstrap: resample the ladder points with replacement, refit, predict, repeat ~1000 times, take percentiles. You cannot do it analytically because the model is nonlinear and the noise model is unspecified. Then interpret against scale: roughly, 0.01 nats is a third more compute, so a CI wider than 0.05 nats cannot support a recipe decision.
Q: I will give you 3% of the flagship budget for ablations. How do you spend it?
Geometrically spaced budgets to maximize the spread of log C, since extrapolation variance
goes as (x* − x̄)² / Σ(x − x̄)². Keep the smallest run inside the scaling regime — below ~10M
non-embedding parameters the behaviour is different and points there bias the fit. Give the
largest single run about half the ladder budget, because it anchors the extrapolation. And
replicate one budget so I can measure my noise rather than assume it.
Q: Two recipes. How do you decide?
Fit a law for each on the same ladder design with the same estimator, evaluate both at the
target C, and report the delta with confidence intervals plus the crossover point. Then
check three things the loss number cannot see: does the candidate change serving cost, does it
change training stability, and do downstream evals move in the same direction as loss.
Q: Are scaling laws ending? Loss-versus-compute curves have not visibly bent. What people usually observe is benchmark saturation, which is a measurement problem — leaderboards are volatile and can overfit human preference. The two live sources of continued progress are better architectures and new data sources, especially multimodal and synthetic. The real constraint is not the law flattening, it is unique high-quality data running out, which changes the shape of the optimization rather than ending it.
Tips & Takeaways
Tips
- Always check
a + b ≈ 1. Free bug detector on any fitted pair of exponents. - Always fit in log space, always use a robust loss. Two lines of code, enormous robustness gain.
- Always report a confidence interval, and compare its width against 0.01 nats ≈ 33% more compute.
- Pin your fitting code. A library optimizer's default method changing between versions can move your flagship recommendation. Hand-roll or pin exactly.
- Replicate one ladder point. You cannot calibrate uncertainty without measuring noise.
- Write down the recipe alongside the law. A law without its recipe is uninterpretable in six months.
- Spread beats density. One more point an order of magnitude out beats three more points near your existing cluster.
- Say the extrapolation distance out loud when presenting: "three decades beyond our largest ablation."
Takeaways
- Pre-training is a one-shot extrapolation problem. Scaling laws exist because of that.
- A law describes your recipe, not the universe. Fix the recipe first.
- Loss forecasting is recipe selection.
L = E + A/N^α + B/D^β— irreducible + capacity + data. At frontier scale ~89% of the loss is irreducible; the whole industry is fighting over the rest.- The IsoFLOPs U-curve has a flat bottom: a gift for serving-driven deviation, a trap for noisy fits.
- Chinchilla beat Kaplan by fixing an experimental-design bug. Methodology is the frontier.
- The estimator you choose changes your flagship recommendation. This is an open problem.
- Where you place ladder points matters more than how many. 13× variance reduction, same cost.
- A forecast without error bars is not a forecast.
- Compare laws at the target and report the crossover, never a single-scale winner.
References
Primary
- Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides — the "Classical Scaling" section: the IsoFLOPs six-step build, the Kaplan/Chinchilla contrast, "UNDERTRAINED!", "Least squares vs MLE… Formalize", "Active learn…"
- Austin et al., How To Scale Your Model — https://jax-ml.github.io/scaling-book/ — do the exercises; he offers interviews for them
Core papers
- Kaplan et al., Scaling Laws for Neural Language Models, 2020 — https://arxiv.org/abs/2001.08361
- Hoffmann et al., Training Compute-Optimal Large Language Models (Chinchilla), 2022 — https://arxiv.org/abs/2203.15556
- Besiroglu et al., Chinchilla Scaling: A Replication Attempt, 2024 — https://arxiv.org/abs/2404.10102 — a replication finding issues in the original's reported fit; read this alongside Chinchilla, it is Chapter 9's problem in the wild
- Clark et al., Unified Scaling Laws for Routed Language Models, 2022 — https://arxiv.org/abs/2202.01169
- Muennighoff et al., Scaling Data-Constrained Language Models, 2023 — https://arxiv.org/abs/2305.16264 (Phase 02)
- Sardana et al., Beyond Chinchilla-Optimal, 2024 — https://arxiv.org/abs/2401.00448 (Phase 02)
- Grattafiori et al., The Llama 3 Herd of Models, 2024 — https://arxiv.org/abs/2407.21783 — their scaling-law methodology section is an unusually clear published account
Statistics
- Efron & Tibshirani, An Introduction to the Bootstrap — the method in Chapter 10
- Huber, Robust Estimation of a Location Parameter, 1964 — the loss in Chapter 8
- Chaloner & Verdinelli, Bayesian Experimental Design, 1995 — the formal version of Chapter 11
Lab 01 — IsoFLOPs Ladder, Law Fitting & the Forecast
Build the machine that turns a few cheap experiments into a defensible prediction about an experiment you can only run once.
The problem
You have a flagship budget of 1e25 FLOPs and one shot. Before spending it you must state
what you will train and what loss you expect, and defend both. The only tool is a ladder of
small runs and a fitted curve — so the whole job reduces to: fit well, extrapolate honestly,
and know how wrong you might be.
This lab builds that machine end to end, and along the way reproduces the Kaplan → Chinchilla correction from scratch — because the lesson there is not a fact to memorize, it is a failure mode to recognize in your own work.
What you build
| Group | Functions | The idea |
|---|---|---|
| Numerical spine | fit_line, fit_power_law, fit_quadratic, parabola_vertex | A power law is a straight line in log-log; the IsoFLOPs minimum is a parabola vertex |
| Loss model & ladder | parametric_loss, loss_decomposition, schedule_bias, synthetic_ladder | L = E + A/N^α + B/D^β, plus an injectable measurement bias |
| IsoFLOPs | isoflop_optimum, isoflops_ladder, fit_scaling_exponents | The six-step method, ending in the a + b ≈ 1 check |
| Surface fit | residuals, squared_loss, huber_loss, fit_parametric | Log-space, robust, deterministic multi-start |
| Optimum | analytic_optimum, optimum_exponent | The closed form, and why Chinchilla's exponent is ~0.5 |
| Uncertainty | bootstrap_forecast, compute_multiplier_for_loss_delta | Error bars, and how to read them in dollars |
| Design | extrapolation_variance, design_ladder, extrapolation_risk | Where to spend your next ablation |
| Decision | loss_at_budget, compare_recipes, decision_is_supported | Compare laws at the target; report the crossover |
Key concepts
| Concept | Why it is in this lab |
|---|---|
L = E + A/N^α + B/D^β | Irreducible + capacity + data. At frontier scale ~87% is irreducible. |
| The IsoFLOPs U | At fixed C, too-small underfits and too-large starves. Flat bottom. |
a + b ≈ 1 | Falls out of C = 6ND. A free bug detector on any fitted pair. |
| Non-uniform bias | A constant offset is absorbed into E; only a varying one tilts exponents. |
| Huber on log residuals | One diverged run costs least-squares ~4,000× what it costs Huber. |
| Bootstrap CI | The model is nonlinear and the noise model unspecified — so resample. |
| 0.01 nats ≈ 33% compute | The conversion that makes a confidence interval interpretable. |
| Extrapolation variance | Spread beats density: 26× lower variance from placement alone. |
| The crossover | "Candidate wins above 3e24" is actionable; "candidate is better" is not. |
Files
| File | What it is |
|---|---|
| lab.py | Your implementation. Signatures, docstrings and validation contracts given. |
| solution.py | Reference. python solution.py runs the full eleven-part story. |
| test_lab.py | 64 tests: happy path, validation, boundaries, invariants, determinism. |
| requirements.txt | pytest only. Pure stdlib otherwise. |
Run
pytest test_lab.py -v # your lab.py — red until you implement
LAB_MODULE=solution pytest test_lab.py -v # the reference — must be green (64 passed)
python solution.py # the worked example
Where to start
Work top to bottom; each group depends on the one above.
- The spine.
fit_line→fit_power_law→fit_quadratic→parabola_vertex. Gettest_fit_power_law_recovers_exponentandtest_fit_quadratic_recovers_exact_parabolapassing before anything else. parametric_lossandsynthetic_ladder. Note the budget constraint: every point on an IsoFLOPs curve satisfiesC = 6ND, and a test checks it.- The IsoFLOPs method, ending at
test_exponents_sum_to_one. fit_parametric. Deterministic multi-start coordinate descent.test_fit_parametric_is_deterministicis not optional garnish — a fit that moves between runs cannot support a nine-figure decision.- The rest.
The money test is test_nonuniform_bias_tilts_the_fitted_exponent, paired with its control
test_uniform_offset_is_absorbed_into_E_and_changes_nothing. Together they are the
Kaplan → Chinchilla lesson: a uniform measurement error changes nothing, a non-uniform one
changes your flagship recommendation.
The traps:
fit_quadraticon collinearxvalues → must raiseValueError, notZeroDivisionError.- A downward-opening parabola has a maximum → must raise, because it means your sweep is wrong.
- Bootstrap resamples can be degenerate → skip them, and fail loudly if too few succeed.
- Do not recompute
C = 6NDto regroup a ladder. Floating-point rounding will split one budget into two groups. Carry the budget label. (A test comment calls this out.) - Everything random goes through a seeded
random.Random(seed).
Success criteria
-
LAB_MODULE=solution pytest test_lab.py -v→ 64 passed. -
Your
lab.pyreaches 64 passed. -
python solution.pyruns, and you can explain all eleven sections. - You can explain why a uniform loss offset changes nothing and a non-uniform one changes everything.
-
You can state why
a + b ≈ 1must hold and what it catches. - You can produce a forecast with a CI and say whether it is tight enough to decide on.
- You can justify a ladder design in terms of extrapolation variance and dollars.
How this maps to the real stack
| This lab | The real thing | Where the miniature lies |
|---|---|---|
synthetic_ladder | Actual training runs on a cluster, days each | Real losses have autocorrelated noise, run-to-run seed variance, and occasional divergence. Ours is i.i.d. Gaussian. Real ladders also vary the recipe imperfectly across scales. |
fit_parametric | scipy.optimize.minimize with L-BFGS-B, or JAX + optax | Real fits use gradients and converge faster. Ours is hand-rolled so it is deterministic and inspectable — and so a library upgrade cannot silently move your recommendation. |
isoflop_optimum | The same parabola fit, on real measured losses | Identical in spirit. Real practice uses more sizes per budget (6–10) and often fits in log-loss space. |
bootstrap_forecast | Bootstrap or Bayesian posterior over law parameters | Real work often uses a hierarchical model that shares information across budgets. Ours treats points as exchangeable, which understates correlation within a budget. |
compare_recipes | An internal go/no-go memo and a review meeting | The arithmetic is the easy part. The real version also carries downstream eval deltas, serving-cost deltas, and stability risk — see the six-point checklist in WARMUP Chapter 12. |
CHINCHILLA constants | Hoffmann et al. Table 3 | These published values have been contested; see Besiroglu et al. Use them as a plausible ground truth for the simulator, not as gospel. |
What is not a lie: the IsoFLOPs procedure, the a + b = 1 identity, the analytic optimum
derivation, and the extrapolation-variance formula. Those are exact, and they are what you will
be asked about.
Extensions
- Fit against published data. Pull the
(N, D, loss)table from any open model report (Llama 3's scaling section is unusually complete) and fit your law to it. Compare your extrapolation to what they actually shipped. - Add a hierarchical noise model. Points within one budget share a run's seed and infrastructure; model that correlation and see how the confidence interval widens. This is Feinberg's "formal stats model" suggestion, at a tractable scale.
- Implement true optimal design. Instead of scoring fixed candidate designs, search for
the next
Cthat minimizes posterior predictive variance at the flagship. That is the "active learn" bullet on his slide, and it is a publishable exercise. - Add a third axis. Extend to
L(N, D, U)with unique tokens, then compare against Phase 02's data-constrained law. - Compare estimators properly. Run a simulation study: generate 1,000 ladders from known truth, fit with each estimator, and report bias and variance of the recovered exponents. That is a real, small, honest paper.
Interview / resume bullets
- "Implemented the full IsoFLOPs scaling-law methodology from scratch — parabola fitting for
per-budget optima, power-law regression for
N_opt(C)andD_opt(C)with ana + b = 1consistency check, and robust (Huber-on-log) fitting of the parametricL(N, D)surface — producing flagship forecasts with bootstrap confidence intervals." - "Reproduced the Kaplan→Chinchilla correction from first principles by simulating a schedule-mismatch measurement bias, demonstrating that a uniform loss offset is absorbed into the irreducible term while a non-uniform one tilts the fitted exponents and shifts the compute-optimal recommendation."
- "Built optimal-experimental-design tooling for scaling ladders, showing ~26× lower extrapolation variance from point placement at identical cost, and sized a ladder at under 2% of flagship budget for 5.2 decades of coverage."
- Interview-ready: "I would not train both and compare. I would fit a law for each recipe over a shared ladder, evaluate both at the target FLOP count, and report the delta with confidence intervals plus the crossover — because curves cross, and 'better' silently assumes a scale."
Phase 02 — Mixture of Experts From Scratch
The architecture that lets you add knowledge without adding cost per token — and the one that turns a modelling decision into a memory problem, a communication problem, a stability problem, and a data problem all at once. Feinberg's slides put both halves in one breath: "MoE scaling laws are better, but have implications for token hunger. We're running out of internet!" And the serving half, from the interview: an MoE "uses a lot more parameters," so when you shard it, "that token might live on the first TPU, but it needs to go to the last TPU."
Why this phase exists
MoE is the single most consequential architecture decision in modern pre-training, and it is the one most often described in a way that makes it sound free. It is not free. It is a cost-shifting device, and a pre-training lead has to know exactly where the cost lands:
| You gain | You pay in |
|---|---|
| More parameters at the same FLOPs per token | HBM — every expert must be resident |
| Better loss per unit of compute | Communication — all-to-all every layer if you shard experts |
| Specialization across the corpus | Stability — routers collapse by default |
| Dropped tokens — capacity is a fixed buffer, and overflow is silent | |
Data hunger — MoE's compute-optimal D is larger, and unique data is the scarce resource |
Every one of those five is something you implement and measure in this phase's lab. And the last two are the ones that surprise people: a production MoE can be silently dropping 5–10% of its tokens, and no exception is ever raised.
This phase also sets up the pivot in Feinberg's Flash 2.0 story. The communication cost you compute here is precisely the wall his team hit — and the fix was not a better kernel but changing which axis you shard along, covered in the transcript dissection, Claim 12.
Concept map
token (d_model vector)
│
┌────────▼────────┐
│ ROUTER │ d_model x n_experts — tiny, and the
│ logits→softmax │ most fragile part of the whole model
└────────┬────────┘
│ top-k, renormalized gates
┌────────────────────┼────────────────────┐
▼ ▼ ▼
┌──────────┐ ┌──────────┐ ┌──────────┐
│ EXPERT 0 │ ... │ EXPERT i │ ... │ EXPERT E │ only k of these run
└──────────┘ └──────────┘ └──────────┘
│ │ │
└────────── weighted sum by gate ─────────┘
│
+ shared expert (runs for EVERY token)
│
▼
output
CONSTRAINTS BOLTED ON TOP
├─ capacity = cf × tokens × k / E overflow is DROPPED (silently)
├─ L_balance = E · Σ f_i·P_i 1.0 = perfect; larger = collapsing
├─ L_z = mean(logsumexp²) keeps logits from saturating the softmax
└─ comms = 2 all-to-alls per layer if experts are sharded across chips
What you will be able to do
- Implement a top-k router with renormalized gates, and explain what not renormalizing silently does to the layer's output scale.
- Write the Switch load-balance loss and explain why it multiplies a non-differentiable
token fraction by a differentiable probability mass — and why balanced routing gives
exactly
1.0. - Write the router z-loss and explain the two distinct failures it prevents (bf16 overflow, and a saturated softmax whose gradient vanishes).
- Compute expert capacity, and quantify the drop-rate / wasted-compute trade-off that tuning
capacity_factoractually buys. - Explain why a fully-dropped token is a silent quality loss, and how a shared expert removes that failure mode entirely.
- Separate total parameters (memory) from active parameters (FLOPs) and never confuse them again.
- Compute the all-to-all traffic of naive expert parallelism and state, in seconds, why it is fatal for an interactive product.
- Simulate router collapse and demonstrate the auxiliary loss preventing it.
The lab
| Lab | What you build |
|---|---|
| Lab 01 — Router, Load Balancing, Capacity & the MoE Layer | A complete MoE layer in pure stdlib: stable softmax, top-k router with gate renormalization, Switch load-balance loss, router z-loss, capacity/drop/pad accounting, a full forward pass with shared expert and residual fallback, total-vs-active parameter counting, expert-parallel communication cost, and a router-collapse simulator that shows the aux loss working |
Success criteria. LAB_MODULE=solution pytest test_lab.py -v → 53 passed;
python solution.py prints the eight-part worked example, including the collapse
demonstration.
Deliverables checklist
-
Balanced routing gives a load-balance loss of exactly
1.0in my implementation. -
My softmax survives logits of 1000 without producing
NaN. - I can state the capacity formula and predict the drop rate for a given imbalance.
-
I have watched a router collapse with
aux_weight = 0and recover with0.01. -
I can compute total vs active parameters and say which goes into
6ND. - I can compute the all-to-all bytes for a realistic MoE and convert it to seconds.
- I can explain what a shared expert buys, in one sentence.
Key takeaways
- Parameters scale with
E; FLOPs scale withk. That single sentence is what MoE is. top_k == n_expertsis a dense model wearing an MoE costume. Sparsity ratio 1.0.- Routers collapse by default. Rich-get-richer is the natural dynamic; the auxiliary loss is what stops it, and its coefficient (~0.01) is one of the most finicky hyperparameters in pre-training.
L_balance = 1.0means perfectly balanced. Memorize the calibration point — it makes the dashboard readable at a glance.- Dropped tokens are silent. No exception, no log line; the token rides the residual and
your loss is quietly worse. Watch
drop_ratelike a hawk. - Capacity factor trades dropping against wasted compute. There is no setting that avoids both.
- A shared expert guarantees every token gets some FFN, which removes the fully-dropped failure mode at the cost of always-on compute.
- Naive expert parallelism costs seconds of pure network time. That is the wall, and the fix is to change the sharding axis, not the kernel.
- MoE trades a compute problem for a data problem. Better scaling law, larger optimal
D, and unique tokens are what is actually running out.
References
- Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated MoE Layer, 2017 — https://arxiv.org/abs/1701.06538
- Fedus, Zoph & Shazeer, Switch Transformers, 2021 — https://arxiv.org/abs/2101.03961 (the load-balance loss, capacity factor,
k=1routing) - Zoph et al., ST-MoE: Designing Stable and Transferable Sparse Expert Models, 2022 — https://arxiv.org/abs/2202.08906 (the router z-loss)
- Lepikhin et al., GShard, 2020 — https://arxiv.org/abs/2006.16668 (expert parallelism and the all-to-all)
- Clark et al., Unified Scaling Laws for Routed Language Models, 2022 — https://arxiv.org/abs/2202.01169 (the MoE scaling law Feinberg cites)
- DeepSeek-AI et al., DeepSeek-V3 Technical Report, 2024 — https://arxiv.org/abs/2412.19437 (shared experts, fine-grained experts, auxiliary-loss-free balancing)
- Jiang et al., Mixtral of Experts, 2024 — https://arxiv.org/abs/2401.04088
- Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides
Warmup Guide — Mixture of Experts From Scratch
How to read this. No prior knowledge of MoE, routing, or sparse models is assumed. Everything is built from a dense feed-forward layer upward: what it is, why MoE exists, how the mechanism works underneath, what it costs in production, and the misconception people carry. If you already know MoE, start at Chapter 5 and do not skip Chapter 7 or Chapter 9.
Table of Contents
- Chapter 1: The Problem MoE Solves
- Chapter 2: The Dense Feed-Forward Layer, First
- Chapter 3: The MoE Layer
- Chapter 4: The Router, In Full
- Chapter 5: Router Collapse — the Default Failure
- Chapter 6: The Router Z-Loss
- Chapter 7: Capacity, Dropping, and the Silent Quality Loss
- Chapter 8: Total vs Active Parameters
- Chapter 9: The Communication Wall
- Chapter 10: MoE Scaling Laws and "Running Out of Internet"
- Chapter 11: The Modern Refinements
- Lab Walkthrough
- Success Criteria
- Interview Q&A
- Tips & Takeaways
- References
Chapter 1: The Problem MoE Solves
The welded constraint
In a dense transformer, every parameter participates in every token. Want the model to know more? Add parameters. But now every token costs more to process — during training and forever afterward during serving.
From Phase 00: C = 6ND for training, 2N FLOPs per token for inference. Both scale linearly
with N. Capacity and cost are welded together.
That is a genuine problem, because the two things you want are in tension:
- Knowledge and capability want many parameters.
- Latency and serving cost want few parameters.
The observation MoE exploits
Here is the insight: not every token needs every parameter. The weights that help predict the next token in a Python function are largely not the weights that help with 18th-century French poetry. In a dense model, both sets fire for both inputs, and most of that computation is wasted on any individual token.
Conditional computation breaks the weld: store many parameters, but activate only a small, input-dependent subset per token.
DENSE MoE
────── ───
params = 8 units params = 64 units (8x the knowledge)
FLOPs = 8 units/token FLOPs = 8 units/token (same cost!)
memory = 64 units <- the bill arrives here
Why it exists historically
The idea is old (Jacobs et al., 1991). Shazeer et al. (2017) made it work at scale for
language, Switch Transformers (2021) simplified it to k = 1 and made it stable, GShard
(2020) built the distributed machinery, and by 2024 essentially every frontier model is an
MoE — including, as Feinberg notes, the Gemini 2.0 series.
The misconception
"MoE gives you a bigger model for free."
It gives you a bigger model for the same FLOPs. It costs you memory, communication, stability, dropped tokens, and data hunger. This entire warmup is an accounting of where that bill lands.
Chapter 2: The Dense Feed-Forward Layer, First
You cannot understand what MoE replaces without knowing what it replaces.
What it is
Each transformer block has two sub-layers: attention (which routes information between positions) and a feed-forward network / MLP (which transforms each position independently).
def dense_ffn(token, w_in, w_out):
"""Expand to d_ff, apply a nonlinearity, project back to d_model."""
hidden = matvec(w_in, token) # d_model -> d_ff (usually d_ff ~ 4*d_model)
activated = [max(0.0, h) for h in hidden]
return matvec(w_out, activated) # d_ff -> d_model
Why it is the target
From Phase 00's parameter count, for a typical block:
attention projections : 41,943,040 (24%)
MLP : 135,266,304 (76%) <- MoE replaces THIS
norms : 8,192 ( 0%)
Two reasons the MLP is the right target:
- It is where the parameters are — about three-quarters of the block.
- It is position-independent. Each token goes through the MLP alone, with no interaction between positions. So you can route each token to a different expert without breaking anything. Attention cannot be split this way, because it is precisely the part that mixes positions.
This is the single best reason MoE targets the MLP and not attention, and it is worth being able to say out loud in an interview.
The misconception
"The MLP is just a nonlinearity; attention is where the intelligence is."
Empirically, most factual knowledge lives in MLP weights — the model-editing literature (ROME, MEMIT) locates and edits facts specifically there. Attention decides what to look at; the MLP decides what to think about it.
Chapter 3: The MoE Layer
The structure
Replace the one MLP with E parallel MLPs ("experts") plus a small router.
token
│
┌────────▼────────┐
│ ROUTER │ E scores, one per expert
│ (d_model x E) │
└────────┬────────┘
│ softmax, take top-k, renormalize
┌────────┬───────┼───────┬────────┐
▼ ▼ ▼ ▼ ▼
E0 E1 E2 ... E7
· ✓ · · ✓ only 2 of 8 run
└───────┬────────┘
▼ weighted sum by gate
output
The forward pass, exactly
def moe_layer(token, router_weights, experts, top_k=2):
# 1. score every expert
logits = matvec(router_weights, token) # E scores
probs = softmax(logits)
# 2. pick the best k, renormalize their gates to sum to 1
chosen = sorted(range(len(probs)), key=lambda i: -probs[i])[:top_k]
total = sum(probs[i] for i in chosen)
gates = {i: probs[i] / total for i in chosen}
# 3. run ONLY those k, and combine by gate weight
out = [0.0] * len(token)
for i, g in gates.items():
contribution = dense_ffn(token, experts[i]["w_in"], experts[i]["w_out"])
out = [o + g * c for o, c in zip(out, contribution)]
return out
The k choice
k | Name | Trade-off |
|---|---|---|
| 1 | Switch routing | Cheapest. Sharpest specialization. Hardest to train — a single routing mistake has no backup. |
| 2 | The common default | Two experts give the gradient two paths and the token a fallback. Used by Mixtral, DeepSeek, most production models. |
| 4–8 | Fine-grained MoE | Used with many small experts (DeepSeek-V3 uses 8 of 256). More combinations, better specialization, more routing overhead. |
E | Dense | Sparsity ratio 1.0. You have built a dense model with extra steps. |
Why renormalize the gates
This is subtle and worth being precise about. Suppose the softmax gives expert 6 a probability of 0.575 and expert 4 a probability of 0.151, and you take the top 2.
- Renormalized: gates become
0.79and0.21, summing to 1. The layer output has a consistent scale regardless of how confident the router was. - Raw: gates stay
0.575and0.151, summing to0.726. The layer output is scaled down by 27% — and by a different amount for every token, depending on how much probability mass leaked to the experts you did not pick.
That varying scale interacts badly with the residual stream and layer norms. Almost all production implementations renormalize. The lab implements both so you can see the difference, and a test asserts it.
Chapter 4: The Router, In Full
What it is
The smallest and most important component. Just d_model × E parameters — for a model with
d_model = 7168 and E = 256, that is 1.8M parameters, versus ~44M for a single expert. It
is 0.04% of the layer and it decides everything.
Numerical stability is not optional here
def softmax(xs):
"""Max-subtraction is what stops exp() overflowing to inf."""
m = max(xs)
exps = [math.exp(x - m) for x in xs]
s = sum(exps)
return [e / s for e in exps]
Why this matters more for a router than elsewhere: router logits are unconstrained and
there is a positive feedback loop pushing them up (Chapter 5). In bf16, exp(90) already
overflows. A naive softmax gives NaN, the NaN propagates through the residual stream, and
your 40-day training run is dead. The lab tests this with logits of 1000.
Subtracting the max changes nothing mathematically —
softmax(x) = softmax(x - c)for any constantc— and changes everything numerically. A test asserts this shift-invariance.
Ties and determinism
Two experts can score identically (especially early in training, when weights are near zero).
If your tie-break is dict ordering or an unstable sort, the same input routes differently on
different runs, and you lose reproducibility — which you need for debugging a run that costs
$30M. Break ties by expert index, always.
Chapter 5: Router Collapse — the Default Failure
This is the chapter that matters most. A router left alone will collapse, and it is not a rare pathology — it is the natural dynamic.
The mechanism
expert i is randomly slightly better at step 0
│
▼
it wins slightly more tokens
│
▼
it receives more gradient, so it trains faster
│
▼
it becomes genuinely better
│
▼
it wins even more tokens ──────────┐
▲ │
└────────────────────────┘
RICH GET RICHER
Within a few thousand steps you have one expert doing everything and E − 1 dead ones. You
paid for E× the memory, E× the checkpoint size, and E× the sharding complexity, and you
are running a dense model.
The fix: an auxiliary load-balancing loss
def load_balance_loss(assignments, all_probs, n_experts):
"""Switch-style: L_aux = E * sum_i f_i * P_i
f_i = fraction of token-SLOTS routed to expert i (discrete -> NO gradient)
P_i = mean router PROBABILITY on expert i (continuous -> differentiable)
"""
total_slots = sum(len(a) for a in assignments)
f = [0.0] * n_experts
for row in assignments:
for e in row:
f[e] += 1.0 / total_slots
P = [0.0] * n_experts
for probs in all_probs:
for i, p in enumerate(probs):
P[i] += p / len(all_probs)
return n_experts * sum(fi * Pi for fi, Pi in zip(f, P))
Why the product f_i · P_i, and not something simpler?
This design is genuinely clever and the reasoning is a great interview answer.
You want to penalize imbalance in f — the actual token counts. But f comes from a
top-k operation, which is discrete: it has zero gradient almost everywhere, so you cannot
backpropagate through it.
P — the mean softmax probability — is differentiable, but on its own it is a weak signal:
the router could keep P uniform while its top-k choices remain lopsided.
Multiplying them gives you both: the loss is scaled by the real imbalance (f) while the
gradient flows through P. When expert i is over-subscribed, f_i is large, so the
gradient pushing P_i down is large.
The calibration point — memorize this
Under perfect balance, f_i = P_i = 1/E for every i:
$$ L_{\text{aux}} = E \sum_{i=1}^{E} \frac{1}{E}\cdot\frac{1}{E} = E \cdot E \cdot \frac{1}{E^2} = 1.0 $$
1.0 means perfectly balanced. Larger is worse. That single fact makes an MoE training
dashboard readable at a glance — and the lab has a test asserting exactly 1.0.
Measured in the lab:
balanced : L_aux=1.0000 max/mean=1.00
collapsed: L_aux=7.4400 max/mean=8.00 dead=7
Watching it happen
The lab simulates the feedback loop directly:
aux_weight=0.000: max/mean after 600 steps = 4.10 COLLAPSED
aux_weight=0.002: max/mean after 600 steps = 1.43 healthy
aux_weight=0.005: max/mean after 600 steps = 1.06 healthy
aux_weight=0.010: max/mean after 600 steps = 1.00 healthy
The coefficient is a real hyperparameter
L_total = L_task + alpha * L_balance + gamma * L_z
alpha ≈ 0.01 is the common value, and it is finicky:
- Too small → collapse. You lose the entire benefit of MoE.
- Too large → you force tokens to experts that are wrong for them, purely to satisfy the balance constraint. Quality drops. The router is now optimizing bookkeeping instead of prediction.
What to put on the dashboard:
max_over_meanutilization andL_balance. Ifmax_over_meanclimbs past ~3, intervene now — not at the next checkpoint. This is exactly the sort of thing the five-person Flash 2.0 rotation was watching for.
Chapter 6: The Router Z-Loss
What it is
$$ L_z = \frac{1}{T}\sum_{t} \left(\log \sum_i e^{z_{t,i}}\right)^2 $$
A penalty on the magnitude of the router logits, introduced in ST-MoE (Zoph et al., 2022).
def router_z_loss(all_logits):
return sum(logsumexp(row) ** 2 for row in all_logits) / len(all_logits)
Why it exists — two distinct failures
1. Numerical overflow. Router logits are unconstrained and the collapse dynamic pushes them
up. In bf16, exp(90) overflows. NaN in the router means NaN in the residual stream means
a dead run.
2. A saturated softmax has no gradient. This one is subtler and more damaging. If logits
are [100, 2, 1], the softmax is [1.0, 0.0, 0.0] to floating-point precision. The
derivative of softmax is p_i(δ_ij − p_j) — when p_i is 1 or 0, that is zero. The router
stops learning. Routing freezes into whatever pattern it happened to have, permanently.
The lab shows the magnitude growing sharply:
logit scale x 1.0 -> z_loss= 9.999
logit scale x 5.0 -> z_loss= 160.348
logit scale x20.0 -> z_loss= 2527.342
gamma ≈ 1e-3 is typical — small, because you are only trying to keep logits in a sane range,
not to shape the distribution.
The general lesson, which transfers well beyond MoE: whenever a small number of unconstrained logits control a discrete decision, add a magnitude penalty. The same reasoning shows up in attention (QK-norm) for the same reason.
Chapter 7: Capacity, Dropping, and the Silent Quality Loss
The constraint
Accelerators want fixed-shape tensors. You cannot allocate "however many tokens happened to route here" — the shape has to be known ahead of time. So each expert gets a fixed buffer:
$$ \text{capacity} = \text{capacity_factor} \times \frac{\text{tokens} \times k}{E} $$
The fraction is the count each expert would get under perfect balance. capacity_factor
(typically 1.0–2.0) is the safety margin for imperfect balance.
The two failure directions
capacity per expert ────────────────────────►
┌──────────────────────────────────────┐
│████████████████░░░░░░░░░░░░░░░░░░░░░░│ under-subscribed:
│ real tokens PADDING (waste) │ you compute on zeros
└──────────────────────────────────────┘
┌──────────────────────────────────────┐
│██████████████████████████████████████│▓▓▓▓ over-subscribed:
│ buffer full │DROP overflow is discarded
└──────────────────────────────────────┘
Measured in the lab on a realistically imbalanced batch:
capacity_factor=1.00: cap= 16 dropped= 12 ( 9.4%) padded= 12 buffer_util=90.6%
capacity_factor=1.25: cap= 20 dropped= 3 ( 2.3%) padded= 35 buffer_util=78.1%
capacity_factor=2.00: cap= 32 dropped= 0 ( 0.0%) padded=128 buffer_util=50.0%
There is no setting that avoids both. Going from 1.0 to 2.0 eliminates dropping and halves your effective compute utilization. That trade-off is the hyperparameter.
Why dropping is the dangerous one
Trace what happens to a dropped token in the lab's forward pass:
if not got_any:
# Every slot dropped: the token skips the FFN entirely and rides the
# residual. No error, no log line. This is the silent failure.
acc = list(token)
fully_dropped += 1
No exception. No warning. No log line. The token passes through the layer unchanged on the residual connection. Training continues. Your loss is very slightly worse than it should be, and nothing tells you why.
This is a genuine, live source of quality loss in production MoE models — and it is why
drop_rate belongs on the dashboard next to max_over_mean.
What a shared expert buys
A shared expert (DeepSeek-V3's design) runs for every token, unconditionally, in addition to the routed ones:
if shared_expert is not None:
acc = [a + c for a, c in zip(acc, expert_forward(shared_expert, token))]
Two things fall out:
- No token can ever be fully dropped. The failure mode above is structurally eliminated.
- The routed experts can specialize harder, because the shared expert absorbs the common, general-purpose transformation that every token needs. You are no longer forcing every expert to independently learn the basics.
The cost: that expert's FLOPs are paid on every token, always. It is a small, always-on tax that buys a large reduction in variance.
Chapter 8: Total vs Active Parameters
The arithmetic most commonly botched in modern LLM discussion, and it is off by 10–20×.
def moe_parameter_counts(d_model, d_ff, n_experts, top_k, shared_experts=0):
one_expert = 2 * d_model * d_ff
router = d_model * n_experts
total = router + (n_experts + shared_experts) * one_expert # -> MEMORY
active = router + (top_k + shared_experts) * one_expert # -> FLOPs
return {"total": total, "active": active, "sparsity_ratio": total / active}
The rule, and there are no exceptions:
| Question | Use |
|---|---|
How much compute to train? C = 6ND | active |
How much compute to serve? 2N/token | active |
| How much HBM do I need? | total |
| How big is the checkpoint? | total |
| How do I shard it? | total |
A 700B-total / 52B-active model trains like a 52B model and stores like a 700B one. Get this backwards and your capacity plan is wrong by an order of magnitude.
From the lab, at toy scale:
top_k=1: total= 2112 active= 320 sparsity=6.60x
top_k=2: total= 2112 active= 576 sparsity=3.67x
top_k=8: total= 2112 active= 2112 sparsity=1.00x
Note the last row: top_k == n_experts gives sparsity exactly 1.0. A dense model wearing an
MoE costume. The lab tests this boundary explicitly, because it is the sanity check that
proves your accounting is right.
Chapter 9: The Communication Wall
Now the cost that shaped Feinberg's Flash 2.0 story.
Why sharding is forced
All E experts must be resident in HBM. A 700B-parameter MoE in bf16 is 1.4 TB. A TPU v5e has
16 GB; an H100 has 80 GB. The model cannot fit on one chip, so experts are distributed
across chips — expert parallelism.
What that costs, per layer
Layer ℓ: token lives on chip 0
router says "expert 37", which lives on chip 4
──> send the activation to chip 4 [NETWORK]
chip 4 computes
──> send the result back to chip 0 [NETWORK]
Layer ℓ+1: router says "expert 12" (chip 1) ... [NETWORK] [NETWORK]
... for every one of ~60 layers.
Feinberg's description: "that token might live on the first TPU, but it needs to go to the last TPU." The collective is an all-to-all — every chip sends a different slice to every other chip — which is the most expensive collective there is, and its cost, he notes, "increases dramatically with N."
The arithmetic
def expert_parallel_comm_bytes(n_tokens, d_model, n_layers, top_k, bytes_per_elem=2):
per_hop = n_tokens * d_model * bytes_per_elem * top_k
return 2 * n_layers * per_hop # dispatch + combine, every layer
From the lab, for a realistic prefill:
8192 tokens, d_model=8192, 60 layers, top_k=2
bytes across the interconnect : 32.2 GB
time at 100 Gbps : 2.578 s
Two and a half seconds of pure network time before a single useful FLOP. For an interactive product with a sub-second budget (Phase 00's napkin math), that is fatal.
The fix, in one sentence
Stop sharding experts across chips. Shard layers instead, and stream chunks of the prompt through the resulting pipeline so the transfers overlap with computation. That is pipelined prefill, credited in the interview to Geng Yan, and it is what made an MoE Gemini Flash servable.
The full treatment — including why it works for prefill and not for decode — is in the transcript dissection, Claim 12.
The transferable lesson, and it is the real one: the fix was not a better kernel or a better model. It was changing which axis you shard along, chosen with knowledge of which phase of inference you are in. That is what "inference co-design" means in practice.
Chapter 10: MoE Scaling Laws and "Running Out of Internet"
Feinberg's slide states both halves:
"MoE scaling laws are better, but have implications for token hunger. We're running out of internet! ... Notice relative data hunger compared to dense! At same active param count and fixed 100B token training, MoE 64E improves on dense."
Decoded
The good half. For a fixed compute budget, an MoE reaches a lower loss than a dense model. The scaling law is strictly better. At equal active parameters and equal tokens, 64 experts beats dense. This is why everyone switched.
The bill. The compute-optimal D for an MoE is larger than for a dense model. MoE has
more capacity to fill, and filling capacity takes tokens. So MoE converts "we have compute"
into "we need more unique tokens" — and unique, high-quality tokens are the resource that is
actually scarce.
Hence: MoE trades a compute problem for a data problem. Which is a good trade right up until you run out of data, and then it is not.
That is the bridge to the rest of Feinberg's agenda — multimodal data, synthetic data, and the
data-constrained scaling law L(N, U, R) where U is unique tokens and R is repeats. He
notes that data work is where he spends "probably half my focus this year so far." The
pre-training lead's time goes to data, not architecture.
The reference for the routed-model scaling law is Clark et al., Unified Scaling Laws for Routed Language Models (2022). The data-constrained law is Muennighoff et al. (2023). Both are on his slide's resource list.
Chapter 11: The Modern Refinements
What production MoE looks like in 2024–2025, beyond the Switch baseline.
| Refinement | What it does | Why |
|---|---|---|
| Fine-grained experts | Many small experts (256) instead of few large ones (8), with higher k | More combinations of experts per token → better specialization at the same active parameter count |
| Shared experts | One or more experts that always run | Removes the fully-dropped failure mode; lets routed experts specialize harder (Chapter 7) |
| Auxiliary-loss-free balancing | A per-expert bias added to router logits, adjusted by observed load, instead of an aux loss | The aux loss damages quality by forcing wrong routing. A bias term balances load without adding a gradient that fights the task loss. (DeepSeek-V3) |
| Expert-choice routing | Invert it: each expert picks its top tokens, rather than each token picking experts | Balance is guaranteed by construction — no aux loss, no dropping. Cost: it needs the whole batch at once, so it does not work for autoregressive decode. |
| Dropless MoE | Variable-size expert buffers via block-sparse kernels | Eliminates dropping and padding entirely — at the price of needing custom kernels (which is Phase 11's territory, and a live example of "missing kernels kill good ideas") |
| Upcycling | Initialize MoE experts from a trained dense checkpoint | Skips the expensive early phase; the experts start from something that already works |
Notice the pattern across that table. Four of the six refinements exist to fix problems created by the two mechanisms in Chapters 5 and 7 — the auxiliary loss and the capacity buffer. Understand those two deeply and the rest of the literature reads as commentary.
Lab Walkthrough
Lab 01 — Router, Load Balancing, Capacity & the MoE Layer
Implement in this order:
softmax,logsumexp,matvec,relu. Gettest_softmax_survives_huge_logitspassing first — max-subtraction is the lesson, not a detail.router_logits,route_token,route_batch. Watch therenormalizeflag; both behaviours are tested. Break ties by expert index.load_balance_loss. Aim for the calibration point:test_balanced_routing_gives_aux_loss_of_exactly_one. If you do not get exactly1.0, the formula is wrong.router_z_loss,expert_utilization.expert_capacity,apply_capacity. Note thatcapacity_factor < 1.0must raise — it drops tokens even under perfect balance, which is never intended.moe_forward. The residual fallback for fully-dropped tokens is the important line.moe_parameter_counts,expert_parallel_comm_bytes,comm_seconds.simulate_collapse. The money test.
The traps:
- A naive softmax gives
NaNon large logits and the test will catch you. - The load-balance loss divides
fby total slots (tokens × k), not by token count. apply_capacityis first-come-first-served in token order — deterministic, and the tests depend on it.top_k == n_expertsmust give sparsity exactly1.0; if it does not, your router or parameter count is wrong.- Everything random goes through a seeded
random.Random(seed).
Success Criteria
-
LAB_MODULE=solution pytest test_lab.py -v→ 53 passed. -
Your
lab.pyreaches 53 passed. -
python solution.pyruns and you can explain all eight sections. -
Balanced routing gives a load-balance loss of exactly
1.0in your code. -
You can explain why the aux loss multiplies
fbyPrather than penalizingfalone. -
You have watched a router collapse at
aux_weight = 0and recover at0.01. - You can state the capacity formula and the drop/pad trade-off from memory.
-
You can compute total vs active parameters and say which goes into
6ND. - You can compute all-to-all bytes for a realistic MoE and convert to seconds.
Interview Q&A
Q: What is a Mixture of Experts and what does it buy you?
Replace each transformer block's MLP with E parallel MLPs plus a small router that picks the
top-k per token. Parameters scale with E; FLOPs scale with k. So you get far more
capacity at the same compute per token. It targets the MLP rather than attention for two
reasons: the MLP is ~76% of the block's parameters, and it is position-independent, so routing
each token separately does not break anything — attention is precisely the part that mixes
positions and cannot be split this way.
Q: What breaks first when you train an MoE?
The router collapses. It is a rich-get-richer loop: one expert is randomly slightly better, so
it wins more tokens, so it trains more, so it becomes genuinely better, so it wins more. Within
a few thousand steps you have one live expert and E−1 dead ones — you paid for E× the memory
to run a dense model. The fix is an auxiliary load-balancing loss, L = E·Σ f_i·P_i, with a
coefficient around 0.01.
Q: Why does that loss multiply a token fraction by a probability?
Because f, the fraction of tokens routed to each expert, comes from a top-k, which is
discrete and has no gradient. P, the mean softmax probability, is differentiable but a weak
signal on its own — the router could keep P uniform while its actual choices stay lopsided.
Multiplying gives you a loss scaled by the real imbalance with gradient flowing through the
differentiable part. And it is calibrated: perfect balance gives exactly 1.0.
Q: What is capacity factor and what happens if you get it wrong?
Each expert gets a fixed buffer of cf × tokens × k / E slots, because hardware needs
fixed-shape tensors. Too low and tokens overflow and are dropped — silently: the token
skips the FFN, rides the residual, no error is raised, and your quality is quietly worse. Too
high and buffers are padded with zeros and you waste compute. There is no setting that avoids
both, so you tune it against measured drop rate. A shared expert that runs for every token
eliminates the fully-dropped case structurally.
Q: A 700B-parameter MoE with 8 of 256 experts active. How much compute to train it?
6ND with N = active parameters, not total. If active is ~52B, it trains like a 52B
dense model. But it stores like a 700B model — 1.4 TB in bf16 — which is what dictates your
sharding, your checkpoint size, and your HBM budget. Active for FLOPs, total for memory.
Q: Why is MoE hard to serve? All experts must be HBM-resident, so a large MoE cannot fit on one chip and you shard experts across chips. Then every layer, every token must be sent to whichever chip holds its expert and the result sent back — two all-to-all collectives per layer. For an 8k-token prefill on a 60-layer model that is tens of gigabytes and seconds of pure network time, which is fatal for an interactive product. The fix Google used for Flash 2.0 was to shard layers instead of experts and stream prompt chunks through the resulting pipeline, so the transfers hide behind computation. Notably it works for prefill — which is compute-bound and has thousands of tokens to chunk — and not for decode, which produces one token at a time and has nothing to overlap with.
Q: If MoE is strictly better per FLOP, why isn't everything an MoE? Three reasons. The memory and sharding complexity is real. The training instability is real — you have added a discrete decision to a system that otherwise has none. And the compute-optimal token count is larger for an MoE, so it converts a compute problem into a data problem — and unique high-quality data is the resource actually running out.
Q: What is the router z-loss for?
It penalizes the magnitude of router logits. Two failures it prevents: numerical overflow —
router logits are unconstrained, the collapse dynamic pushes them up, and exp() overflows in
bf16, putting NaN through the residual stream; and softmax saturation — once the softmax is
effectively one-hot, its derivative p(1−p) is zero, the router stops learning, and routing
freezes permanently.
Tips & Takeaways
Tips
- Put
max_over_meananddrop_rateon the dashboard. They are the two numbers that tell you an MoE is failing, and both fail silently otherwise. - Memorize
L_aux = 1.0means balanced. It makes the metric readable instantly. - Always ask "total or active?" the moment someone quotes an MoE parameter count.
- Break router ties deterministically. Reproducibility on a $30M run is not negotiable.
- Never write a softmax without max-subtraction. Especially in a router.
- When you see a new MoE paper, ask which of the two core problems it is fixing — balance or capacity. Most of the literature is one or the other.
- Compute the all-to-all bytes before you commit to an architecture, not after.
Takeaways
- Parameters scale with
E; FLOPs scale withk. That is MoE, entire. - MoE targets the MLP because it holds ~76% of the parameters and is position-independent.
- Renormalizing the top-k gates keeps the output scale consistent; not doing so silently attenuates the layer by a per-token amount.
- Routers collapse by default. The auxiliary loss is what prevents it, and
1.0is balanced. - The
f·Pproduct exists to get a gradient through a discrete decision. - The z-loss prevents overflow and softmax saturation — two different failures.
- Capacity trades dropping against wasted compute. Dropping is silent, which makes it worse.
- A shared expert eliminates fully-dropped tokens and lets the routed experts specialize.
- Active for FLOPs, total for memory. Off by 10–20× if you swap them.
- Naive expert parallelism costs seconds of network time. The fix is a different sharding axis, not a better kernel.
- MoE trades a compute problem for a data problem — and data is what is running out.
References
Foundational
- Jacobs et al., Adaptive Mixtures of Local Experts, 1991 — the original idea
- Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated MoE Layer, 2017 — https://arxiv.org/abs/1701.06538
- Lepikhin et al., GShard: Scaling Giant Models with Conditional Computation, 2020 — https://arxiv.org/abs/2006.16668 (expert parallelism, capacity, the all-to-all)
- Fedus, Zoph & Shazeer, Switch Transformers, 2021 — https://arxiv.org/abs/2101.03961 (
k=1, the load-balance loss, capacity factor)
Stability and refinements
- Zoph et al., ST-MoE: Designing Stable and Transferable Sparse Expert Models, 2022 — https://arxiv.org/abs/2202.08906 (the router z-loss)
- Zhou et al., Mixture-of-Experts with Expert Choice Routing, 2022 — https://arxiv.org/abs/2202.09368
- Gale et al., MegaBlocks: Efficient Sparse Training with Mixture-of-Experts, 2022 — https://arxiv.org/abs/2211.15841 (dropless MoE via block-sparse kernels)
- Komatsuzaki et al., Sparse Upcycling, 2022 — https://arxiv.org/abs/2212.05055
Production models
- Jiang et al., Mixtral of Experts, 2024 — https://arxiv.org/abs/2401.04088
- DeepSeek-AI et al., DeepSeek-V3 Technical Report, 2024 — https://arxiv.org/abs/2412.19437 (fine-grained + shared experts, auxiliary-loss-free balancing)
Scaling
- Clark et al., Unified Scaling Laws for Routed Language Models, 2022 — https://arxiv.org/abs/2202.01169
- Muennighoff et al., Scaling Data-Constrained Language Models, 2023 — https://arxiv.org/abs/2305.16264
Primary source for this phase's framing
- Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides
- Developing Dev interview · video
Lab 01 — Router, Load Balancing, Capacity & the MoE Layer
Build a complete Mixture-of-Experts layer in pure Python, then break it on purpose: watch the router collapse, watch tokens get dropped silently, and watch the auxiliary loss save it.
The problem
MoE is usually explained as "more parameters, same FLOPs." That framing is true and useless, because it omits every part that makes MoE hard: routers collapse by default, capacity buffers drop tokens without raising anything, and sharding experts across chips costs seconds of network time per forward pass.
This lab implements the mechanism and all three failure modes, so that "MoE" stops being a word and becomes something you have debugged.
What you build
| Group | Functions | The idea |
|---|---|---|
| Primitives | softmax, logsumexp, matvec, relu | Max-subtraction is the lesson, not a detail |
| Router | router_logits, route_token, route_batch | Top-k, gate renormalization, deterministic ties |
| Aux losses | load_balance_loss, router_z_loss, expert_utilization | What stops collapse, and the dashboard numbers |
| Capacity | expert_capacity, apply_capacity | Fixed buffers, dropping, padding |
| The layer | make_expert, expert_forward, moe_forward, total_training_loss | Full forward pass with shared expert and residual fallback |
| Accounting | moe_parameter_counts, expert_parallel_comm_bytes, comm_seconds | Total vs active; the communication wall |
| The failure | simulate_collapse | Rich-get-richer, and the aux loss preventing it |
Key concepts
| Concept | Why it is in this lab |
|---|---|
| Max-subtracted softmax | Router logits grow without bound; naive exp gives NaN and kills the run |
| Gate renormalization | Without it the layer output is silently attenuated by a per-token amount |
L_aux = E·Σ f_i·P_i | Discrete f supplies the signal, differentiable P supplies the gradient |
L_aux = 1.0 | The calibration point: exactly balanced. Larger is worse. |
| Router z-loss | Prevents bf16 overflow and softmax saturation — two distinct failures |
| Capacity factor | Trades dropped tokens against padded (wasted) compute. No free setting. |
| Silent dropping | A fully-dropped token rides the residual. No error. Quality quietly drops. |
| Shared expert | Structurally eliminates the fully-dropped case |
| Total vs active | Active → FLOPs (6ND); total → HBM. Off by 10–20× if swapped. |
| All-to-all cost | 2 collectives per layer under expert parallelism — the Flash 2.0 wall |
Files
| File | What it is |
|---|---|
| lab.py | Your implementation. Signatures, docstrings and validation contracts given. |
| solution.py | Reference. python solution.py runs an eight-part worked example. |
| test_lab.py | 53 tests: happy path, validation, boundaries, invariants, determinism. |
| requirements.txt | pytest only. Pure stdlib otherwise. |
Run
pytest test_lab.py -v # your lab.py — red until you implement
LAB_MODULE=solution pytest test_lab.py -v # the reference — must be green (53 passed)
python solution.py # the worked example
Where to start
softmax→ gettest_softmax_survives_huge_logitsgreen. Logits of 1000 must not produceNaN. This is the most common real-world MoE crash.route_token→test_renormalized_gates_sum_to_oneand itsrawcounterpart. Both behaviours are tested because both exist in the wild.load_balance_loss→ aim fortest_balanced_routing_gives_aux_loss_of_exactly_one. If you do not get exactly1.0, your formula is wrong. Rememberfdivides by total slots (tokens × k), not token count.apply_capacity→ first-come-first-served in token order; the tests depend on that determinism.moe_forward→ the residual fallback for fully-dropped tokens is the line that matters.simulate_collapse→ the money test.
The traps:
- Naive softmax →
NaNon large logits. Tested. - Dividing
fby token count instead of slot count → yourL_auxwill not be1.0. capacity_factor < 1.0must raise — it drops tokens even under perfect balance.top_k == n_expertsmust give sparsity exactly1.0. If not, your accounting is off.- Unstable tie-breaking in the router destroys reproducibility. Break by expert index.
Success criteria
-
LAB_MODULE=solution pytest test_lab.py -v→ 53 passed. -
Your
lab.pyreaches 53 passed. -
python solution.pyruns and you can explain all eight sections. -
You have seen
aux_weight=0.0collapse tomax/mean = 4.10and0.01hold at1.00. - You can state the capacity formula and predict a drop rate from an imbalance.
- You can convert an MoE's all-to-all traffic into seconds and say why it is fatal.
How this maps to the real stack
| This lab | The real thing | Where the miniature lies |
|---|---|---|
route_token / route_batch | Megatron-LM MoELayer, DeepSpeed-MoE, Mixtral's MoeBlock, JAX/Flax MoE | Real routers run fused on-device with the top-k in a kernel. Ours is a Python loop — same math, ~10⁶× slower. |
load_balance_loss | switch_load_balancing_loss in Megatron/DeepSpeed | Identical formula. Real implementations compute it per-device and all-reduce, which introduces subtleties about whether balance is global or per-shard — a genuine source of production bugs. |
apply_capacity | GShard/Switch dispatch-and-combine via one-hot matmuls or gather/scatter | Real dispatch builds a [E, capacity, d_model] tensor with one-hot matmuls so it stays differentiable and fused. Ours is a Python list walk. |
expert_forward | A grouped/batched GEMM over experts, or block-sparse kernels (MegaBlocks) | Real experts run as one batched matmul, not a loop. This is exactly where "missing kernels kill good ideas" bites — dropless MoE needs custom block-sparse kernels to exist at all. |
expert_parallel_comm_bytes | NCCL/all_to_all in Megatron; XLA collectives on TPU | Ours ignores topology (a torus vs a fat-tree changes the constant a lot), overlap, and compression. Directionally right, not a substitute for profiling. |
simulate_collapse | The real dynamic, observed on training dashboards | A crude ODE-ish model, tuned to show the qualitative behaviour. Real collapse depends on data ordering, init, and LR schedule. |
What is not a lie: the load-balance formula and its 1.0 calibration, the capacity
formula, the total-vs-active distinction, and the all-to-all byte count. Those are exact, and
those are what get asked about.
Extensions
- Implement expert-choice routing. Invert the assignment — each expert picks its top tokens. Balance becomes guaranteed by construction, no aux loss, no dropping. Then discover why it cannot be used for autoregressive decode.
- Implement auxiliary-loss-free balancing (DeepSeek-V3): maintain a per-expert bias added to router logits, nudged up or down by observed load. Compare final quality against the aux loss at matched balance — the claim is that the aux loss costs quality by forcing wrong routing, and you can measure that.
- Add a real gradient. Wire this into Phase 03's autograd (or PyTorch) and actually train a tiny MoE on a toy task. Watch the aux-loss coefficient matter.
- Model the pipelined-prefill fix. Extend
expert_parallel_comm_byteswith a pipelined alternative that shards layers instead of experts, and reproduce the order-of-magnitude improvement in the transcript dissection's Claim 12. - Sweep the capacity factor against real routing entropy. Generate routings at varying
levels of imbalance and plot drop rate vs
capacity_factor— that curve is the one you actually use to pick the hyperparameter.
Interview / resume bullets
- "Implemented a Mixture-of-Experts layer from first principles in pure Python — numerically stable top-k routing with gate renormalization, Switch-style load-balancing loss, ST-MoE router z-loss, capacity-factor dispatch with drop/pad accounting, shared experts, and expert-parallel communication cost modelling — verified by 53 tests covering boundary and determinism cases."
- "Reproduced and quantified the three canonical MoE failure modes: router collapse via rich-get-richer feedback (and its prevention by an auxiliary loss), silent token dropping through fixed capacity buffers, and the all-to-all communication wall of naive expert parallelism (32 GB and ~2.6 s per 8k-token prefill on a 60-layer model)."
- Interview-ready: "Active parameters for FLOPs, total parameters for memory. A 700B/52B MoE trains like a 52B model and stores like a 700B one — and confusing the two makes your capacity plan wrong by an order of magnitude."
Phase 03 — Roofline, MFU & Inference Co-Design
The phase where a product requirement becomes an architecture constraint. Feinberg's second research vertical is building "neural architectures that are efficient to run inference on" by choosing network topology, "shapes of the matrices" and "attention shapes, num heads" that "fully utilize the hardware." This phase gives you the arithmetic that makes those choices defensible instead of aesthetic — and reproduces the napkin math from his talk that turns "we want a real-time agent" into "the model must be smaller."
Why this phase exists
Two models with identical parameter counts can differ by 3× in serving speed, purely from shape choices made before training. And after training, the shape is frozen forever. So the decisions in this phase are among the few in the entire pipeline that are genuinely irreversible.
There is also a widespread confusion this phase kills. Feinberg spends part of the interview correcting people who see a "low" MFU number and conclude incompetence. To hit 100% you would need to be "doing a bunch of matmuls in a loop without reading any memory" — which is not a neural network, because real nets "have to apply activation functions or do attention or write intermediate outputs." 35% MFU is an accounting identity, not a grade. Once you can produce the breakdown, that number stops being a scoreboard and becomes an agenda.
The third thing this phase establishes is the split that governs all of serving:
PREFILL DECODE
processes the whole prompt at once generates one token at a time
COMPUTE bound MEMORY-BANDWIDTH bound
high arithmetic intensity 1–50 FLOP/byte (far below any ridge)
fix it with pipelining, chunking fix it with batching, GQA, quantization
They are two different machines. Applying one's optimization to the other is the most common serving mistake there is, and the reason prefill/decode disaggregation exists.
Concept map
PRODUCT REQUIREMENT HARDWARE
"respond in < 1 second" peak FLOP/s, HBM bandwidth, HBM capacity
│ │
└──────────────┬─────────────────────────┘
▼
┌─────────────────┐
│ ROOFLINE │ arithmetic intensity = FLOPs / bytes
│ ridge point = │ below ridge -> reduce BYTES
│ peak / bandwidth│ above ridge -> reduce FLOPs
└────────┬────────┘
│
┌────────────────┴────────────────┐
▼ ▼
┌───────────┐ ┌───────────┐
│ PREFILL │ compute-bound │ DECODE │ memory-bound
│ 2N/token │ │ re-reads │
└─────┬─────┘ │ ALL weights│
│ └─────┬─────┘
└──────────────┬──────────────────┘
▼
┌──────────────────────┐
│ LATENCY NAPKIN │ prefill + decode + scaffolding <= budget?
└──────────┬───────────┘
│ if no...
┌──────────────┼──────────────┐
▼ ▼ ▼
more chips smaller model better shapes
(expensive, (Flash!) (GQA, tiles,
finite) depth/width)
What you will be able to do
- Compute arithmetic intensity and place any operation on a roofline; say whether to attack bytes or FLOPs.
- Compute the ridge point of any accelerator from its spec sheet, and explain why H100's (~296 FLOP/byte) is higher than A100's (~154) despite being the faster chip.
- Produce an MFU budget — decompose wall-clock into matmul / vector / memory / comms / optimizer — and read off the biggest lever.
- Distinguish MFU from HFU and explain why HFU is always higher (~33% for free).
- Reproduce Feinberg's napkin math: Llama3-70B on v5e, ~5.8 s of prefill on one chip, a 4×4 station to get under the 0.5 s API limit.
- Show that at batch 1, decode costs ~3.8× prefill — and explain why that changes the provisioning answer.
- Quantify three co-design levers: tile quantization, GQA group size, and depth-vs-width.
- Prove that decode stays memory-bound even at batch 1024, and explain what follows from that.
The lab
| Lab | What you build |
|---|---|
| Lab 01 — Roofline, MFU Budget & the Latency Napkin | The full co-design toolkit: roofline analysis with ridge points and a "which lever" verdict, MFU/HFU accounting and budget decomposition, prefill/decode latency models, a chip-count solver for a latency budget, and the shape levers — tile efficiency, GQA saving, decode arithmetic intensity vs batch, depth-vs-width serial cost |
Success criteria. LAB_MODULE=solution pytest test_lab.py -v → 50 passed;
python solution.py reproduces the talk's ~5.8 s / 4×4 result and shows the chip count falling
as the model shrinks.
Deliverables checklist
- I can compute a ridge point and say what it means physically.
- I can produce an MFU budget and name the biggest lever from it.
- I always ask "MFU or HFU?" when shown a utilization number.
- I have reproduced the ~5.8 s single-chip prefill figure.
- I can explain why decode beats prefill at batch 1 and what that does to provisioning.
- I can state the GQA saving as a ratio and in concurrent requests.
- I can explain why decode never becomes compute-bound at realistic batch sizes.
Key takeaways
- Arithmetic intensity decides which resource you are fighting. It is a property of the algorithm, not the chip.
- Below the ridge, cut bytes; above it, cut FLOPs. Everything else is detail.
- MFU of 35% is an identity, not a grade. The breakdown is your agenda.
- HFU ≥ MFU, always. Activation checkpointing counted as useful work is ~33% for free.
- Prefill and decode are different machines. Never apply one's fix to the other.
- Decode is memory-bound at every realistic batch size — even 1024. That single fact explains batching, GQA, quantization and speculative decoding.
n_kv_headsis the biggest single serving lever, it costs almost no quality, and it is frozen at pre-training time.- Depth is serial and costs latency; width is parallel. As wide as quality allows, as shallow as quality tolerates.
- When the latency budget cannot be met, the answer is a smaller model, not more chips. That is why Flash and Flash-Lite exist, and why this is a pre-training concern.
References
- Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides — the "Small Model Customers" and "Inference-optimized Scaling" sections, and the Llama3-70B/v5e napkin math
- Williams, Waterman & Patterson, Roofline: An Insightful Visual Performance Model, CACM 2009 — the original model
- Pope et al., Efficiently Scaling Transformer Inference, 2022 — https://arxiv.org/abs/2211.05102 — the definitive treatment of prefill/decode arithmetic
- Austin et al., How To Scale Your Model — https://jax-ml.github.io/scaling-book/ — do the exercises
- 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
- Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022 — https://arxiv.org/abs/2205.05198 — where the MFU/HFU distinction is made precise
- Kwon et al., Efficient Memory Management for LLM Serving with PagedAttention (vLLM), 2023 — https://arxiv.org/abs/2309.06180
Warmup Guide — Roofline, MFU & Inference Co-Design
How to read this. No prior knowledge of performance modelling is assumed. Arithmetic intensity, the roofline, MFU, memory-bound versus compute-bound, and the prefill/decode split are all built from nothing. If you already profile kernels, start at Chapter 4 and do not skip Chapter 6 or Chapter 8.
Table of Contents
- Chapter 1: The Question This Phase Answers
- Chapter 2: A Chip Is Several Machines
- Chapter 3: Arithmetic Intensity and the Roofline
- Chapter 4: MFU — an Accounting Identity, Not a Grade
- Chapter 5: Prefill and Decode Are Two Different Machines
- Chapter 6: The Latency Napkin
- Chapter 7: Shape Co-Design, Lever by Lever
- Chapter 8: The Decode Wall
- Chapter 9: What Co-Design Actually Looks Like as a Job
- Lab Walkthrough
- Success Criteria
- Interview Q&A
- Tips & Takeaways
- References
Chapter 1: The Question This Phase Answers
Feinberg's talk poses a product scenario and then does the arithmetic live. Paraphrasing his setup: a web-interaction agent with a 128k context but only 8k incremental per turn, 128 decode tokens to emit an action, no more than a second of latency between actions, and 250 ms of that already consumed by "scaffolding, load balancing, request validation, kv cache retrieval" — which he flags as an optimistic allowance.
Then: Llama3-70B on v5e chips. What happens?
His answer, in one line from the slide: "Uh oh… 5.7 seconds for 1 chip." And the conclusion: to hit a 0.5 s API limit on prefill alone you already need a 4×4 station of v5e.
That calculation is the entire phase. It is how a product requirement becomes an architecture constraint, and it is why the pre-training team cares about serving at all. The job title says "pre-training," but the objective function has a serving term in it, and this chapter is that term.
Why this belongs to pre-training and not to the serving team
Because the levers that matter most are frozen when training starts:
n_kv_heads— decided before training, changes serving throughput by up to 64×- depth vs width at fixed
N— decided before training, changes decode latency directly - matrix dimensions — decided before training, changes how well the hardware tiles
- total parameter count — decided before training, and it is the dominant term in everything
A serving engineer inherits all of these. Only the pre-training lead can choose them. That is what "inference co-design" means, and why it is one of Feinberg's three verticals.
Chapter 2: A Chip Is Several Machines
The mental model that fixes most confusion
People treat an accelerator as one number ("an H100 does a petaflop"). It is not. It is several units with wildly different throughputs, and your model must use all of them:
┌──────────────────────────────────────────────────────────┐
│ ACCELERATOR │
│ │
│ ┌────────────────┐ VERY fast. This is the number │
│ │ MATMUL UNIT │ on the spec sheet. │
│ │ ~990 TFLOP/s │ (tensor cores / systolic array) │
│ └────────────────┘ │
│ │
│ ┌────────────────┐ 50-100x SLOWER than matmul. │
│ │ VECTOR UNIT │ gelu, softmax, norms, residuals. │
│ └────────────────┘ Counted in nobody's "peak". │
│ │
│ ┌────────────────┐ ~3.35 TB/s. Every activation │
│ │ HBM │ in and out crosses this. │
│ └────────────────┘ │
│ │
│ ┌────────────────┐ ~0.05-0.9 TB/s. Collectives. │
│ │ INTERCONNECT │ Slowest by far. │
│ └────────────────┘ │
└──────────────────────────────────────────────────────────┘
"Peak FLOP/s" describes the matmul unit alone. Everything else your model does is time during which that unit sits idle. Chapter 4 turns this observation into an accounting identity.
The number to internalize
Bandwidth is the scarce one, and it has been getting relatively scarcer for two decades. Compute has grown faster than memory bandwidth every generation — which is why more and more workloads are memory-bound over time, and why the roofline's ridge point keeps moving right.
for chip, (peak, hbm, bw, watts) in HARDWARE.items():
print(f"{chip:10s} {peak/1e12:6.0f} TFLOP/s {bw/1e12:5.2f} TB/s "
f"-> ridge {peak/bw:6.0f} FLOP/byte")
H100 990 TFLOP/s 3.35 TB/s -> ridge 296 FLOP/byte
A100-80 312 TFLOP/s 2.03 TB/s -> ridge 154 FLOP/byte
TPU v5e 197 TFLOP/s 0.82 TB/s -> ridge 241 FLOP/byte
TPU v5p 459 TFLOP/s 2.77 TB/s -> ridge 166 FLOP/byte
A counter-intuitive fact worth holding: the H100's ridge point (296) is higher than the A100's (154), even though the H100 is the better chip. Compute improved 3.2×; bandwidth only 1.7×. Newer hardware is harder to keep fed, so memory-bound workloads get relatively worse on it, not better. Anyone who says "just wait for better GPUs" has not looked at this column.
Chapter 3: Arithmetic Intensity and the Roofline
Arithmetic intensity
$$ I = \frac{\text{FLOPs performed}}{\text{bytes moved}} $$
"How much arithmetic do I get per byte I drag out of memory?" It is a property of the algorithm and its blocking, not of the chip.
def arithmetic_intensity(flops, bytes_moved):
if bytes_moved <= 0:
raise ValueError("bytes_moved must be positive")
return flops / bytes_moved
The ridge point
$$ I_{\text{ridge}} = \frac{\text{peak FLOP/s}}{\text{HBM bandwidth}} $$
The intensity at which the two resources are exactly balanced.
I < I_ridge→ memory-bound. The matmul unit is starved. Adding FLOPs is free; removing bytes is what helps.I > I_ridge→ compute-bound. Memory keeps up. Removing FLOPs is what helps.
The roofline itself
achievable
FLOP/s
peak ┤ ┌────────────────────────── compute-bound (flat roof)
│ ╱
│ ╱
│ ╱ slope = HBM bandwidth
│ ╱
│ ╱ memory-bound (sloped roof)
│ ╱
└───────┴──────────────────────────────────► arithmetic intensity
▲
ridge point (~296 FLOP/byte on H100)
def roofline_throughput(chip, intensity):
peak, _hbm, bw, _w = HARDWARE[chip]
return min(peak, bw * intensity) # the whole model, in one line
That single min() is the entire roofline model. Its power is not the formula; it is that it
tells you which lever to pull, which stops you optimizing the wrong thing for a week.
Three operations, placed
big matmul (8192^3) AI= 2730.7 compute 100.0% of peak -> reduce FLOPs
decode, batch=1 AI= 1.0 memory 0.3% of peak -> reduce bytes moved
layernorm (8k x 8k) AI= 1.2 memory 0.4% of peak -> reduce bytes moved
Read the middle row again. Batch-1 decode achieves 0.3% of an H100's peak — and that is not a bug, it is the arithmetic. You are reading 140 GB of weights to do 140 GFLOPs of work. No kernel engineer can fix that; only batching, quantization, or a smaller model can.
The third row is why kernel fusion exists: a layernorm does almost no arithmetic but has to read and write the whole activation tensor. Fusing it into the neighbouring matmul's epilogue removes the round-trip entirely.
The misconception
"Our MFU is low because our kernels are bad."
Maybe. But if your operation sits at intensity 1 on a chip whose ridge is 296, no kernel can be good. The roofline tells you the ceiling before you start optimizing, which is precisely its value: it distinguishes "we implemented this badly" from "this is what the algorithm costs."
Chapter 4: MFU — an Accounting Identity, Not a Grade
The definition
$$ \text{MFU} = \frac{\text{useful model FLOPs per second}}{\text{peak FLOP/s} \times \text{chips}} $$
Where "useful model FLOPs" is the 6ND arithmetic the model requires — not including
recomputation.
Why it is never near 100%
Feinberg addresses this directly, because people on social media see a number like 35% and conclude someone is incompetent. His point: to reach 100% you would need to be "doing a bunch of matmuls in a loop without reading any memory," and that is not a neural network. Real nets "have to apply activation functions or do attention or write intermediate outputs."
So build the budget:
def mfu_budget(matmul_s, vector_s, memory_s, comms_s, optimizer_s):
parts = {"matmul": matmul_s, "vector": vector_s, "memory": memory_s,
"comms": comms_s, "optimizer": optimizer_s}
total = sum(parts.values())
out = {f"{k}_fraction": v / total for k, v in parts.items()}
out["mfu"] = matmul_s / total
losses = {k: v for k, v in parts.items() if k != "matmul"}
out["biggest_lever"] = max(losses, key=losses.get)
return out
MFU (matmul busy) 35.7%
lost to vector 16.1%
lost to memory 21.4%
lost to comms 17.9%
lost to optimizer 8.9%
-> biggest lever: memory. That is your next week's work.
35.7% is not a failure. It is a sum that adds to 100%. And notice what you now have that you did not have from the single number: an ordered work queue.
| If the biggest loss is… | Do this |
|---|---|
| comms | overlap collectives with compute; change the sharding plan |
| memory | fuse kernels; use FlashAttention; reduce activation traffic |
| vector | fuse norms and activations into matmul epilogues |
| optimizer | shard optimizer state (ZeRO-1); use a cheaper optimizer; fuse the update |
Feinberg's framing of the goal connects straight back to co-design: the aim is "choosing shapes for this neural net that fully saturate all of those hardware units" — not maximizing one number.
MFU vs HFU — the trap
HFU (Hardware FLOPs Utilization) counts recomputation from activation checkpointing as
useful work. Full checkpointing takes hardware FLOPs per token from ~6N to ~8N:
def hfu(model_flops, seconds, chip, n_chips, recompute_factor=8/6):
return mfu(model_flops * recompute_factor, seconds, chip, n_chips)
MFU (honest) 11.8%
HFU (counts recompute as useful) 15.7%
Same run. Same hardware. ~33% apart, for free.
HFU is always ≥ MFU. When a blog post, a vendor benchmark, or a colleague quotes a utilization number, ask which one. A third of the difference between two teams' reported numbers is often just this definition.
Two more comparability traps. (1) MFU is only comparable within a fixed hardware/precision class — an FP8 number uses a different (larger) denominator than a bf16 one. (2) "Peak with sparsity" doubles the denominator again and assumes a 2:4 structured sparse model you almost certainly do not have.
Chapter 5: Prefill and Decode Are Two Different Machines
This is the most important structural fact in serving, and everything in Chapters 6–8 follows from it.
Prefill
You have the user's whole prompt. Process all T tokens at once. Every weight you read is used
for T tokens, so arithmetic intensity is high.
def prefill_seconds(n_tokens, n_params, chip, n_chips, mfu_frac=1.0):
"""2N FLOPs per token, forward only. COMPUTE bound."""
peak, _hbm, _bw, _w = HARDWARE[chip]
return (2.0 * n_params * n_tokens) / (n_chips * peak * mfu_frac)
Decode
You produce one token at a time. To produce it you must read every weight in the model,
and then do 2N FLOPs with them. Intensity ≈ 1.
def decode_seconds(n_tokens, n_params, chip, n_chips, bytes_per_param=2):
"""Every generated token re-reads EVERY weight. MEMORY-BANDWIDTH bound."""
_peak, _hbm, bw, _w = HARDWARE[chip]
return n_tokens * n_params * bytes_per_param / (n_chips * bw)
The table to memorize
| Prefill | Decode | |
|---|---|---|
| Work per step | whole prompt | one token |
| Bound by | compute | memory bandwidth |
| Arithmetic intensity | hundreds–thousands | ~1–50 |
| Helped by | more FLOP/s, pipelining, chunking | more bandwidth, batching, fewer bytes |
| Hurt by | long prompts (T² attention) | large models, large KV cache |
| Parallelism that works | pipelining (lots of tokens to chunk) | batching (nothing to pipeline) |
Why this is not a technicality
Because the optimizations are opposite:
- Pipelining works beautifully for prefill (thousands of tokens to chunk, so transfers hide
behind compute) and is actively harmful for decode — it adds
S−1serial hops to every single token, with no parallel work to overlap. - Batching is the primary decode lever and does comparatively little for prefill, which is already compute-saturated.
That asymmetry is exactly why prefill/decode disaggregation exists — running the two phases on separately provisioned hardware — and it is the reason the Flash 2.0 fix was pipelined prefill specifically, not pipelined inference in general.
Chapter 6: The Latency Napkin
Now reproduce the calculation from the talk, end to end.
The setup
- 8k incremental prefill per turn
- 128 decode tokens (enough to emit one agent action)
- 1.0 s total budget
- 250 ms of it already spent on scaffolding
- Llama3-70B (140 GB in bf16) on TPU v5e (197 TFLOP/s, 16 GB, 819 GB/s)
The first thing to check: does it even fit?
def weights_fit_chips(n_params, chip, bytes_per_param=2):
_peak, hbm, _bw, _w = HARDWARE[chip]
return math.ceil(n_params * bytes_per_param / hbm)
weights_fit_chips(70e9, "TPU v5e") → 9 chips, just to hold the model. Before a single
consideration of speed. Capacity and latency push in the same direction, which is a recurring
theme.
The sweep
1 chips: prefill 5.822s decode 21.880s total 27.952s MISS
4 chips: prefill 1.455s decode 5.470s total 7.176s MISS
16 chips: prefill 0.364s decode 1.368s total 1.981s MISS
64 chips: prefill 0.091s decode 0.342s total 0.683s OK
128 chips: prefill 0.045s decode 0.171s total 0.466s OK
The prefill column reproduces the talk exactly. One chip gives 5.8 s, matching his "5.7 seconds for 1 chip"; a 4×4 = 16-chip station brings prefill to 0.36 s, under the 0.5 s API limit he quotes.
The part most summaries get wrong
Look at the decode column. At batch 1, generating 128 tokens costs ~3.8× more than prefilling 8192 of them. 8192 tokens of prefill is 1.1 PFLOPs of compute; 128 tokens of decode is 17.9 TB of memory traffic. The chip is fast at the former and slow at the latter.
So while the prefill station needs 16 chips, the full turn needs ~64. The talk's 4×4 figure is about prefill specifically — and the honest full-budget answer is larger.
The caveat that keeps this honest: batch-1 decode is the worst case. Real serving batches many concurrent requests, which amortizes the weight read across all of them — that is precisely why continuous batching exists. But an interactive agent that must respond within a second cannot always wait to fill a batch. The tension between latency (small batch) and throughput (large batch) is the central trade-off of serving, and this table is where you feel it.
The conclusion that pays for the whole phase
N= 70.0B -> 64 chips for a 1.0s turn (9 just to hold weights)
N= 35.0B -> 32 chips for a 1.0s turn (5 just to hold weights)
N= 17.0B -> 16 chips for a 1.0s turn (3 just to hold weights)
N= 8.0B -> 8 chips for a 1.0s turn (1 just to hold weights)
Halving N halves both columns. Prefill FLOPs are linear in N; decode bytes are linear
in N. There is no shape trick that beats simply having a smaller model.
That is the economic case for Flash and Flash-Lite, stated as arithmetic — and it is why a pre-training lead owns it. The serving team cannot make the model smaller. Only the recipe can.
Chapter 7: Shape Co-Design, Lever by Lever
Three levers, all frozen at pre-training time.
Lever 1 — tile quantization
Matmul units work on fixed tiles (128×128 on a TPU systolic array; multiples of 8/16/64 for tensor cores). A dimension that is not a multiple gets padded, and you pay for the padding.
d_ff= 11000 -> 99.9% of the padded matmul is real work
d_ff= 11008 -> 100.0%
d_ff= 8192 -> 100.0%
d_ff= 4097 -> 97.0%
d_ff= 4096 -> 100.0%
Individually small. But it applies to d_model, d_ff, head dimension, expert count, and
every shard boundary, and it compounds multiplicatively across dozens of layers. This is
why production models use dimensions like 4096, 8192, 11008 rather than round decimal numbers.
The worst case is one element past a boundary:
tile_efficiency(129, tile=128)= 50.4%. You pay for two tiles to use one and a bit.
Lever 2 — KV heads (the big one)
64 kv heads (group 1): 687.2 GB of KV cache 1.0x (MHA)
8 kv heads (group 8): 85.9 GB of KV cache 8.0x smaller
1 kv heads (group 64): 10.7 GB of KV cache 64.0x smaller
Grouped-Query Attention shares one KV head across a group of query heads. Quality barely moves; the cache shrinks by exactly the group factor.
Translate that into what a serving team actually cares about — concurrent requests on 8×H100 serving a 70B model at 8k context: 22 with MHA, 182 with GQA-8. An 8× throughput difference from one integer in a config file, chosen months before the model exists.
This is the single clearest example of inference co-design there is. It is also irreversible: the KV projections have different shapes, so you cannot convert afterwards without retraining.
Lever 3 — depth vs width
At fixed N you can be deep and narrow or shallow and wide.
- Depth is serial. Layer
k+1cannot start until layerkfinishes. Depth costs decode latency directly, and adds pipeline stages that must synchronize. - Width is parallel and produces larger, more efficient matmuls.
Deeper models are often slightly better per parameter, so there is a real trade. The co-design heuristic: as wide as quality allows, as shallow as quality tolerates.
Chapter 8: The Decode Wall
The most consequential single fact in LLM serving.
def decode_batch_intensity(n_params, n_layers, n_kv_heads, d_head, seq_len, batch):
"""Weights are read ONCE and shared across the batch; KV is per request."""
flops = 2.0 * n_params * batch
kv = kv_cache_bytes(n_layers, n_kv_heads, d_head, seq_len, batch)
return flops / (n_params * 2 + kv)
H100 ridge point: 296 FLOP/byte
batch= 1: AI= 1.0 FLOP/byte MEMORY bound ( 0.3% of the way to the ridge)
batch= 8: AI= 6.9 FLOP/byte MEMORY bound ( 2.3%)
batch= 64: AI= 28.7 FLOP/byte MEMORY bound ( 9.7%)
batch= 256: AI= 43.3 FLOP/byte MEMORY bound (14.7%)
batch= 1024: AI= 49.6 FLOP/byte MEMORY bound (16.8%)
Even at batch 1024, decode reaches only 17% of the way to the ridge. Decode does not become compute-bound at any batch size you would actually run.
Why batching stops helping
Look at the shape of that curve — 1 → 8 gives 6.9×, but 256 → 1024 gives only 1.15×. The reason is in the denominator:
bytes = (weights, FIXED) + (KV cache, GROWS WITH BATCH)
At small batch, the fixed weight term dominates and batching amortizes it beautifully. At large batch, the KV term dominates and grows with the batch, so the ratio saturates. The KV cache is what limits batching, which is why PagedAttention, prefix sharing, and GQA are all ultimately about the same thing.
Everything that follows from this one fact
| Technique | What it really does |
|---|---|
| Continuous batching | raise intensity by amortizing the weight read |
| GQA / MQA | shrink the KV term so batching keeps working longer |
| Quantization | halve or quarter the bytes; the dominant term |
| Speculative decoding | verify k tokens per weight-read pass — raises intensity directly |
| PagedAttention | fit more requests in the same HBM, so batches can be bigger |
| Smaller models | attack the dominant term head-on |
Every one of those is an attack on bytes, not FLOPs. That is what "memory-bound" means in practice, and why Feinberg's third vertical (quantization) is an energy and bandwidth lever before it is a memory-capacity one.
Chapter 9: What Co-Design Actually Looks Like as a Job
Pulling it together into the workflow a pre-training lead runs.
1. Get the product's latency budget and expected traffic shape.
"sub-second turns, 8k incremental context, 128-token actions"
2. Compute the napkin for candidate model sizes.
-> a table of (N, chips needed, chips to hold weights)
3. Find the largest N that meets the budget at acceptable cost.
-> this is now a CONSTRAINT on the pre-training recipe
4. Feed that back into the scaling law (Phase 01).
"we can afford 17B active parameters; what is the best 17B we can train
with our compute budget, and how much loss do we give up vs Chinchilla?"
5. Choose shapes inside that budget.
n_kv_heads (biggest lever), tile-aligned dimensions, depth/width,
dense vs MoE (Phase 02: active params drive FLOPs, total drives memory)
6. Verify with a roofline + MFU budget on the real hardware, then iterate.
Step 4 is the seam where this phase meets Phase 01, and it is the actual job. The scaling
law says "for this compute budget, N = 92B minimizes loss." The napkin says "92B needs 3
chips just to hold and blows the latency budget." So you deliberately undershoot Chinchilla,
train a smaller model on more tokens, and accept slightly worse loss for dramatically cheaper
serving.
That deliberate deviation is inference-aware scaling, and it is why Feinberg's slide notes that Chinchilla-style scaling "ignores inference cost."
Lab Walkthrough
Lab 01 — Roofline, MFU Budget & the Latency Napkin
Implement in this order:
arithmetic_intensity,ridge_point,roofline_throughput,roofline_report. The roofline is onemin(); the value is theleververdict it returns.mfu,hfu,mfu_budget.test_hfu_is_always_at_least_mfuis the one to internalize.prefill_seconds,decode_seconds,interactive_latency,chips_for_latency_budget,weights_fit_chips.tile_efficiency,kv_cache_bytes,gqa_saving,decode_batch_intensity,depth_vs_width.
The money tests:
test_prefill_reproduces_the_talks_number— ~5.8 s on one v5e chip.test_a_4x4_station_brings_prefill_under_the_half_second_limit.test_a_smaller_model_needs_fewer_chips— the whole economic argument, in one assertion.test_decode_stays_memory_bound_even_at_huge_batch.
The traps:
- Prefill is
2Nper token (forward only), not6N. That is training. - Decode time is driven by bytes, not FLOPs. If your
decode_secondshaspeakin it, it is wrong. chips_for_latency_budgetmust returnNonerather than loop forever — "no amount of hardware fixes this" is a real and important answer.- A budget below the scaffolding overhead is impossible; raise rather than return something.
Success Criteria
-
LAB_MODULE=solution pytest test_lab.py -v→ 50 passed. -
Your
lab.pyreaches 50 passed. -
python solution.pyruns and you can explain all nine sections. - You can compute a ridge point and explain why H100's is higher than A100's.
- You can produce an MFU budget and name the biggest lever from it.
- You reflexively ask "MFU or HFU?" when shown a utilization figure.
- You can reproduce the ~5.8 s single-chip prefill number and the 4×4 conclusion.
- You can explain why decode beats prefill at batch 1, and the batching caveat.
- You can state the GQA saving as a ratio and as concurrent requests.
Interview Q&A
Q: What is the roofline model and what is it for?
Plot achievable throughput against arithmetic intensity — FLOPs per byte moved. Below the ridge
point (peak FLOP/s ÷ bandwidth) you are memory-bound and throughput is bandwidth × intensity;
above it you are compute-bound and capped at peak. Its value is that it tells you which lever
to pull before you spend a week optimizing: below the ridge, reduce bytes; above it, reduce
FLOPs. It also gives you a ceiling, which distinguishes "we implemented this badly" from "this
is what the algorithm costs."
Q: Our training run is at 38% MFU. Is that bad? No — that is the normal band. MFU is the fraction of peak matmul throughput achieved, and a transformer is not pure matmul: it runs vector ops (norms, activations, softmax), moves activations to and from HBM, runs collectives, and executes the optimizer step. Each of those is time the matmul unit is idle. The useful move is to decompose it — matmul / vector / memory / comms / optimizer — because the breakdown is an optimization agenda. And I would check whether that 38% is MFU or HFU: HFU counts activation recomputation as useful work and is always higher, typically by about a third.
Q: Why is decoding so much slower than prefill per token?
Different bottlenecks. Prefill processes the whole prompt at once, so each weight read is
amortized across thousands of tokens — arithmetic intensity is in the hundreds and it is
compute-bound. Decode produces one token at a time, so it reads every weight in the model to
do 2N FLOPs — intensity around 1, memory-bandwidth-bound. On a 70B model that is 140 GB of
traffic per token. It is not a kernel problem; it is arithmetic.
Q: How would you make decode faster? Attack bytes, not FLOPs. Batching, to amortize the weight read across concurrent requests — the single biggest lever. Quantization, to halve or quarter the dominant term. GQA or MQA, to shrink the KV cache, which is what limits how large a batch you can fit. Speculative decoding, which verifies several tokens per weight-read pass. And ultimately a smaller model. Notably, batching saturates: past a few hundred, the KV term grows with the batch and the intensity curve flattens.
Q: You need sub-second agent turns from a 70B model. Walk me through it.
Napkin first. 8k incremental prefill at 2N per token is 1.1 PFLOPs; on a v5e at 197 TFLOP/s
that is ~5.8 s on one chip, so 16 chips gets prefill under half a second. But 128 decode tokens
at batch 1 is 17.9 TB of memory traffic — about 3.8× the prefill cost — so the full turn needs
around 64 chips. Also, 140 GB of weights on 16 GB chips means 9 chips just to hold the model.
Then the real conclusion: 64 chips per concurrent conversation is not a viable product, so the
answer is not more hardware, it is a smaller model. Halving N halves both prefill and decode.
That is the economic argument for a Flash-class model, and it is a pre-training decision, not a
serving one.
Q: What is inference co-design and why does the pre-training team own it?
Choosing architecture shapes with the serving target in mind — matrix dimensions that tile
cleanly onto the hardware, KV-head sharing, depth versus width, dense versus MoE. The
pre-training team owns it because these are all frozen the moment training starts. A serving
engineer inherits n_kv_heads; they cannot change it without retraining. And it is a huge
lever: 64 KV heads versus 8 is an 8× difference in concurrent requests at essentially no
quality cost.
Q: Why is the H100's ridge point higher than the A100's?
Because compute scaled faster than bandwidth: 3.2× versus 1.7× between those generations. The
ridge is peak ÷ bandwidth, so it moved right. The practical consequence is that memory-bound
workloads — which is all of decode — get relatively worse on newer hardware, not better. This
is a long-running trend and it is why bandwidth-saving techniques keep gaining importance.
Tips & Takeaways
Tips
- Compute arithmetic intensity before optimizing anything. It tells you the ceiling and the lever in one number.
- Always ask "MFU or HFU?" and "at what precision?" Both change the denominator.
- Never quote a utilization number without the breakdown. The breakdown is the useful part.
- Do the napkin before choosing a model, not after. It takes two minutes and routinely changes the answer.
- Check "does it even fit?" first. Capacity often binds before latency does.
- Treat prefill and decode as separate systems with separate budgets and separate hardware provisioning.
- Write down
n_kv_headsin the design doc with its serving justification. It is the highest-leverage irreversible number in the config. - Prefer tile-aligned dimensions. Free performance; costs nothing but attention.
Takeaways
- A chip is several machines; "peak" describes only the matmul unit.
- Arithmetic intensity decides which resource you are fighting. Below the ridge, cut bytes.
- Newer hardware has a higher ridge point — memory-bound work gets relatively worse.
- MFU of 35% is an accounting identity. The decomposition is your work queue.
- HFU ≥ MFU always. Ask which.
- Prefill is compute-bound; decode is memory-bound. Two machines, opposite optimizations.
- Pipelining helps prefill and hurts decode. This is why disaggregation exists.
- Decode never reaches the ridge at any realistic batch size — even 1024.
- The KV cache is what makes batching saturate, which is why GQA matters so much.
- When latency cannot be met, the answer is a smaller model. That is a pre-training decision.
References
- Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides — "Small Model Customers", "Why Do Real-time Use Cases Imply Smaller Models", and the v5e napkin math
- Developing Dev interview · video — the MFU discussion and the inference co-design vertical
- Williams, Waterman & Patterson, Roofline: An Insightful Visual Performance Model for Multicore Architectures, CACM 2009
- Pope et al., Efficiently Scaling Transformer Inference, 2022 — https://arxiv.org/abs/2211.05102
- Austin et al., How To Scale Your Model — https://jax-ml.github.io/scaling-book/
- Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022 — https://arxiv.org/abs/2205.05198 (MFU vs HFU)
- Shazeer, Fast Transformer Decoding (MQA), 2019 — https://arxiv.org/abs/1911.02150
- Ainslie et al., GQA, 2023 — https://arxiv.org/abs/2305.13245
- Dao et al., FlashAttention, 2022 — https://arxiv.org/abs/2205.14135 (the canonical memory-traffic win)
- Kwon et al., PagedAttention / vLLM, 2023 — https://arxiv.org/abs/2309.06180
- Leviathan et al., Fast Inference from Transformers via Speculative Decoding, 2022 — https://arxiv.org/abs/2211.17192
- Jouppi et al., In-Datacenter Performance Analysis of a Tensor Processing Unit, 2017 — https://arxiv.org/abs/1704.04760
Lab 01 — Roofline, MFU Budget & the Latency Napkin
Build the arithmetic that turns "we want a real-time agent" into "the model must be smaller" — and reproduce the napkin math from Feinberg's Princeton talk.
The problem
A product manager says "sub-second responses." A serving engineer says "we'll add GPUs." Both are guessing. The roofline and the latency napkin turn that conversation into arithmetic, and the arithmetic usually says something neither of them expected: no amount of hardware fixes this; the model has to be smaller.
That conclusion is why a pre-training lead owns inference co-design. The levers that matter —
n_kv_heads, depth vs width, tile alignment, parameter count — are all frozen the moment
training starts.
What you build
| Group | Functions | The idea |
|---|---|---|
| Roofline | arithmetic_intensity, ridge_point, roofline_throughput, roofline_report | FLOPs per byte; below the ridge cut bytes, above it cut FLOPs |
| Utilization | mfu, hfu, mfu_budget | 35% is an identity, not a grade; the breakdown is an agenda |
| Latency | prefill_seconds, decode_seconds, interactive_latency, chips_for_latency_budget, weights_fit_chips | The napkin, end to end |
| Co-design | tile_efficiency, kv_cache_bytes, gqa_saving, decode_batch_intensity, depth_vs_width | The three irreversible levers |
Key concepts
| Concept | Why it is in this lab |
|---|---|
| Arithmetic intensity | Property of the algorithm; decides which resource you fight |
| Ridge point | peak ÷ bandwidth. H100 ≈ 296 — and higher than A100's 154 |
| MFU vs HFU | HFU counts recomputation as useful. Always ≥ MFU, ~33% for free |
| MFU budget | matmul / vector / memory / comms / optimizer → your work queue |
Prefill 2N/token | Compute-bound. 6N is training; using it here is a 3× error |
| Decode = bytes | Every token re-reads every weight. 140 GB per token at 70B |
| The 4×4 station | 16 v5e chips to get 8k prefill under 0.5 s — the talk's conclusion |
| GQA ratio = group size | 64→8 KV heads is exactly 8× less cache, ~0 quality cost |
| The decode wall | Even batch 1024 reaches only 17% of the ridge |
Files
| File | What it is |
|---|---|
| lab.py | Your implementation. Signatures, docstrings and validation contracts given. |
| solution.py | Reference. python solution.py runs a nine-part worked example. |
| test_lab.py | 50 tests: happy path, validation, boundaries, invariants, determinism. |
| requirements.txt | pytest only. Pure stdlib otherwise. |
Run
pytest test_lab.py -v # your lab.py — red until you implement
LAB_MODULE=solution pytest test_lab.py -v # the reference — must be green (50 passed)
python solution.py # the worked example
Where to start
arithmetic_intensity→ridge_point→roofline_throughput. The roofline is onemin(). Gettest_roofline_is_continuous_at_the_ridgegreen — the two roof segments must meet exactly.mfu_budget, thentest_perfect_mfu_when_only_matmul_runs. That boundary case is Feinberg's point: 100% requires a pure matmul loop, which is not a neural network.prefill_seconds/decode_seconds. Thentest_prefill_reproduces_the_talks_number— if you get ~5.8 s you have it right.chips_for_latency_budget. Note it must returnNonewhen the budget is unreachable.- The co-design levers.
The traps:
- Prefill is
2Nper token, forward only.6Nis training — a 3× error. decode_secondsmust not containpeak. Decode is driven by bytes and bandwidth. If peak FLOP/s appears in it, the model is wrong.chips_for_latency_budgetmust terminate and returnNonerather than loop. "More hardware does not fix this" is a real answer and the function has to be able to give it.- A budget below the scaffolding overhead is impossible — raise, don't return a number.
gqa_savingmust rejectn_kv_heads > n_query_headsand non-divisible configurations.
Success criteria
-
LAB_MODULE=solution pytest test_lab.py -v→ 50 passed. -
Your
lab.pyreaches 50 passed. -
python solution.pyruns and you can explain all nine sections. - You reproduce ~5.8 s single-chip prefill and the 4×4 station conclusion.
- You can explain why decode beats prefill at batch 1, and the batching caveat.
- You can state why H100's ridge point is higher than A100's, and why that matters.
How this maps to the real stack
| This lab | The real thing | Where the miniature lies |
|---|---|---|
roofline_report | NVIDIA Nsight Compute's roofline; Intel Advisor | Real tools measure achieved bytes and FLOPs from hardware counters. Ours computes the analytical intensity, which is the ceiling, not the achieved value. |
mfu / mfu_budget | Megatron-LM and MaxText log MFU per step; profiler timelines give the breakdown | Real breakdowns come from kernel traces, where ops overlap. Ours assumes serial phases, so it over-attributes time. Directionally right, and the right mental model. |
prefill_seconds / decode_seconds | vLLM / TensorRT-LLM benchmarks; llm-analysis | Ours ignores attention's T² term (fine at 8k, wrong at 128k), kernel launch overhead, scheduling, and the fact that real bandwidth utilization is 60–90% of peak. Use it to size, not to promise SLAs. |
gqa_saving | The num_key_value_heads field in any HF config | Exact. This one is not an approximation. |
tile_efficiency | XLA / cuBLAS padding behaviour | Real libraries pick among several tile sizes and may pad differently. The effect is real; the exact constant is not. |
HARDWARE table | Vendor spec sheets | Peak numbers are marketing maxima at ideal clocks. Sustained throughput under thermal load is 5–15% lower. |
What is not a lie: the ridge-point definition, 2N per token for prefill, the KV-cache
formula, the GQA ratio, and the fact that decode is bandwidth-bound. Those are exact, and those
are what get asked about.
Extensions
- Add the attention
T²term. At 128k context the sequence-dependent attention matmuls dominate prefill (Phase 00, Break 1). Extendprefill_secondsand watch the napkin's conclusions change completely for long-context products. - Model continuous batching properly. Add a queue with arrival rates and compute p50/p99 latency versus throughput. That curve — the latency/throughput frontier — is what serving teams actually optimize, and this lab only shows its two endpoints.
- Add speculative decoding. Model an acceptance rate
αand a draft model of sizeN_d; compute the effective intensity gain and find where it stops paying. - Validate against real hardware. Run a 7B model in vLLM, measure tokens/sec at several
batch sizes, and compare to
decode_batch_intensity. The gap between your model and reality is the lesson. - Build the co-design search. Given a latency budget and a compute budget, search over
(N, n_layers, d_model, n_kv_heads)for the configuration that maximizes predicted quality (using Phase 01's scaling law) subject to meeting the napkin. That is the actual job, and it is a genuinely good portfolio piece.
Interview / resume bullets
- "Built a roofline and MFU-accounting toolkit for LLM training and serving — arithmetic intensity versus ridge point with an explicit optimize-bytes-or-FLOPs verdict, MFU/HFU disambiguation, and wall-clock decomposition across matmul, vector, memory, collective and optimizer time to rank optimization work."
- "Reproduced Google DeepMind's published inference-scaling analysis from first principles: ~5.8 s single-chip prefill for a 70B model on TPU v5e, a 4×4 prefill station to meet a 0.5 s API limit, and the batch-1 result that decode costs ~3.8× prefill — concluding quantitatively that meeting an interactive latency budget requires a smaller model rather than more accelerators."
- "Quantified the irreversible inference co-design levers set at pre-training time: GQA group size (8× concurrent-request throughput at 8 KV heads versus 64), tile-aligned matrix dimensions, and depth-versus-width serial latency cost."
- Interview-ready: "Decode never becomes compute-bound at any realistic batch size — even 1024 reaches 17% of the ridge — because the KV cache grows with the batch while the weight read does not. That one fact explains batching, GQA, quantization and speculative decoding."