Frontier Pre-Training Lead

← All Roles

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

DocumentWhat it gives you
The RoleWho 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, DissectedEvery 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 StandardWhat every lab in this track guarantees
GlossaryEvery term, one line each
Cheat SheetThe 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

#PhaseWhat you buildMaps to
00FLOPs & Memory AlgebraA 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
01Scaling Laws I — Kaplan → ChinchillaIsoFLOPs ladder, parabola/power-law/parametric fitting, bootstrap forecasts, optimal ladder design, the baseline-vs-candidate crossoverThe six-step IsoFLOPs slide; "UNDERTRAINED!" · 64 tests
02Mixture of Experts From ScratchRouter, 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
03Roofline, MFU & Inference Co-DesignRoofline 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.

PhaseWhat it will buildCovered today in
Scaling Laws II — beyond L(N, D)L(N, U, R) data-constrained optimizer, inference-aware lifetime cost, routed & distillation lawsClaims 15–17
Sharding & ParallelismCollective cost models, DP/FSDP/TP/PP/EP planner, the 4×4 mesh questionClaim 12
Pipeline Prefill & ServingAll-to-all vs pipelined-prefill simulator, bubble analysis, prefill/decode disaggregationClaim 12
Distillation at ScaleKL/temperature losses, top-k teacher-logit store, teacher-compute allocationClaims 3, 17
Quantization & Power EconomicsAffine/group/vector quantizers, outlier handling, the energy & TCO modelClaim 5
Training Stability & SREGoodput model, spike detector, checkpoint-policy optimizerClaim 13
Post-Training: SFT, RLHF, PPOSFT masking, Bradley-Terry reward model, PPO step with KL controlClaim 1
Kernels & Tile DSLsA tiny tile-DSL interpreter, online softmax, memory-traffic accountingClaims 1, 19
Research as an MDPPortfolio planner, value of information, calibration scoringClaim 2
Capstone — Mini Frontier LabBudget → recipe → forecast → compress → serve → the go/no-go memothe whole track

What every phase contains

FileVoiceWhat it gives you
README.mdthe syllabuswhy the phase exists, the concept map, the lab spec, deliverables, key takeaways
WARMUP.mdthe professorzero 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 workbencha 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.

WeeksWorkMilestone
1Read jd.md + TRANSCRIPT-DISSECTED.mdYou can explain MoE, MFU, 6ND, distillation and the Flash 2.0 fix to a colleague
2Phase 00You can size any model on any cluster on a napkin
3–4Phase 01You can fit a scaling law and defend a flagship forecast with error bars
5–6Phase 02You can implement an MoE layer and debug a collapsing router
7–8Phase 03You can do the latency napkin and defend an architecture shape
9+The four artifacts belowThe 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):

  1. Handwritten How To Scale Your Model exercises, on video.
  2. A transformer from scratch — with a FLOP counter and a memory profile, not just a forward pass.
  3. A real, benchmarked contribution to vLLM / SGLang / TensorRT.
  4. 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.



Key takeaways for the whole track

  1. You get one shot. Everything upstream of the flagship run exists to make one irreversible decision defensible.
  2. A scaling law describes your recipe, not the universe. Change the recipe, refit the law.
  3. Loss forecasting is recipe selection. Compare fitted curves at the target FLOP count.
  4. Optimize lifetime cost, not training cost — while knowing the inference term is unforecastable.
  5. Nothing is free; costs move. MoE trades compute for memory, communication, instability, and data hunger.
  6. Prefill and decode are different machines. Compute-bound vs bandwidth-bound.
  7. Power is the bill. Quantization is an energy lever before it is a memory lever.
  8. Distillation is variance reduction, and its hard part is storage.
  9. Goodput is a research deliverable. 78% vs 94% on the same hardware is days of a 40-day run.
  10. 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

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.

WhenWhereWhat he actually did
→2017Princeton, COS + SML (CS + Statistics & Machine Learning)3D CNN for MRI segmentation, for connectome reconstruction
2017–2018UC Berkeley RISE Lab, ML Systems PhD (with Ion Stoica, Joey Gonzalez, Mike Franklin)Model-based deep RL on MuJoCo (with Sergey Levine)
2018–2019Dropped outLeft the PhD for a startup
2019–2021Sisu Data (with Peter Bailis), Head of MLEfficient database cubing with FDR control, via a custom distributed lasso engine
2021–2022Google CerebraQuantizing the Ads DNN for pCTR serving efficiency
2022→Google Brain → Google DeepMindOptimizer work (with Elad Hazan) → inference-efficient LLMsFlash pre-training lead → Distinguished Engineer, area lead for Gemini pre-training

Four things to extract from that table

  1. 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.
  2. 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.
  3. 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.
  4. 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.

VerticalThe one-line problemWhy it is hardPhase
DistillationPush "statistics about the underlying dataset" from a teacher into a studentThe 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-designPick "shapes of the matrices", "attention shapes, num heads" so the net "fully saturates" every hardware unitQuality and MFU pull in opposite directions; the right answer differs per chip generation.Phase 05
QuantizationShrink 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 ≈ 6ND and 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.

#SkillPhase
1Derive transformer FLOPs and memory from shapes; 6ND; KV-cache sizing00
2Fit and extrapolate a scaling law; IsoFLOPs; know why Kaplan and Chinchilla disagreed01
3Data-constrained (L(N,U,R)), inference-aware, routed and distillation scaling laws02
4Implement an MoE layer: router, top-k, auxiliary losses, capacity factor, drop/pad03
5DP / FSDP / TP / PP / EP; collective cost models; GSPMD-style sharding annotations04
6Roofline & MFU; arithmetic intensity; prefill vs decode; the latency napkin math05
7Pipelined prefill for MoE; disaggregated serving; hiding comms behind compute06
8Distillation objectives; teacher-logit storage at scale; capacity gap; variance reduction07
9Quantization from FP32 → 4-bit; activation quant; power/TCO arithmetic08
10Training stability: loss spikes, goodput, checkpointing, iterator bugs, SDC09
11SFT and RLHF/PPO mechanics well enough to debug them on old hardware10
12Kernel/tile-DSL thinking: why ThunderKittens-style abstractions exist11
13Research as an MDP: value-of-information, kill criteria, portfolio planning12
14Put it together: budget → recipe → forecast → run → compress → serve → memo13

8. How He Says To Get Hired

He is unusually concrete about this, so it is worth listing exactly.

  1. 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."
  2. 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.
  3. 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."
  4. Be able to traverse a citation tree. Knowing which papers are the high-value ones for a topic is itself the skill.
  5. 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.)
  6. 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:

  1. Quote it (briefly) or state it plainly.
  2. Decode it for someone with zero background — every term defined from nothing.
  3. Go under the hood — the actual mechanism, with math and runnable code.
  4. 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

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:

  1. 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.
  2. 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.
  3. 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 p values 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:

StrategyStorage for 10T tokensCostFidelity
Full logits~5 EBimpossibleperfect
Top-k logits (k=64) + renormalized tail mass~10 TBcheapvery good
Online distillation (teacher runs alongside student)0teacher FLOPs every step, foreverperfect
Sampled sequences (teacher generates text, student trains on it)~20 TB textone-timeweakest

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 formax_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):

OperationEnergyRelative
8-bit integer add0.03 pJ
32-bit float add0.9 pJ30×
32-bit float multiply3.7 pJ123×
Read 32 bits from on-chip SRAM5 pJ167×
Read 32 bits from off-chip DRAM640 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 1e24three 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 * D is 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

  1. Attention is excluded. The QKᵀ and attn·V matmuls cost roughly 12 · L · n_ctx · d_model FLOPs per token, which does not scale with N. 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), where B=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.
  2. MoE. N must 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?"
  3. Embeddings. The input embedding is a lookup (~free); the output unembedding is a real matmul (2 · d_model · vocab per token). At small N with a 256k vocab this is a large fraction — a common source of wrong small-model FLOP counts.
  4. Activation checkpointing adds a partial extra forward pass, pushing you toward ~8ND of hardware FLOPs while the model FLOPs stay 6ND. (This is the MFU/HFU gap from Claim 6.)

Takeaway. C = 6ND is 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 N gives 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:

MethodWhat it minimizesBias
Least squares on LΣ (L_pred − L_obs)²dominated by large-loss (small-model) points
Least squares on log Lrelative errortreats all scales equally
Huber loss on log Lrelative error, outlier-robustwhat Chinchilla actually used
MLE with an explicit noise modellikelihood under stated assumptionsrequires 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−1 serial 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:

FailureWhat it looks likeWhy it is brutal
Loss spikeloss jumps 2→8 in one stepMay recover, may permanently damage the model. You have minutes to decide whether to roll back.
Hardware failureone chip of 10,000 diesSynchronous 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 erroringThe worst one. No crash, no alert; the model just gets subtly worse. Detected only by cross-replica checksums.
Data iterator bugwrong shard, an off-by-one, a repeated segmentThis is the one he names. Silent. You can burn days of compute training on the wrong data.
Stragglersone slow nodeEveryone waits at the all-reduce barrier; throughput collapses to the slowest node.
Checkpoint corruptionthe save itself failsYour 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

  1. 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.
  2. 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.
  3. The two phases want different hardware allocations — which is the entire argument for prefill/decode disaggregation, and the setup for Claim 12.
  4. Therefore: make the model smaller. Halving N halves both columns linearly. This is the entire economic case for Flash and Flash-Lite, and it is why his job exists.
  5. 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_inf and 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:

"D was 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 corpus
  • R = 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:

FlavorWhat it meansHow you do it
Adding pointsfill in the quality-vs-size curve at a new sizetrain another model at a new N
Increasing the slopemake the whole curve better — more quality per parameterbetter 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:

  1. C = 6ND. Derive it. It converts money into models.
  2. 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.
  3. A scaling law is a property of your recipe, not of nature. Change the recipe, refit.
  4. Loss forecasting is recipe selection. Compare fitted laws at the target FLOP count, not runs at small scale.
  5. Chinchilla beat Kaplan by fixing an experimental-design bug, not by having a better idea. Methodology is the frontier.
  6. Chinchilla-optimal is the wrong objective if you serve at scale. Optimize lifetime cost — but know that D_inf is unforecastable (Jevons, market expansion).
  7. MoE buys capacity with memory, communication, instability, and data hunger. Nothing is free; it just moves.
  8. The Flash 2.0 unlock was changing the sharding axis — layers instead of experts, at prefill — so communication hides behind compute.
  9. Prefill and decode are different machines. Compute-bound vs bandwidth-bound. Never apply one's optimization to the other.
  10. MFU of 35% is an accounting identity, not a failure. Read the breakdown; it is your agenda.
  11. 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.
  12. Distillation is variance reduction; a better teacher reduces bias. The teacher's shape is the signal, not its argmax.
  13. Distillation's hard part is storage, not the loss. 10T tokens × 256k vocab is 5 EB. Top-k or online.
  14. Data is where the pre-training lead's time goes — half his year. Not architecture.
  15. D is not iid. L(N, U, R): unique tokens and repeats are different axes. ~4 epochs is nearly free; 32 is mostly waste.
  16. Synthetic data helps in the Stein's-paradox sense — biased but lower variance.
  17. 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.
  18. Research is an MDP. Buy information before outcomes. Write kill criteria in advance. Diversify across premises, not implementations.
  19. You can do real pre-training research on one GPU: kernels, vector quantization, test-time-compute trade-offs, and the statistics of scaling laws.
  20. 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

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 algebra6ND, 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)

FileContract
README.mdthe 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.pylearner implementation with focused # TODO markers; signatures and docstrings already in place; never a blank file
solution.pycomplete reference; python solution.py runs a worked example and prints output; deterministic
test_lab.pypositive, negative, boundary, numerical-stability and determinism tests; runnable against either module via LAB_MODULE
requirements.txtusually 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.approx or 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; 2N forward and 4N backward sum to 6N.
  • 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 × forward FLOPs; dequant(quant(x)) is within the group's error bound; pipeline utilization is M/(M+S−1); goodput ≤ 1
  • deterministic output — same seed → same bytes

The two teaching docs (per phase)

DocumentVoiceWhat it gives you
README.mdthe syllabuswhy the phase exists, concept map, lab spec, deliverables checklist, key takeaways
WARMUP.mdthe professorzero-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.md opens 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 pointpeak 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 lossL = 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 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

QuantityFormulaNotes
Matmul cost2·m·k·n FLOPsthe 2 is multiply + add
Linear layer, per token2 × params FLOPsforward only
Training computeC = 6ND2N forward + 4N backward
Inference compute2N per tokenprefill 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·Lthe part 6ND omits
Unembedding6·d_model·V per tokenlarge fraction at small N
Chinchilla splitN = √(C / 6r), D = rNr ≈ 20 tokens/param
Lifetime cost6N·D_train + 2N·D_infthe objective Chinchilla ignores
Scaling lawL = E + A/N^α + B/D^βirreducible + capacity + data
Compute-optimal NN_opt ∝ C^(β/(α+β))= 0.5 exactly when α = β
Consistency checka + b ≈ 1falls out of C = 6ND — free bug detector

Memory

TermBytes per parameterNotes
bf16 weights2
bf16 gradients2
fp32 master + Adam m + v12
Mixed-precision Adam total16before a single activation
Adafactor4 (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

Chipbf16 peakHBMBandwidthWatts
H100 SXM990 TFLOP/s80 GB3.35 TB/s700
A100-80312 TFLOP/s80 GB2.03 TB/s400
TPU v5e197 TFLOP/s16 GB819 GB/s170
TPU v5p459 TFLOP/s95 GB2.77 TB/s600

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 MFU35–55% (not a failure — an accounting identity)
MFU vs HFUHFU counts recomputation as useful work; always ≥ MFU. Ask which.
0.01 nats of loss33% 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 placementspread beats density: ~26× lower extrapolation variance, same cost
Goodput swing64% 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

BreakErrorFix
Attention at long context~4% at 2k, ~15% at 8k, ~245% at 128k, ~20× at 1Madd 12·L·T·d_model per token
MoE10–20×use active parameters, not total
Embeddingsup to 89% at small Nreport non-embedding N
Activation checkpointing6ND~8ND hardware FLOPsthat 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 kthat is MoE, entire
L_balance = 1.0 means perfectly balancedlarger is worse — the dashboard number
top_k == n_expertssparsity 1.0 — a dense model in an MoE costume
Active → FLOPs; total → memoryoff by 10–20× if swapped
Dropped tokens are silentthey ride the residual; no error is raised
max_over_mean > 3the router is collapsing — intervene now
Expert-parallel comms2 all-to-alls per layer; ~32 GB / ~2.6 s for an 8k prefill on 60 layers

Prefill vs decode — never confuse them

PrefillDecode
Processesthe whole prompt at onceone token at a time
Bound bycomputememory bandwidth
Arithmetic intensityhigh1–43 FLOP/byte (far below the ridge)
Optimizationpipelining, chunkingbatching, GQA, quantization, speculation
Parallelismpipelined prefill workspipelining 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.73C^0.5
10× compute →5.37× params, 1.86× data~3.2× each
Methodone run per size, intermediate lossesseparate runs, properly decayed LR
Resultmodels were UNDERTRAINEDsmaller 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

VerticalThe problemThe hard part
Distillationpush teacher statistics into a studentstorage: 10T tokens × 256k vocab = 5 EB. Top-k or online.
Inference co-designshapes that saturate every hardware unitquality vs MFU pull opposite ways; differs per chip generation
QuantizationFP32 → 4 bitsoutliers; 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ᵀ and dW = 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

  1. Loss spike — decide in minutes: skip the batch, lower LR, or roll back.
  2. Silent data corruption — no crash, model quietly degrades. Cross-replica checksums.
  3. Data-iterator bug — wrong shard or an off-by-one. Days of compute on the wrong data.
  4. Router collapse (MoE) — watch max_over_mean; above 3 is an emergency.
  5. 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 = 6ND is 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 with 6ND.
  • Phase 02 optimizes lifetime cost = 6ND (train) + 2N·D_inf (serve).
  • Phase 03 asks what N even 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

  1. Derive 2N forward / 4N backward / 6N total from the definition of a matmul, in under two minutes, on a whiteboard.
  2. 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.
  3. State the four regimes where 6ND breaks (attention at long context, MoE, embeddings, activation checkpointing) and estimate the error in each.
  4. Compute training memory: weights + gradients + optimizer states + activations, and say which term dominates at which scale.
  5. Compute KV-cache size and show how GQA changes it by the group factor.
  6. Convert between: chips × days ↔ FLOPs ↔ (N, D) ↔ dollars ↔ joules.
  7. Answer "1,000 H100 for 30 days — what do you train?" with a number and a defence.

The lab

LabWhat you build
Lab 01 — Transformer FLOPs, Memory & Budget CalculatorA 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 6ND without 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

  1. A matmul costs 2 × (number of weights) FLOPs per token. Everything else follows.
  2. Backward is exactly 2× forward because each forward matmul becomes two backward matmuls (dX and dW).
  3. 6ND excludes attention. Fine at 2k context, badly wrong at 128k.
  4. For MoE, N is active parameters. Total parameters determine memory, not FLOPs.
  5. Memory, not FLOPs, is usually what stops you. Adam in mixed precision costs ~16 bytes per parameter before you have stored a single activation.
  6. The KV cache is the serving wall, and kv_heads is 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

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:

TermMeaningExample
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.

FormatBitsRangePrecisionTypical use
FP3232huge~7 decimal digitsoptimizer states, loss accumulation
TF3219 (in 32 slots)FP32 range~3 digitsNVIDIA matmul default
BF1616FP32 range~3 digitsthe workhorse for training
FP1616small~3 digitsolder training, overflows easily
FP88small~1 digitfrontier training/inference
INT44integer codesweight-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 × n entries.
  • Each entry is a sum of k products.
  • Each product-and-accumulate is 2 FLOPs (one multiply, one add).

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

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

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

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

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

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

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

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

How it works underneath

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

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

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

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

import math

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

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

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

Production significance

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

The misconception

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

No. The 2 is multiply-plus-add, in the forward pass alone. The forward/backward factor is a separate , 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)

SymbolNameTypical valueWhat it is
d or d_modelmodel dimension4096–16384width of the residual stream
d_ff or FFFN hidden dim~4·d (or ~2.7·d for SwiGLU)width of the MLP's middle
n_h or N_headsquery heads32–128parallel attention heads
n_kvkey/value heads1–128fewer than n_h under GQA/MQA
d_h or Hhead dimension64–128usually d / n_h
Llayers32–120depth
Vvocabulary32k–256knumber of distinct tokens
Bbatch sizesequences processed together
Tsequence length2k–1Mtokens 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_h to keep variance stable (without this, large d_h produces enormous logits and a saturated softmax with vanishing gradients).
  • softmax: turns scores into weights that sum to 1 per row.
  • A·V: each position's output is the weighted average of all values.
  • W_o: projects the concatenated heads back to d.
  • MLP: expand to d_ff, apply a nonlinearity, project back. The nonlinearity is the only reason depth helps — without it, a stack of linear layers collapses to a single linear layer.

The misconception

"Attention is where all the computation is."

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


Chapter 4: Counting the Parameters

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

Per layer

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

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

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

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


p = params_per_layer(d_model=4096, d_ff=11008, n_heads=32,
                     n_kv_heads=8, d_head=128, gated=True)
for k, v in p.items():
    print(f"{k:6s} {v:>14,}")
attn      41,943,040
mlp      135,266,304
norms          8,192
total    177,217,536

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

Whole model

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


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

The embedding trap

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

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

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


Chapter 5: The 6ND Derivation

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

Step 1 — Forward: 2N FLOPs per token

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

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

Step 2 — Backward: 4N FLOPs per token

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

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

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

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

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

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

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

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

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

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

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

Step 3 — Add

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

Over D tokens:

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

The three numbers to keep separate

QuantityFLOPs per tokenWhen you use it
Forward only2Ninference / prefill; teacher forward in distillation
Backward4N
Full training step6Npre-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:

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

Chapter 6: The Exact Per-Step Count

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

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

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

The MLP term: 18BTDF

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

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

Over B·T tokens: 18BTDF. ✓

The attention-projection term: 24BTDNH

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

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

Over B·T tokens: 24BTDNH. ✓

The factoring

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

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

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

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

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

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

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


cfg = dict(B=8, T=8192, d_model=8192, d_ff=28672,
           n_heads=64, d_head=128, n_layers=80)
r = exact_step_flops(**cfg)
for k, v in r.items():
    print(f"{k:10s} {v:>22,}  ({v/r['total']:5.1%})")
mlp        22,166,154,415,964,160  (63.6%)
attn_proj   8,444,249,301,319,680  (24.2%)
attn_seq    4,222,124,650,659,840  (12.1%)
total      34,832,528,367,943,680  (100.0%)

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

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


Chapter 7: Where 6ND Breaks

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

Break 1 — Attention at long context

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

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

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

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

for T in (2048, 8192, 32768, 131072, 1048576):
    f = attention_fraction(70e9, 80, 8192, T)
    print(f"T={T:>9,}  attention = {f:6.1%} of FLOPs   "
          f"(6ND error {f/(1-f):7.1%})")
T=    2,048  attention =   3.7% of FLOPs   (6ND error    3.8%)
T=    8,192  attention =  13.3% of FLOPs   (6ND error   15.3%)
T=   32,768  attention =  38.0% of FLOPs   (6ND error   61.4%)
T=  131,072  attention =  71.1% of FLOPs   (6ND error  245.4%)
T=1,048,576  attention =  95.2% of FLOPs   (6ND error 1963.4%)

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

Break 2 — Mixture of Experts

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

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

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

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

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

Break 3 — Embeddings

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

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

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

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

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

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


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

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

The four consumers

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

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

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

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


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

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

Activations

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

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

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

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

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

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


Chapter 9: The KV Cache — The Serving Wall

What it is

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

The size

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

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

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

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

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

Why this is the serving constraint

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

Available for KV = (total HBM) − (weights) − (workspace)
Max concurrent   = Available / (KV bytes per request)
def max_concurrent(total_hbm_gb, n_params, kv_per_request_gb,
                   bytes_per_param=2, workspace_gb=10):
    weights_gb = n_params * bytes_per_param / 1e9
    free = total_hbm_gb - weights_gb - workspace_gb
    if free <= 0:
        return 0
    return int(free / kv_per_request_gb)

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

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

Under the hood — why decode is memory-bound

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

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

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

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


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

Now answer Feinberg's opening question end to end.

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

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

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

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

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

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

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

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

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

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

for k, (ok, detail) in sanity_checks(N, D, "H100", 1000).items():
    print(f"  [{'OK ' if ok else 'FAIL'}] {k:24s} {detail}")
  [OK ] memory fits (sharded)    1480 GB / 80000 GB
  [OK ] data available           needs 1.85T tokens
  [OK ] servable                 92B params, 185 GB in bf16
  [OK ] energy                   504.0 MWh, ~$60,480 at $0.12/kWh

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


Chapter 11: Money, Watts and the Other Units

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

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

for k, v in cost_report("H100", 1000, 30, price_per_chip_hour=2.50).items():
    print(f"{k:20s} {v:>14,.2f}")
rental $               1,800,000.00
energy kWh               604,800.00
energy $                  72,576.00
FLOPs @40% MFU  1.026e+24
$ per 1e21 FLOPs           1,753.65

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

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


Lab Walkthrough

Lab 01 — Transformer FLOPs, Memory & Budget Calculator

You will implement, in order:

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

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

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


Success Criteria

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

Interview Q&A

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

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

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

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

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

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

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


Tips & Takeaways

Tips

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

Takeaways

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

References

Primary

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

Papers

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

Hardware

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

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

GroupFunctionsThe idea
The primitivematmul_flops, padded_matmul_flops2mkn; and what tile padding costs you
Parametersparams_per_layer, total_params, moe_parameter_splitGQA/MQA, gated MLPs, the embedding trap, active vs total
FLOPstraining_flops, inference_flops, training_flops_exact, unembedding_flops, attention_flop_fraction6ND, 2N, the exact per-shape count, and where 6ND breaks
Memoryactivation_bytes, training_memory, kv_cache_bytes, max_concurrent_requests, decode_arithmetic_intensity, ridge_pointAdam's 16 bytes/param, ZeRO stages, the serving wall
Budgetsbudget_to_flops, flops_to_days, chinchilla_split, lifetime_flops, cost_report, budget_reportchips × days ↔ FLOPs ↔ (N, D) ↔ dollars ↔ watts

Key concepts

ConceptWhy it is in this lab
2mknA linear layer costs 2 × params per token. The root of everything.
6NDForward 2N + backward 4N. The currency conversion of the field.
18BTDF + 24BTDNHThe slide identity — 6ND written out in shapes. A test asserts it.
Active vs total parametersMoE: active for FLOPs, total for memory. Off by 13× if you swap them.
Non-embedding NA 256k vocab is 89% of a small model. Scaling ladders must exclude it.
16 bytes/paramMixed-precision Adam, before a single activation.
ZeRO stagesOptimizer states are 75% of training memory — shard those first.
KV cache2·L·n_kv·d_h·T·B·b. The serving wall; n_kv is the lever.
Ridge pointpeak FLOP/s ÷ HBM bandwidth. Below it you are memory-bound.
Lifetime FLOPs6N·D_train + 2N·D_inf. Chinchilla's blind spot.

Files

FileWhat it is
lab.pyYour implementation. Signatures, docstrings and validation contracts are given; the arithmetic is yours.
solution.pyReference. python solution.py prints a twelve-part worked example.
test_lab.py75 tests: happy path, validation, boundaries, invariants, determinism.
requirements.txtpytest 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

  1. matmul_flops and get test_matmul_flops_is_2x_weight_count passing. That single identity is the root of 6ND; if it is not obvious to you yet, re-read WARMUP Chapter 2.
  2. params_per_layer, then test_gqa_shrinks_only_kv_projections. It forces you to notice that GQA touches W_k/W_v and leaves W_q/W_o alone.
  3. training_flops_exact, then test_exact_step_reproduces_the_slide_identity. This is the money test — it checks your per-shape accounting against Feinberg's slide.
  4. 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.py reaches 75 passed.
  • python solution.py runs and you can explain every one of its twelve sections.
  • Without the calculator, you can derive 6ND on paper in under two minutes.
  • You can state the four places 6ND breaks and the error in each.
  • You have run budget_report 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 and why.

How this maps to the real stack

This labThe real thingWhere the miniature lies
training_flops_exacttorch.utils.flop_counter.FlopCounterMode; JAX cost_analysis(); NVIDIA NsightReal counters trace the actual graph, so they catch fused ops, custom kernels, and recomputation. Ours assumes a textbook architecture.
training_memorytorch.cuda.max_memory_allocated(); DeepSpeed's memory estimator; torch.distributed.fsdpIgnores allocator fragmentation, NCCL buffers, CUDA context (~1 GB), and per-framework overheads. Real usage runs 10–30% above this.
kv_cache_bytesvLLM's PagedAttention block manager; TensorRT-LLM's KVCacheManagerReal engines page the cache in fixed blocks, so the true footprint rounds up to block granularity and can share prefixes across requests.
chinchilla_splitNobody ships a library for this — every lab has an internal versionUses 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_intensityThe roofline analysis in Nsight Compute; llm-analysisAssumes perfect bandwidth utilization. Real kernels achieve 60–90% of peak HBM bandwidth.
HARDWARE tableVendor spec sheetsPeak 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:

  1. Validate against a real model. Load a small HF checkpoint, count its parameters with sum(p.numel() for p in model.parameters()), and check total_params matches. Then run FlopCounterMode on a forward pass and check inference_flops. Any mismatch is a lesson.
  2. Add pipeline/tensor-parallel memory. Extend training_memory with tp_degree and pp_degree, and model the activation memory that pipeline stages must hold in flight. Phase 04 will need it.
  3. Add an MoE memory model. Expert-parallel sharding puts different experts on different chips — extend moe_parameter_split to report per-device memory given ep_degree.
  4. 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 fixed C. That is Phase 01, and doing it here first makes Phase 01 trivial.
  5. 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 = 6ND with 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ᵀ and dW = 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:

  1. You are extrapolating three orders of magnitude beyond your largest data point. Any functional form fits your ladder; they diverge wildly at the target.
  2. 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.
  3. 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

  1. 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).
  2. Run an IsoFLOPs sweep: fix C, vary N, fit a parabola in log N, extract the vertex.
  3. Fit power laws N_opt ∝ C^a, D_opt ∝ C^b and check a + b ≈ 1 (which it must, since C = 6ND).
  4. Explain precisely why Kaplan got a ≈ 0.73 and Chinchilla got a ≈ 0.5 — the LR-decay bias — and demonstrate the bias numerically.
  5. Fit the parametric form with a robust loss on log L, and show that least-squares on raw L gives different (worse) exponents.
  6. Produce a forecast with a confidence interval, via bootstrap over the ladder points.
  7. Compare a baseline recipe against a candidate recipe as fitted laws at the target C, and report the crossover.
  8. Choose ladder design points to minimize extrapolation variance instead of filling a grid.

The lab

LabWhat you build
Lab 01 — IsoFLOPs Ladder, Law Fitting & the ForecastA 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 ≈ 1 is 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 C and reported the crossover.
  • I can justify where to place my next ladder run in terms of information gain.

Key takeaways

  1. A scaling law is a property of a recipe. Fix the recipe first; the law describes that.
  2. Loss forecasting is recipe selection. Compare fitted curves at the target, not runs at small scale.
  3. The U-curve at fixed C is 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).
  4. a + b ≈ 1 falls out of C = 6ND. If your fitted exponents violate it, you have a bug.
  5. Reading loss mid-run is biased, systematically against training on more data. This one methodological error set the industry's strategy for two years.
  6. The estimator matters. Least-squares on L, least-squares on log L, Huber on log L, and MLE give different exponents on identical data — and therefore different flagship recommendations. Feinberg's own listed open problem: "Formalize."
  7. 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

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 — parameters
  • D — training tokens
  • E, 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:

TermNameMeaningWhat drives it to zero
Eirreducible lossThe entropy of language itself. Even a perfect model cannot predict the next token with certainty.Nothing. It is a floor.
A/N^αcapacity termYour model is too small to represent the true function.More parameters.
B/D^βdata termYou 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 to 1/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:

  1. It is smooth. Accuracy is a step function — the gradient is zero almost everywhere. Cross-entropy has a gradient everywhere.
  2. 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.
  3. 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, huge D: oceans of data, no capacity to absorb it. The A/N^α term dominates.
  • Huge N, tiny D: enormous capacity, starved of examples. The B/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=1e9 and N=3e9 differed 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 1e19 to 122 at 1e25. 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 particular E/A/B values 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 L weights 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 L weights relative error, treating all scales evenly.
  • Huber on log L does 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 scipy upgrade can move your N_opt by 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:

  1. 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?
  2. Derive the maximum-likelihood estimator under that model.
  3. 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 widthVerdict
< 0.01 natsTight. You can make a confident recipe call.
0.01–0.05Usable for go/no-go, not for choosing between similar recipes.
> 0.05Your 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

  1. Budget. The ladder should cost a small fraction of the flagship — a few percent is typical, and Phase 00's $ per 1e21 FLOPs makes that concrete.
  2. 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.
  3. Spread. Chapter 9: variance is driven by the spread of log C.
  4. 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:

  1. Derive L*(flops) baseline
  2. L*(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:

  1. fit_line, fit_power_law, fit_quadratic, parabola_vertex — the numerical spine. fit_quadratic uses Cramer's rule on the 3×3 normal equations; it must raise on a degenerate (collinear) design rather than divide by zero.
  2. synthetic_ladder — generate (N, D, L) with seeded noise and an injectable lr_decay_bias. This is what lets you reproduce the Kaplan bug.
  3. isoflop_optimum — sweep, fit the parabola, return the vertex.
  4. fit_scaling_exponents — steps 5–6, plus the a + b ≈ 1 consistency check.
  5. residuals / huber_loss / fit_parametric — the surface fit. Deterministic coordinate descent; same input must give byte-identical output.
  6. analytic_optimum — the closed form. A test checks it agrees with the numerical sweep.
  7. bootstrap_forecast — seeded resampling, percentile interval.
  8. extrapolation_variance / design_ladder — the optimal-design tooling.
  9. 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_quadratic on 2 points, or on collinear points → must raise, not ZeroDivisionError.
  • 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 -v all green.
  • Your lab.py passes after the TODOs.
  • python solution.py reproduces 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 ≈ 1 must 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

  1. Pre-training is a one-shot extrapolation problem. Scaling laws exist because of that.
  2. A law describes your recipe, not the universe. Fix the recipe first.
  3. Loss forecasting is recipe selection.
  4. 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.
  5. The IsoFLOPs U-curve has a flat bottom: a gift for serving-driven deviation, a trap for noisy fits.
  6. Chinchilla beat Kaplan by fixing an experimental-design bug. Methodology is the frontier.
  7. The estimator you choose changes your flagship recommendation. This is an open problem.
  8. Where you place ladder points matters more than how many. 13× variance reduction, same cost.
  9. A forecast without error bars is not a forecast.
  10. 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

GroupFunctionsThe idea
Numerical spinefit_line, fit_power_law, fit_quadratic, parabola_vertexA power law is a straight line in log-log; the IsoFLOPs minimum is a parabola vertex
Loss model & ladderparametric_loss, loss_decomposition, schedule_bias, synthetic_ladderL = E + A/N^α + B/D^β, plus an injectable measurement bias
IsoFLOPsisoflop_optimum, isoflops_ladder, fit_scaling_exponentsThe six-step method, ending in the a + b ≈ 1 check
Surface fitresiduals, squared_loss, huber_loss, fit_parametricLog-space, robust, deterministic multi-start
Optimumanalytic_optimum, optimum_exponentThe closed form, and why Chinchilla's exponent is ~0.5
Uncertaintybootstrap_forecast, compute_multiplier_for_loss_deltaError bars, and how to read them in dollars
Designextrapolation_variance, design_ladder, extrapolation_riskWhere to spend your next ablation
Decisionloss_at_budget, compare_recipes, decision_is_supportedCompare laws at the target; report the crossover

Key concepts

ConceptWhy it is in this lab
L = E + A/N^α + B/D^βIrreducible + capacity + data. At frontier scale ~87% is irreducible.
The IsoFLOPs UAt fixed C, too-small underfits and too-large starves. Flat bottom.
a + b ≈ 1Falls out of C = 6ND. A free bug detector on any fitted pair.
Non-uniform biasA constant offset is absorbed into E; only a varying one tilts exponents.
Huber on log residualsOne diverged run costs least-squares ~4,000× what it costs Huber.
Bootstrap CIThe model is nonlinear and the noise model unspecified — so resample.
0.01 nats ≈ 33% computeThe conversion that makes a confidence interval interpretable.
Extrapolation varianceSpread beats density: 26× lower variance from placement alone.
The crossover"Candidate wins above 3e24" is actionable; "candidate is better" is not.

Files

FileWhat it is
lab.pyYour implementation. Signatures, docstrings and validation contracts given.
solution.pyReference. python solution.py runs the full eleven-part story.
test_lab.py64 tests: happy path, validation, boundaries, invariants, determinism.
requirements.txtpytest 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.

  1. The spine. fit_linefit_power_lawfit_quadraticparabola_vertex. Get test_fit_power_law_recovers_exponent and test_fit_quadratic_recovers_exact_parabola passing before anything else.
  2. parametric_loss and synthetic_ladder. Note the budget constraint: every point on an IsoFLOPs curve satisfies C = 6ND, and a test checks it.
  3. The IsoFLOPs method, ending at test_exponents_sum_to_one.
  4. fit_parametric. Deterministic multi-start coordinate descent. test_fit_parametric_is_deterministic is not optional garnish — a fit that moves between runs cannot support a nine-figure decision.
  5. 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_quadratic on collinear x values → must raise ValueError, not ZeroDivisionError.
  • 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 = 6ND to 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.py reaches 64 passed.
  • python solution.py runs, 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 ≈ 1 must 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 labThe real thingWhere the miniature lies
synthetic_ladderActual training runs on a cluster, days eachReal 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_parametricscipy.optimize.minimize with L-BFGS-B, or JAX + optaxReal 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_optimumThe same parabola fit, on real measured lossesIdentical in spirit. Real practice uses more sizes per budget (6–10) and often fits in log-loss space.
bootstrap_forecastBootstrap or Bayesian posterior over law parametersReal work often uses a hierarchical model that shares information across budgets. Ours treats points as exchangeable, which understates correlation within a budget.
compare_recipesAn internal go/no-go memo and a review meetingThe 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 constantsHoffmann et al. Table 3These 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

  1. 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.
  2. 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.
  3. Implement true optimal design. Instead of scoring fixed candidate designs, search for the next C that minimizes posterior predictive variance at the flagship. That is the "active learn" bullet on his slide, and it is a publishable exercise.
  4. Add a third axis. Extend to L(N, D, U) with unique tokens, then compare against Phase 02's data-constrained law.
  5. 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) and D_opt(C) with an a + b = 1 consistency check, and robust (Huber-on-log) fitting of the parametric L(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 gainYou pay in
More parameters at the same FLOPs per tokenHBM — every expert must be resident
Better loss per unit of computeCommunication — all-to-all every layer if you shard experts
Specialization across the corpusStability — 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

  1. Implement a top-k router with renormalized gates, and explain what not renormalizing silently does to the layer's output scale.
  2. 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.
  3. Write the router z-loss and explain the two distinct failures it prevents (bf16 overflow, and a saturated softmax whose gradient vanishes).
  4. Compute expert capacity, and quantify the drop-rate / wasted-compute trade-off that tuning capacity_factor actually buys.
  5. Explain why a fully-dropped token is a silent quality loss, and how a shared expert removes that failure mode entirely.
  6. Separate total parameters (memory) from active parameters (FLOPs) and never confuse them again.
  7. Compute the all-to-all traffic of naive expert parallelism and state, in seconds, why it is fatal for an interactive product.
  8. Simulate router collapse and demonstrate the auxiliary loss preventing it.

The lab

LabWhat you build
Lab 01 — Router, Load Balancing, Capacity & the MoE LayerA 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 -v53 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.0 in 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 = 0 and recover with 0.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

  1. Parameters scale with E; FLOPs scale with k. That single sentence is what MoE is.
  2. top_k == n_experts is a dense model wearing an MoE costume. Sparsity ratio 1.0.
  3. 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.
  4. L_balance = 1.0 means perfectly balanced. Memorize the calibration point — it makes the dashboard readable at a glance.
  5. Dropped tokens are silent. No exception, no log line; the token rides the residual and your loss is quietly worse. Watch drop_rate like a hawk.
  6. Capacity factor trades dropping against wasted compute. There is no setting that avoids both.
  7. A shared expert guarantees every token gets some FFN, which removes the fully-dropped failure mode at the cost of always-on compute.
  8. 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.
  9. 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=1 routing)
  • 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

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:

  1. It is where the parameters are — about three-quarters of the block.
  2. 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

kNameTrade-off
1Switch routingCheapest. Sharpest specialization. Hardest to train — a single routing mistake has no backup.
2The common defaultTwo experts give the gradient two paths and the token a fallback. Used by Mixtral, DeepSeek, most production models.
4–8Fine-grained MoEUsed with many small experts (DeepSeek-V3 uses 8 of 256). More combinations, better specialization, more routing overhead.
EDenseSparsity 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.79 and 0.21, summing to 1. The layer output has a consistent scale regardless of how confident the router was.
  • Raw: gates stay 0.575 and 0.151, summing to 0.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 constant c — 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_mean utilization and L_balance. If max_over_mean climbs 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:

  1. No token can ever be fully dropped. The failure mode above is structurally eliminated.
  2. 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:

QuestionUse
How much compute to train? C = 6NDactive
How much compute to serve? 2N/tokenactive
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.

RefinementWhat it doesWhy
Fine-grained expertsMany small experts (256) instead of few large ones (8), with higher kMore combinations of experts per token → better specialization at the same active parameter count
Shared expertsOne or more experts that always runRemoves the fully-dropped failure mode; lets routed experts specialize harder (Chapter 7)
Auxiliary-loss-free balancingA per-expert bias added to router logits, adjusted by observed load, instead of an aux lossThe 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 routingInvert it: each expert picks its top tokens, rather than each token picking expertsBalance 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 MoEVariable-size expert buffers via block-sparse kernelsEliminates 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")
UpcyclingInitialize MoE experts from a trained dense checkpointSkips 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:

  1. softmax, logsumexp, matvec, relu. Get test_softmax_survives_huge_logits passing first — max-subtraction is the lesson, not a detail.
  2. router_logits, route_token, route_batch. Watch the renormalize flag; both behaviours are tested. Break ties by expert index.
  3. load_balance_loss. Aim for the calibration point: test_balanced_routing_gives_aux_loss_of_exactly_one. If you do not get exactly 1.0, the formula is wrong.
  4. router_z_loss, expert_utilization.
  5. expert_capacity, apply_capacity. Note that capacity_factor < 1.0 must raise — it drops tokens even under perfect balance, which is never intended.
  6. moe_forward. The residual fallback for fully-dropped tokens is the important line.
  7. moe_parameter_counts, expert_parallel_comm_bytes, comm_seconds.
  8. simulate_collapse. The money test.

The traps:

  • A naive softmax gives NaN on large logits and the test will catch you.
  • The load-balance loss divides f by total slots (tokens × k), not by token count.
  • apply_capacity is first-come-first-served in token order — deterministic, and the tests depend on it.
  • top_k == n_experts must give sparsity exactly 1.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.py reaches 53 passed.
  • python solution.py runs and you can explain all eight sections.
  • Balanced routing gives a load-balance loss of exactly 1.0 in your code.
  • You can explain why the aux loss multiplies f by P rather than penalizing f alone.
  • You have watched a router collapse at aux_weight = 0 and recover at 0.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_mean and drop_rate on the dashboard. They are the two numbers that tell you an MoE is failing, and both fail silently otherwise.
  • Memorize L_aux = 1.0 means 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

  1. Parameters scale with E; FLOPs scale with k. That is MoE, entire.
  2. MoE targets the MLP because it holds ~76% of the parameters and is position-independent.
  3. Renormalizing the top-k gates keeps the output scale consistent; not doing so silently attenuates the layer by a per-token amount.
  4. Routers collapse by default. The auxiliary loss is what prevents it, and 1.0 is balanced.
  5. The f·P product exists to get a gradient through a discrete decision.
  6. The z-loss prevents overflow and softmax saturation — two different failures.
  7. Capacity trades dropping against wasted compute. Dropping is silent, which makes it worse.
  8. A shared expert eliminates fully-dropped tokens and lets the routed experts specialize.
  9. Active for FLOPs, total for memory. Off by 10–20× if you swap them.
  10. Naive expert parallelism costs seconds of network time. The fix is a different sharding axis, not a better kernel.
  11. 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

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

GroupFunctionsThe idea
Primitivessoftmax, logsumexp, matvec, reluMax-subtraction is the lesson, not a detail
Routerrouter_logits, route_token, route_batchTop-k, gate renormalization, deterministic ties
Aux lossesload_balance_loss, router_z_loss, expert_utilizationWhat stops collapse, and the dashboard numbers
Capacityexpert_capacity, apply_capacityFixed buffers, dropping, padding
The layermake_expert, expert_forward, moe_forward, total_training_lossFull forward pass with shared expert and residual fallback
Accountingmoe_parameter_counts, expert_parallel_comm_bytes, comm_secondsTotal vs active; the communication wall
The failuresimulate_collapseRich-get-richer, and the aux loss preventing it

Key concepts

ConceptWhy it is in this lab
Max-subtracted softmaxRouter logits grow without bound; naive exp gives NaN and kills the run
Gate renormalizationWithout it the layer output is silently attenuated by a per-token amount
L_aux = E·Σ f_i·P_iDiscrete f supplies the signal, differentiable P supplies the gradient
L_aux = 1.0The calibration point: exactly balanced. Larger is worse.
Router z-lossPrevents bf16 overflow and softmax saturation — two distinct failures
Capacity factorTrades dropped tokens against padded (wasted) compute. No free setting.
Silent droppingA fully-dropped token rides the residual. No error. Quality quietly drops.
Shared expertStructurally eliminates the fully-dropped case
Total vs activeActive → FLOPs (6ND); total → HBM. Off by 10–20× if swapped.
All-to-all cost2 collectives per layer under expert parallelism — the Flash 2.0 wall

Files

FileWhat it is
lab.pyYour implementation. Signatures, docstrings and validation contracts given.
solution.pyReference. python solution.py runs an eight-part worked example.
test_lab.py53 tests: happy path, validation, boundaries, invariants, determinism.
requirements.txtpytest 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

  1. softmax → get test_softmax_survives_huge_logits green. Logits of 1000 must not produce NaN. This is the most common real-world MoE crash.
  2. route_tokentest_renormalized_gates_sum_to_one and its raw counterpart. Both behaviours are tested because both exist in the wild.
  3. load_balance_loss → aim for test_balanced_routing_gives_aux_loss_of_exactly_one. If you do not get exactly 1.0, your formula is wrong. Remember f divides by total slots (tokens × k), not token count.
  4. apply_capacity → first-come-first-served in token order; the tests depend on that determinism.
  5. moe_forward → the residual fallback for fully-dropped tokens is the line that matters.
  6. simulate_collapse → the money test.

The traps:

  • Naive softmax → NaN on large logits. Tested.
  • Dividing f by token count instead of slot count → your L_aux will not be 1.0.
  • capacity_factor < 1.0 must raise — it drops tokens even under perfect balance.
  • top_k == n_experts must give sparsity exactly 1.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.py reaches 53 passed.
  • python solution.py runs and you can explain all eight sections.
  • You have seen aux_weight=0.0 collapse to max/mean = 4.10 and 0.01 hold at 1.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 labThe real thingWhere the miniature lies
route_token / route_batchMegatron-LM MoELayer, DeepSpeed-MoE, Mixtral's MoeBlock, JAX/Flax MoEReal routers run fused on-device with the top-k in a kernel. Ours is a Python loop — same math, ~10⁶× slower.
load_balance_lossswitch_load_balancing_loss in Megatron/DeepSpeedIdentical 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_capacityGShard/Switch dispatch-and-combine via one-hot matmuls or gather/scatterReal 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_forwardA 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_bytesNCCL/all_to_all in Megatron; XLA collectives on TPUOurs ignores topology (a torus vs a fat-tree changes the constant a lot), overlap, and compression. Directionally right, not a substitute for profiling.
simulate_collapseThe real dynamic, observed on training dashboardsA 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

  1. 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.
  2. 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.
  3. 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.
  4. Model the pipelined-prefill fix. Extend expert_parallel_comm_bytes with a pipelined alternative that shards layers instead of experts, and reproduce the order-of-magnitude improvement in the transcript dissection's Claim 12.
  5. 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

  1. Compute arithmetic intensity and place any operation on a roofline; say whether to attack bytes or FLOPs.
  2. 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.
  3. Produce an MFU budget — decompose wall-clock into matmul / vector / memory / comms / optimizer — and read off the biggest lever.
  4. Distinguish MFU from HFU and explain why HFU is always higher (~33% for free).
  5. 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.
  6. Show that at batch 1, decode costs ~3.8× prefill — and explain why that changes the provisioning answer.
  7. Quantify three co-design levers: tile quantization, GQA group size, and depth-vs-width.
  8. Prove that decode stays memory-bound even at batch 1024, and explain what follows from that.

The lab

LabWhat you build
Lab 01 — Roofline, MFU Budget & the Latency NapkinThe 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 -v50 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

  1. Arithmetic intensity decides which resource you are fighting. It is a property of the algorithm, not the chip.
  2. Below the ridge, cut bytes; above it, cut FLOPs. Everything else is detail.
  3. MFU of 35% is an identity, not a grade. The breakdown is your agenda.
  4. HFU ≥ MFU, always. Activation checkpointing counted as useful work is ~33% for free.
  5. Prefill and decode are different machines. Never apply one's fix to the other.
  6. Decode is memory-bound at every realistic batch size — even 1024. That single fact explains batching, GQA, quantization and speculative decoding.
  7. n_kv_heads is the biggest single serving lever, it costs almost no quality, and it is frozen at pre-training time.
  8. Depth is serial and costs latency; width is parallel. As wide as quality allows, as shallow as quality tolerates.
  9. 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

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_ridgememory-bound. The matmul unit is starved. Adding FLOPs is free; removing bytes is what helps.
  • I > I_ridgecompute-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
commsoverlap collectives with compute; change the sharding plan
memoryfuse kernels; use FlashAttention; reduce activation traffic
vectorfuse norms and activations into matmul epilogues
optimizershard 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

PrefillDecode
Work per stepwhole promptone token
Bound bycomputememory bandwidth
Arithmetic intensityhundreds–thousands~1–50
Helped bymore FLOP/s, pipelining, chunkingmore bandwidth, batching, fewer bytes
Hurt bylong prompts ( attention)large models, large KV cache
Parallelism that workspipelining (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−1 serial 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+1 cannot start until layer k finishes. 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

TechniqueWhat it really does
Continuous batchingraise intensity by amortizing the weight read
GQA / MQAshrink the KV term so batching keeps working longer
Quantizationhalve or quarter the bytes; the dominant term
Speculative decodingverify k tokens per weight-read pass — raises intensity directly
PagedAttentionfit more requests in the same HBM, so batches can be bigger
Smaller modelsattack 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:

  1. arithmetic_intensity, ridge_point, roofline_throughput, roofline_report. The roofline is one min(); the value is the lever verdict it returns.
  2. mfu, hfu, mfu_budget. test_hfu_is_always_at_least_mfu is the one to internalize.
  3. prefill_seconds, decode_seconds, interactive_latency, chips_for_latency_budget, weights_fit_chips.
  4. 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 2N per token (forward only), not 6N. That is training.
  • Decode time is driven by bytes, not FLOPs. If your decode_seconds has peak in it, it is wrong.
  • chips_for_latency_budget must return None rather 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.py reaches 50 passed.
  • python solution.py runs 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_heads in 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

  1. A chip is several machines; "peak" describes only the matmul unit.
  2. Arithmetic intensity decides which resource you are fighting. Below the ridge, cut bytes.
  3. Newer hardware has a higher ridge point — memory-bound work gets relatively worse.
  4. MFU of 35% is an accounting identity. The decomposition is your work queue.
  5. HFU ≥ MFU always. Ask which.
  6. Prefill is compute-bound; decode is memory-bound. Two machines, opposite optimizations.
  7. Pipelining helps prefill and hurts decode. This is why disaggregation exists.
  8. Decode never reaches the ridge at any realistic batch size — even 1024.
  9. The KV cache is what makes batching saturate, which is why GQA matters so much.
  10. 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

GroupFunctionsThe idea
Rooflinearithmetic_intensity, ridge_point, roofline_throughput, roofline_reportFLOPs per byte; below the ridge cut bytes, above it cut FLOPs
Utilizationmfu, hfu, mfu_budget35% is an identity, not a grade; the breakdown is an agenda
Latencyprefill_seconds, decode_seconds, interactive_latency, chips_for_latency_budget, weights_fit_chipsThe napkin, end to end
Co-designtile_efficiency, kv_cache_bytes, gqa_saving, decode_batch_intensity, depth_vs_widthThe three irreversible levers

Key concepts

ConceptWhy it is in this lab
Arithmetic intensityProperty of the algorithm; decides which resource you fight
Ridge pointpeak ÷ bandwidth. H100 ≈ 296 — and higher than A100's 154
MFU vs HFUHFU counts recomputation as useful. Always ≥ MFU, ~33% for free
MFU budgetmatmul / vector / memory / comms / optimizer → your work queue
Prefill 2N/tokenCompute-bound. 6N is training; using it here is a 3× error
Decode = bytesEvery token re-reads every weight. 140 GB per token at 70B
The 4×4 station16 v5e chips to get 8k prefill under 0.5 s — the talk's conclusion
GQA ratio = group size64→8 KV heads is exactly 8× less cache, ~0 quality cost
The decode wallEven batch 1024 reaches only 17% of the ridge

Files

FileWhat it is
lab.pyYour implementation. Signatures, docstrings and validation contracts given.
solution.pyReference. python solution.py runs a nine-part worked example.
test_lab.py50 tests: happy path, validation, boundaries, invariants, determinism.
requirements.txtpytest 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

  1. arithmetic_intensityridge_pointroofline_throughput. The roofline is one min(). Get test_roofline_is_continuous_at_the_ridge green — the two roof segments must meet exactly.
  2. mfu_budget, then test_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.
  3. prefill_seconds / decode_seconds. Then test_prefill_reproduces_the_talks_number — if you get ~5.8 s you have it right.
  4. chips_for_latency_budget. Note it must return None when the budget is unreachable.
  5. The co-design levers.

The traps:

  • Prefill is 2N per token, forward only. 6N is training — a 3× error.
  • decode_seconds must not contain peak. Decode is driven by bytes and bandwidth. If peak FLOP/s appears in it, the model is wrong.
  • chips_for_latency_budget must terminate and return None rather 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_saving must reject n_kv_heads > n_query_heads and non-divisible configurations.

Success criteria

  • LAB_MODULE=solution pytest test_lab.py -v → 50 passed.
  • Your lab.py reaches 50 passed.
  • python solution.py runs 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 labThe real thingWhere the miniature lies
roofline_reportNVIDIA Nsight Compute's roofline; Intel AdvisorReal tools measure achieved bytes and FLOPs from hardware counters. Ours computes the analytical intensity, which is the ceiling, not the achieved value.
mfu / mfu_budgetMegatron-LM and MaxText log MFU per step; profiler timelines give the breakdownReal 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_secondsvLLM / TensorRT-LLM benchmarks; llm-analysisOurs ignores attention's 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_savingThe num_key_value_heads field in any HF configExact. This one is not an approximation.
tile_efficiencyXLA / cuBLAS padding behaviourReal libraries pick among several tile sizes and may pad differently. The effect is real; the exact constant is not.
HARDWARE tableVendor spec sheetsPeak 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

  1. Add the attention term. At 128k context the sequence-dependent attention matmuls dominate prefill (Phase 00, Break 1). Extend prefill_seconds and watch the napkin's conclusions change completely for long-context products.
  2. 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.
  3. Add speculative decoding. Model an acceptance rate α and a draft model of size N_d; compute the effective intensity gain and find where it stops paying.
  4. 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.
  5. 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."