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