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

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.