"""Transformer FLOPs, Memory & Budget Calculator — reference solution. Every function in this module is a piece of the arithmetic a pre-training lead does on a napkin before spending eight figures of compute. Run `python solution.py` for a worked example. UNITS ARE NAMED IN EVERY SIGNATURE. Bytes are bytes, FLOPs are FLOPs, seconds are seconds. Most errors at this level are unit errors, not algebra errors. """ import math # -------------------------------------------------------------------------------------- # Hardware table. Single source of truth — never inline these numbers. # Sources: NVIDIA H100/A100 architecture whitepapers; Google Cloud TPU v5e/v5p docs. # peak_flops is DENSE bf16 (no structured sparsity). # -------------------------------------------------------------------------------------- HARDWARE = { # name peak bf16 FLOP/s HBM bytes HBM bandwidth B/s board 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), } # Bytes per parameter for the standard mixed-precision Adam recipe. # bf16 weights 2 + bf16 grads 2 + fp32 master 4 + fp32 m 4 + fp32 v 4 = 16 OPTIMIZER_STATE_BYTES_PER_PARAM = { "adam": 12, # fp32 master + m + v (weights and grads counted separately) "sgd_momentum": 8, # fp32 master + momentum "adafactor": 4, # factored second moment "sgd": 4, # fp32 master only } # ====================================================================================== # 1. The primitive # ====================================================================================== def matmul_flops(m: int, k: int, n: int) -> int: """FLOPs for an (m x k) @ (k x n) matrix multiply. Each of the m*n outputs is a sum of k products; each product-and-accumulate is 2 FLOPs (one multiply, one add). Hence 2*m*k*n. """ for name, v in (("m", m), ("k", k), ("n", n)): if v <= 0: raise ValueError(f"{name} must be positive, got {v}") return 2 * m * k * n def padded_matmul_flops(m: int, k: int, n: int, tile: int = 128) -> int: """FLOPs the hardware ACTUALLY performs, rounding each dimension up to the tile size. Systolic arrays and tensor cores operate on fixed tiles; a dimension that is not a multiple of the tile is padded, and you pay for the padding. """ if tile <= 0: raise ValueError(f"tile must be positive, got {tile}") up = lambda x: math.ceil(x / tile) * tile return matmul_flops(up(m), up(k), up(n)) # ====================================================================================== # 2. Parameter counting # ====================================================================================== def params_per_layer(d_model: int, d_ff: int, n_heads: int, n_kv_heads: int, d_head: int, gated: bool = True) -> dict: """Parameter count for one transformer block, itemized. Biases are omitted — modern LLMs drop them. `gated=True` means a SwiGLU-style MLP with THREE matrices (up, gate, down) rather than two. `n_kv_heads < n_heads` is GQA; `n_kv_heads == 1` is MQA. """ if n_kv_heads > n_heads: raise ValueError(f"n_kv_heads ({n_kv_heads}) cannot exceed n_heads ({n_heads})") for name, v in (("d_model", d_model), ("d_ff", d_ff), ("n_heads", n_heads), ("n_kv_heads", n_kv_heads), ("d_head", d_head)): if v <= 0: raise ValueError(f"{name} must be positive, got {v}") if n_heads % n_kv_heads != 0: raise ValueError(f"n_heads ({n_heads}) must be divisible by n_kv_heads ({n_kv_heads})") w_q = d_model * n_heads * d_head w_k = d_model * n_kv_heads * d_head 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 = (3 if gated else 2) * d_model * d_ff norms = 2 * d_model # pre-attn norm + pre-mlp norm return {"attn": attn, "mlp": mlp, "norms": norms, "total": attn + mlp + norms} def total_params(n_layers: int, d_model: int, d_ff: int, n_heads: int, n_kv_heads: int, d_head: int, vocab: int, tied_embeddings: bool = False, gated: bool = True) -> dict: """Whole-model parameter count, split into body vs embeddings. `non_embedding` is the number scaling-law work should use as N — see the WARMUP's "embedding trap" table for why. """ if n_layers <= 0: raise ValueError(f"n_layers must be positive, got {n_layers}") if vocab <= 0: raise ValueError(f"vocab must be positive, got {vocab}") per = params_per_layer(d_model, d_ff, n_heads, n_kv_heads, d_head, gated)["total"] body = n_layers * per final_norm = d_model embed = vocab * d_model unembed = 0 if tied_embeddings else vocab * d_model return { "body": body, "embed": embed, "unembed": unembed, "final_norm": final_norm, "non_embedding": body + final_norm, "total": body + final_norm + embed + unembed, } def moe_parameter_split(n_layers: int, d_model: int, d_ff_expert: int, n_experts: int, top_k: int, n_heads: int, n_kv_heads: int, d_head: int, shared_experts: int = 0) -> dict: """Total (memory) vs active (FLOPs) parameters for an MoE transformer. THE most commonly botched arithmetic in modern LLM work: - use `active` in C = 6ND (a token only routes through top_k experts) - use `total` for memory (all experts must live in HBM) """ if top_k > n_experts: raise ValueError(f"top_k ({top_k}) cannot exceed n_experts ({n_experts})") if top_k <= 0 or n_experts <= 0: raise ValueError("top_k and n_experts must be positive") if shared_experts < 0: raise ValueError("shared_experts cannot be negative") attn = params_per_layer(d_model, d_ff_expert, n_heads, n_kv_heads, d_head, gated=True)["attn"] one_expert = 3 * d_model * d_ff_expert router = d_model * n_experts norms = 2 * d_model total_per_layer = attn + router + norms + (n_experts + shared_experts) * one_expert active_per_layer = attn + router + norms + (top_k + shared_experts) * one_expert total = n_layers * total_per_layer active = n_layers * active_per_layer return {"total": total, "active": active, "sparsity_ratio": total / active} # ====================================================================================== # 3. FLOPs # ====================================================================================== def training_flops(n_params: int, n_tokens: int) -> float: """C = 6ND. Forward 2N + backward 4N, per token. For MoE, pass ACTIVE parameters. Excludes sequence-dependent attention matmuls. """ if n_params <= 0 or n_tokens <= 0: raise ValueError("n_params and n_tokens must be positive") return 6.0 * n_params * n_tokens def inference_flops(n_params: int, n_tokens: int) -> float: """2N per token — forward only. Used for prefill and for the serving term of lifetime-cost models.""" if n_params <= 0 or n_tokens <= 0: raise ValueError("n_params and n_tokens must be positive") return 2.0 * n_params * n_tokens def training_flops_exact(batch: int, seq_len: int, d_model: int, d_ff: int, n_heads: int, d_head: int, n_layers: int, n_kv_heads: int = None, gated: bool = True, include_attention_matmuls: bool = True, causal_halving: bool = False) -> dict: """Per-optimizer-step training FLOPs, decomposed by where they go. With n_kv_heads == n_heads and gated=True this reproduces Feinberg's slide identity 18*B*T*D*F + 24*B*T*D*N*H = 6*B*T*(3*D*F + 4*D*N*H) where D=d_model, F=d_ff, N=n_heads, H=d_head. `causal_halving` halves the sequence-dependent term to account for the causal mask. Most published accounting (including Kaplan) leaves it False; FlashAttention with a causal mask does skip those blocks. Pick a convention, state it, be consistent. """ if n_kv_heads is None: n_kv_heads = n_heads if batch <= 0 or seq_len <= 0 or n_layers <= 0: raise ValueError("batch, seq_len and n_layers must be positive") # Reuse the validation in params_per_layer. params_per_layer(d_model, d_ff, n_heads, n_kv_heads, d_head, gated) tokens = batch * seq_len n_mats = 3 if gated else 2 mlp = 6 * tokens * (n_mats * 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 * tokens * (2 * d_model * n_heads * d_head + 2 * d_model * n_kv_heads * d_head) * n_layers attn = 0 if include_attention_matmuls: # QK^T and A@V: 2 matmuls of 2*T*T*d_head per head, forward; x3 for fwd+bwd. attn = 6 * 2 * batch * n_heads * seq_len * seq_len * d_head * n_layers if causal_halving: attn //= 2 return {"mlp": mlp, "attn_proj": proj, "attn_seq": attn, "total": mlp + proj + attn} def unembedding_flops(n_tokens: int, d_model: int, vocab: int) -> float: """Training FLOPs for the output projection (d_model -> vocab). 6*d*V per token.""" if n_tokens <= 0 or d_model <= 0 or vocab <= 0: raise ValueError("all arguments must be positive") return 6.0 * n_tokens * d_model * vocab def attention_flop_fraction(n_params: int, n_layers: int, d_model: int, seq_len: int) -> float: """Fraction of training FLOPs in the sequence-dependent attention matmuls. This is how much `6ND` is missing. ~4% at 2k context; ~95% at 1M. """ if seq_len <= 0: raise ValueError("seq_len must be positive") body = 6.0 * n_params attn = 6.0 * 2 * n_layers * seq_len * d_model return attn / (body + attn) # ====================================================================================== # 4. Memory # ====================================================================================== def activation_bytes(batch: int, seq_len: int, d_model: int, n_layers: int, multiplier: int = 16, bytes_per_elem: int = 2, checkpointing: str = None) -> float: """Bytes of activations held for the backward pass. `multiplier` is how many d_model-sized tensors per layer the framework keeps alive (10-30 in practice, depending on fusion). checkpointing: None -> keep everything "selective" -> keep ~30% (recompute the cheap, memory-heavy ops) "full" -> keep only layer boundaries, recompute the rest """ if checkpointing not in (None, "selective", "full"): raise ValueError(f"unknown checkpointing mode: {checkpointing!r}") if batch <= 0 or seq_len <= 0: raise ValueError("batch and seq_len must be positive") boundary = batch * seq_len * d_model * n_layers * bytes_per_elem raw = boundary * multiplier if checkpointing == "full": return float(boundary) if checkpointing == "selective": return raw * 0.3 return float(raw) def training_memory(n_params: int, activation_bytes_: float = 0.0, optimizer: str = "adam", zero_stage: int = 0, dp_degree: int = 1, weight_bytes: int = 2, grad_bytes: int = 2) -> dict: """Per-device HBM bytes for a training step, itemized. ZeRO/FSDP stages shard the terms across `dp_degree` data-parallel replicas: stage 1 -> optimizer states stage 2 -> + gradients stage 3 -> + parameters """ if optimizer not in OPTIMIZER_STATE_BYTES_PER_PARAM: raise ValueError(f"unknown optimizer: {optimizer!r}") if zero_stage not in (0, 1, 2, 3): raise ValueError(f"zero_stage must be 0-3, got {zero_stage}") if dp_degree < 1: raise ValueError(f"dp_degree must be >= 1, got {dp_degree}") if n_params <= 0: raise ValueError("n_params must be positive") weights = weight_bytes * n_params grads = grad_bytes * n_params opt = OPTIMIZER_STATE_BYTES_PER_PARAM[optimizer] * n_params 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": activation_bytes_, "total": weights + grads + opt + activation_bytes_} def kv_cache_bytes(n_layers: int, n_kv_heads: int, d_head: int, seq_len: int, batch: int, bytes_per_elem: int = 2) -> int: """KV cache size: 2 (K and V) * L * n_kv * d_head * T * B * bytes. This is the serving wall. `n_kv_heads` is the lever, and it is frozen at pre-training time. """ for name, v in (("n_layers", n_layers), ("n_kv_heads", n_kv_heads), ("d_head", d_head), ("seq_len", seq_len), ("batch", batch)): if v <= 0: raise ValueError(f"{name} must be positive, got {v}") return 2 * n_layers * n_kv_heads * d_head * seq_len * batch * bytes_per_elem def max_concurrent_requests(total_hbm_bytes: float, n_params: int, kv_bytes_per_request: float, weight_bytes: int = 2, workspace_bytes: float = 10e9) -> int: """How many simultaneous requests fit after weights and workspace. Returns 0 if the model does not even fit. """ if kv_bytes_per_request <= 0: raise ValueError("kv_bytes_per_request must be positive") free = total_hbm_bytes - n_params * weight_bytes - workspace_bytes if free <= 0: return 0 return int(free / kv_bytes_per_request) def decode_arithmetic_intensity(n_params: int, kv_bytes: float, batch: int, weight_bytes: int = 2) -> float: """FLOPs per byte moved during decode. Weights are read once and shared across the batch; KV is per request. Compare against the accelerator's ridge point (peak_flops / hbm_bandwidth) to see whether you are compute- or memory-bound. For decode the answer is essentially always "memory". """ if batch <= 0: raise ValueError("batch must be positive") flops = 2.0 * n_params * batch bytes_moved = n_params * weight_bytes + kv_bytes return flops / bytes_moved def ridge_point(chip: str) -> float: """peak FLOP/s divided by HBM bandwidth: the arithmetic intensity at which a kernel stops being memory-bound and starts being compute-bound.""" peak, _hbm, bw, _w = _chip(chip) return peak / bw # ====================================================================================== # 5. Budgets # ====================================================================================== def _chip(chip: str): if chip not in HARDWARE: raise ValueError(f"unknown chip {chip!r}; known: {sorted(HARDWARE)}") return HARDWARE[chip] def budget_to_flops(chip: str, n_chips: int, days: float, mfu: float = 0.4) -> float: """Convert a hardware allocation into a total FLOP budget C.""" if n_chips <= 0 or days <= 0: raise ValueError("n_chips and days must be positive") if not 0 < mfu <= 1: raise ValueError(f"mfu must be in (0, 1], got {mfu}") peak, _hbm, _bw, _w = _chip(chip) return n_chips * peak * mfu * days * 86400 def flops_to_days(total_flops: float, chip: str, n_chips: int, mfu: float = 0.4) -> float: """Inverse of budget_to_flops: how long does C take on this cluster?""" peak, _hbm, _bw, _w = _chip(chip) if n_chips <= 0: raise ValueError("n_chips must be positive") return total_flops / (n_chips * peak * mfu) / 86400 def chinchilla_split(total_flops: float, tokens_per_param: float = 20.0): """Chinchilla-optimal (N, D) for a compute budget C. C = 6ND and D = r*N => C = 6r*N^2 => N = sqrt(C / (6r)). The canonical r is ~20 tokens per parameter. """ if total_flops <= 0: raise ValueError("total_flops must be positive") if tokens_per_param <= 0: raise ValueError("tokens_per_param must be positive") n_params = math.sqrt(total_flops / (6.0 * tokens_per_param)) return n_params, tokens_per_param * n_params def lifetime_flops(n_params: int, train_tokens: float, inference_tokens: float) -> float: """6ND to train + 2N per served token, forever. The objective Chinchilla ignores.""" return training_flops(n_params, train_tokens) + 2.0 * n_params * inference_tokens def cost_report(chip: str, n_chips: int, days: float, price_per_chip_hour: float, mfu: float = 0.4, kwh_price: float = 0.12, pue: float = 1.2) -> dict: """Money and energy for a hardware allocation. PUE (Power Usage Effectiveness) accounts for cooling and power delivery on top of the chips themselves; 1.1-1.5 is the normal range for a modern datacenter. """ peak, _hbm, _bw, watts = _chip(chip) if price_per_chip_hour < 0 or kwh_price < 0 or pue < 1: raise ValueError("price, kwh_price must be >= 0 and pue must be >= 1") hours = days * 24 rental = n_chips * hours * price_per_chip_hour energy_kwh = n_chips * watts * pue * hours / 1000.0 flops = budget_to_flops(chip, n_chips, days, mfu) return { "rental_dollars": rental, "energy_kwh": energy_kwh, "energy_dollars": energy_kwh * kwh_price, "flops": flops, "dollars_per_1e21_flops": rental / flops * 1e21, } def budget_report(chip: str, n_chips: int, days: float, mfu: float = 0.4, tokens_per_param: float = 20.0, available_unique_tokens: float = 15e12, price_per_chip_hour: float = 2.50) -> dict: """The full napkin: budget -> (N, D) -> sanity checks. This is the answer to "I give you 1000 H100 for 30 days; what do you train?" """ peak, hbm, _bw, _w = _chip(chip) total_flops = budget_to_flops(chip, n_chips, days, mfu) n_params, n_tokens = chinchilla_split(total_flops, tokens_per_param) mem = training_memory(n_params) cluster_hbm = n_chips * hbm costs = cost_report(chip, n_chips, days, price_per_chip_hour, mfu) checks = { # 70% of cluster HBM, leaving room for activations and fragmentation. "memory_fits_sharded": mem["total"] < cluster_hbm * 0.7, "data_available": n_tokens <= available_unique_tokens, "single_chip_serving": n_params * 2 <= hbm, } return { "total_flops": total_flops, "n_params": n_params, "n_tokens": n_tokens, "training_memory_bytes": mem["total"], "cluster_hbm_bytes": cluster_hbm, "chips_to_hold_weights_bf16": math.ceil(n_params * 2 / hbm), "checks": checks, **costs, } # ====================================================================================== # Worked example # ====================================================================================== def _gb(x): return x / 1e9 def main(): print("=" * 78) print("TRANSFORMER FLOPs / MEMORY / BUDGET CALCULATOR") print("=" * 78) print("\n[1] The primitive: a matmul costs 2*m*k*n") print(f" (1 x 4096) @ (4096 x 16384) = {matmul_flops(1, 4096, 16384):,} FLOPs") print(f" which is exactly 2 x the {4096*16384:,} weights in that layer.") print("\n[2] Parameters of one block (d=4096, d_ff=11008, 32 heads, GQA-8)") p = params_per_layer(4096, 11008, 32, 8, 128) for k, v in p.items(): print(f" {k:6s} {v:>14,}") print(f" MLP is {p['mlp']/p['total']:.1%} of the block -> this is what MoE replaces.") print("\n[3] Backward is EXACTLY 2x forward") fwd = matmul_flops(1024, 4096, 11008) bwd = matmul_flops(1024, 11008, 4096) + matmul_flops(4096, 1024, 11008) print(f" forward {fwd:,} backward {bwd:,} ratio {bwd/fwd:.1f}") print("\n[4] C = 6ND on known runs") for name, N, D in [("GPT-3", 175e9, 300e9), ("Chinchilla", 70e9, 1.4e12), ("Llama-3-70B", 70e9, 15e12), ("Llama-3-8B", 8e9, 15e12)]: print(f" {name:14s} N={N/1e9:6.1f}B D={D/1e12:5.1f}T " f"C={training_flops(N, D):.2e}") print("\n[5] Exact per-step decomposition (B=8, T=8192, 70B-class shapes)") r = training_flops_exact(batch=8, seq_len=8192, d_model=8192, d_ff=28672, n_heads=64, d_head=128, n_layers=80) for k, v in r.items(): print(f" {k:10s} {v:>22,} ({v/r['total']:5.1%})") print("\n[6] Where 6ND breaks: attention at long context") for T in (2048, 8192, 32768, 131072, 1048576): f = attention_flop_fraction(70e9, 80, 8192, T) print(f" T={T:>9,} attention {f:6.1%} of FLOPs 6ND error {f/(1-f):8.1%}") print("\n[7] Where 6ND breaks: MoE (total vs active)") m = moe_parameter_split(n_layers=60, d_model=7168, d_ff_expert=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 <- STORE this (HBM)") print(f" active {m['active']/1e9:7.1f}B <- COMPUTE this (6ND)") print(f" sparsity ratio {m['sparsity_ratio']:.1f}x") print("\n[8] Training memory for a 70B model (mixed-precision Adam)") for stage, dp in ((0, 1), (1, 64), (2, 64), (3, 64)): mm = training_memory(70e9, zero_stage=stage, dp_degree=dp) print(f" ZeRO-{stage} (dp={dp:2d}): {_gb(mm['total']):8.1f} GB/device " f"[w {_gb(mm['weights']):6.1f} | g {_gb(mm['grads']):6.1f} " f"| o {_gb(mm['optimizer']):6.1f}]") print("\n[9] The KV cache wall (80 layers, d_head=128, 8k context)") for label, kv in (("MHA (64 kv)", 64), ("GQA-8 ( 8 kv)", 8), ("MQA ( 1 kv)", 1)): per_req = kv_cache_bytes(80, kv, 128, 8192, 1) conc = max_concurrent_requests(640e9, 70e9, per_req) print(f" {label} {_gb(per_req):6.2f} GB/request " f"{conc:4d} concurrent on 8xH100") print("\n[10] Decode is memory-bound at every realistic batch size") print(f" H100 ridge point: {ridge_point('H100'):.0f} FLOP/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 " f"-> {'MEMORY' if ai < ridge_point('H100') else 'COMPUTE'} bound") print("\n[11] THE QUESTION: 1000 H100 for 30 days. What do you train?") rep = budget_report("H100", 1000, 30) print(f" C = {rep['total_flops']:.3e} FLOPs") print(f" Chinchilla N = {rep['n_params']/1e9:.1f}B parameters") print(f" Chinchilla D = {rep['n_tokens']/1e12:.2f}T tokens") print(f" training memory = {_gb(rep['training_memory_bytes']):,.0f} GB " f"(cluster has {_gb(rep['cluster_hbm_bytes']):,.0f} GB)") print(f" chips to hold weights at serve time: " f"{rep['chips_to_hold_weights_bf16']}") print(f" rental = ${rep['rental_dollars']:,.0f}") print(f" energy = {rep['energy_kwh']/1000:,.1f} MWh " f"(${rep['energy_dollars']:,.0f})") print(f" $ per 1e21 FLOPs = ${rep['dollars_per_1e21_flops']:,.2f}") for k, v in rep["checks"].items(): print(f" [{'OK ' if v else 'FAIL'}] {k}") print("\n[12] The senior move: deliberately undershoot Chinchilla for serving") C = rep["total_flops"] for N in (rep["n_params"], 50e9, 30e9, 15e9): D = C / (6 * N) life_1e15 = lifetime_flops(N, D, 1e15) print(f" N={N/1e9:6.1f}B D={D/1e12:6.2f}T " f"lifetime@1e15 served = {life_1e15:.3e} FLOPs " f"({math.ceil(N*2/80e9)} H100 to serve)") print("\n Smaller model, more tokens, same training budget, far cheaper to serve.") print(" Quantifying that trade is Phase 02.") print("=" * 78) if __name__ == "__main__": main()