Curriculum

Week 9. Bookkeeping: parameters, FLOPs, memory

Phase 2. The modern transformer · week 9 of 24

Learn in the app: tutor, coding problems →

Core: parameters by component, 6N FLOPs, training and inference memory, KV cache size · Depth: exact activation memory, MFU and training cost, total versus active MoE parameters · ≈ 11 h core / 21 h total

This gets its own week because interviewers ask about it almost every time, and almost everyone gets the numbers wrong. Every decision later on starts from these formulas: how many GPUs, what batch size, whether the context fits, whether parallelism pays off (week 14). You need to be able to derive them, not recall them.

Step 1. Parameters

• Parameters by component: embeddings V·D, attention D·D + 2·D·K·H + D·D, FFN 3·D·F, norms. With MHA (K·H = D) attention gives 4D², and SwiGLU with F = 8D/3 gives 3·D·8D/3 = 8D². That is 12D² per layer, plus the input embedding and the head, 2VD: N ≈ 2VD + 12LD² (norms add 2D per layer; we drop them when rounding)
• Sanity check on Llama-3-8B: the formula gives 1.05B + 6.44B = 7.49B, while the real count is 8.03B. The gap is honest: GQA shrinks attention (1.34B instead of 2.15B), and F = 14336 = 3.5D inflates the FFN (5.64B instead of 4.29B)
• Embeddings dominate when 2VD > 12LD², that is, when V > 6LD. In GPT-2 small they are about a third of all parameters

Step 2. FLOPs

In plain terms. On every token, each weight is multiplied once and added once: 2 operations. The backward pass does that work twice: it computes the gradient with respect to the layer input, to pass it further back, and the gradient with respect to the weight itself. Total: 2 + 4 = 6 operations per weight per token. A 1B-parameter model on 20B tokens: 6 · 10⁹ · 2·10¹⁰ = 1.2·10²⁰ FLOPs.

• FLOPs: forward ≈ 2N per token, backward ≈ 4N, total 6N per token (N = number of parameters). The basis is simple: a matmul (b×n)·(n×m) costs 2bnm. Embeddings are not part of N here: they are a row lookup, not a multiplication
  • The attention contribution ~4·B·N·S²·H, and when it starts to dominate the matmuls. This comes from QKᵀ and AV, 2·B·N·S²·H each per layer in the forward pass. Per token, with backward, that is 12·L·S·D. Against 6N ≈ 72LD² the threshold is S ≈ 6D; for Llama-3-8B (budget.py) that is about 26.6k
• The D×V head is a real matmul, another 6VD per token. In flops_per_token it has its own line, lm_head: for Llama-3-8B it is ~3.2 GFLOPs, about 6% of the total
• MoE: total vs active parameters. In an MoE layer (mixture of experts, week 19), instead of one FFN there are E experts, and a router sends each token to k of them. In plain terms. 24 layers, each with 20M of attention and 8 experts of 60M, top-2. Total parameters: 24·(20 + 8·60)M = 12.0B; active per token: 24·(20 + 2·60)M = 3.36B. FLOPs are counted from the active ones: 6 · 3.36·10⁹ ≈ 2·10¹⁰ per token in training. Memory is counted from the total: 24 GB of weights in bf16, 192 GB at 16 bytes for AdamW. In decode with a large batch, tokens spread across all experts: a step reads all 12B but computes like 3.36B. That is why the critical batch (week 13) for MoE is higher by about 12 / 3.36 ≈ 3.6×. DeepSeek-V3 has 671B total parameters and 37B active: the memory of a giant, the arithmetic of a mid-size model

Step 3. Memory

In plain terms. When you train with AdamW, every parameter carries five "cards": the weight in bf16 (2 bytes), the gradient in bf16 (2), an exact fp32 copy of the weight (4), and two Adam moments in fp32 (4 + 4). Total: 16 bytes. A 7B model takes 112 GB before any activations: it no longer fits on a single 80 GB H100.

• Activations: what is stored for backward, and where O(B·L·S·D) comes from. Each matmul stores its input, SwiGLU stores the intermediate B·S·F, and naive attention also stores the B·N·S² matrix (the formula is in 05-ГЛУБИНА). For Llama-3-8B at batch 8 × 8192 that is 384 GiB, versus 120 GiB for the states. The cures are checkpointing (week 2) and microbatches
• Training memory: parameters + gradients + Adam states (m, v) + activations. For fp32 AdamW this is 16 bytes per parameter; for mixed precision you have to count carefully: 2 + 2 + 4 + 4 + 4, again 16. Mixed precision saves on activations and speed, not on states
• Inference memory: weights + KV cache = 2 · B · S · L · K · H · bytes. Derivation: for every token, in every layer, you store a key and a value (2), each with K heads of size H. Llama-3-8B, bf16: 128 KiB per token, 1 GiB per 8k of context; with MHA it would be 4 GiB
• How much the KV cache weighs. Weights do not depend on the batch; the cache grows linearly with it. The cache equals the weights at B = weight bytes / (S · 2 · L · K · H · bytes). For Llama-3-8B at 8k this is 16.06 GB / 1.07 GB ≈ 15 sequences, at 4k about 30. Beyond that, the cache is the main consumer of both memory and step time: decode reads all of it for every new token (week 11, decode_step_time)
• MFU (model FLOPs utilization), the measure of training efficiency: the fraction of peak delivered by useful FLOPs (6N plus attention per token × tokens per second). Recomputation from checkpointing does not count. Well-tuned training reaches 35–50%; training_time uses 45% by default

Common mistakes

• Computing the KV cache with N instead of K: caught by test_gqa_saves_kv_cache_exactly_N_over_K. Forgetting the factor of 2 for K and V: caught by test_kv_cache_formula, test_kv_cache_size_formula
• Including embeddings in 6N: caught by test_flops_rule_of_6N. Using F = 4D for SwiGLU: caught by test_swiglu_ffn_dim_is_8D_over_3
• Counting tied embeddings twice (tie_embeddings in Spec); mixing up GB and GiB (a 7% difference, no test will catch it)
• For MoE, counting 6N from the total parameters and memory from the active ones. It is the other way round: FLOPs from active, memory from total

Code → nanolm/budget.py: a calculator. Input: a model config and hardware (Spec). Output: parameter_breakdown, flops_per_token, training_memory, inference_memory, training_time, report. Run: python -m nanolm.budget --V 128256 --D 4096 --L 32 --N 32 --K 8 --S 8192 --F 14336, tests: pytest tests/test_budget.py -v (the key one is test_llama3_8b_parameter_count). Inside the model the same quantities are computed by NanoLM.num_parameters, flops_per_token, kv_cache_bytes. Check the predictions by measuring with torch.cuda.max_memory_allocated() (on a Mac, torch.mps.current_allocated_memory()).

Math (Track D): D9: estimating the training time of a 7B model from 6N, MFU and GPU count; D10: Stirling and order-of-magnitude estimates with large numbers.

Interview question of the week: "Will full fine-tuning of a 7B model fit on one 80 GB H100?" A 3-minute structure: 16 bytes per parameter → 112 GB of states, no → what to shard (ZeRO/FSDP, week 14) → what to compress (8-bit Adam moments) → activations: checkpointing, microbatches → the alternative: LoRA (week 15).

Sources: Kaplan et al., Scaling Laws (2020), the origin of 6N; Korthikanti et al., *Reducing Activation Recomputation* (2022); Chowdhery et al., PaLM (2022), the origin of MFU; Pope et al., *Efficiently Scaling Transformer Inference* (2022), which counts inference memory and time the same way; DeepSeek-V3 (2024), total and active MoE parameters.

Deeper: 05-ГЛУБИНА, section "Week 9. Accounting: exact formulas instead of 'roughly'".

Week outcomes

• I can derive N ≈ 2VD + 12LD² (for MHA) from the sum over components and say when embeddings dominate.
• I can estimate training time in 5 minutes from 6N per token, the token count and peak FLOPs, accounting for MFU.
• I can compute the KV cache for a given model, batch and context, and the batch at which it catches up with the weights.
• I can separate total and active MoE parameters: which go into FLOPs and which into memory.
• I can write the budget.py calculator and confirm the memory forecast by measurement.
• I can explain why the formula gives 7.5B for Llama-3-8B while the real count is 8.03B.

Self-check

1. Where do the 2 and the 4 in 6N come from?
2. At what context length does attention start to dominate the matmuls, and why?
3. What makes up the 16 bytes per parameter with AdamW in fp32, and what changes with mixed precision?