"""Transformer FLOPs, Memory & Budget Calculator — YOUR implementation. Fill in every `# TODO`. Signatures, docstrings and validation contracts are already here; the arithmetic is yours. Run: pytest test_lab.py -v # red until you implement LAB_MODULE=solution pytest test_lab.py -v # the reference (must be green) python solution.py # the worked example ORDER OF WORK: do `matmul_flops` first and get the "backward is 2x forward" test passing. Everything else in this file is bookkeeping on top of that one primitive. UNITS: bytes are bytes, FLOPs are FLOPs, seconds are seconds. Most errors at this level are unit errors, not algebra errors. Say the unit out loud as you write each line. """ import math # -------------------------------------------------------------------------------------- # Hardware table — given to you. Single source of truth; never inline these numbers. # 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), } 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). Raises ValueError if any dimension is <= 0. """ # TODO: validate m, k, n are positive; raise ValueError naming the offending arg. # TODO: return 2 * m * k * n raise NotImplementedError def padded_matmul_flops(m: int, k: int, n: int, tile: int = 128) -> int: """FLOPs the hardware ACTUALLY performs, rounding each dimension up to `tile`. 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. """ # TODO: validate tile > 0. # TODO: round m, k, n UP to the next multiple of `tile` (math.ceil), then call # matmul_flops on the padded dimensions. raise NotImplementedError # ====================================================================================== # 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. Returns {"attn", "mlp", "norms", "total"}. Biases omitted (modern LLMs drop them). `gated=True` = SwiGLU-style MLP with THREE matrices (up, gate, down), else two. `n_kv_heads < n_heads` is GQA; `n_kv_heads == 1` is MQA. Raises ValueError if n_kv_heads > n_heads, if n_heads is not divisible by n_kv_heads, or if any dimension is <= 0. """ # TODO: validate. n_kv_heads must be <= n_heads AND divide it evenly. # TODO: 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 # TODO: mlp = (3 if gated else 2) * d_model * d_ff # TODO: norms = 2 * d_model (one before attention, one before the MLP) raise NotImplementedError 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. Returns {"body", "embed", "unembed", "final_norm", "non_embedding", "total"}. `non_embedding` = body + final_norm. THIS is the N that scaling-law work should use — see the WARMUP's "embedding trap" table for why (a 256k vocab can be 89% of a small model). """ # TODO: validate n_layers and vocab are positive. # TODO: body = n_layers * params_per_layer(...)["total"] # TODO: embed = vocab * d_model # unembed = 0 if tied_embeddings else vocab * d_model # final_norm = d_model raise NotImplementedError 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. Returns {"total", "active", "sparsity_ratio"}. THE most commonly botched arithmetic in modern LLM work: - `active` goes into C = 6ND (a token routes through only top_k experts) - `total` goes into memory (all experts must live in HBM) Per layer, BOTH counts include: attention projections + router + norms. They differ only in how many experts you count: total -> (n_experts + shared_experts) experts active -> (top_k + shared_experts) experts One expert has 3 * d_model * d_ff_expert parameters (gated MLP). The router has d_model * n_experts parameters. Raises ValueError if top_k > n_experts, if either is <= 0, or shared_experts < 0. """ # TODO: validate. # TODO: reuse params_per_layer(...)["attn"] for the attention term. # TODO: build total_per_layer and active_per_layer, multiply by n_layers. raise NotImplementedError # ====================================================================================== # 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. """ # TODO: validate positives; return 6.0 * n_params * n_tokens raise NotImplementedError def inference_flops(n_params: int, n_tokens: int) -> float: """2N per token — forward only. Prefill, and the serving term of lifetime cost.""" # TODO: validate positives; return 2.0 * n_params * n_tokens raise NotImplementedError 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. Returns {"mlp", "attn_proj", "attn_seq", "total"}. With n_kv_heads == n_heads and gated=True this must reproduce the 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. A test checks exactly this. Hints: tokens = batch * seq_len mlp = 6 * tokens * (n_mats * d_model * d_ff) * n_layers attn_proj = 6 * tokens * (2*d_model*n_heads*d_head + 2*d_model*n_kv_heads*d_head) * n_layers attn_seq = 6 * 2 * batch * n_heads * seq_len**2 * d_head * n_layers (QK^T and A@V are two matmuls of 2*T*T*d_head per head; x3 fwd+bwd) `causal_halving` halves attn_seq (integer division). If n_kv_heads is None, default it to n_heads. """ # TODO: implement. raise NotImplementedError 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.""" # TODO: validate positives; return 6.0 * n_tokens * d_model * vocab raise NotImplementedError 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. body = 6 * n_params attn = 6 * 2 * n_layers * seq_len * d_model return attn / (body + attn) """ # TODO: implement. raise NotImplementedError # ====================================================================================== # 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. boundary = batch * seq_len * d_model * n_layers * bytes_per_elem raw = boundary * multiplier checkpointing: None -> raw "selective" -> raw * 0.3 "full" -> boundary (keep only layer boundaries, recompute the rest) Raises ValueError on an unknown checkpointing mode. """ # TODO: implement. raise NotImplementedError 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. Returns {"weights", "grads", "optimizer", "activations", "total"}. ZeRO/FSDP shards across `dp_degree` data-parallel replicas: stage >= 1 -> optimizer states /= dp_degree stage >= 2 -> + gradients /= dp_degree stage >= 3 -> + parameters /= dp_degree Raises ValueError on an unknown optimizer, zero_stage outside 0-3, or dp_degree < 1. """ # TODO: implement. Use OPTIMIZER_STATE_BYTES_PER_PARAM for the optimizer term. raise NotImplementedError 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, and n_kv_heads is the lever — frozen at pre-training time. """ # TODO: validate all positive; return the product above. raise NotImplementedError 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. free = total_hbm - n_params*weight_bytes - workspace Return 0 if free <= 0 (the model does not fit at all), else int(free / kv_per_req). """ # TODO: implement. raise NotImplementedError 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. flops = 2 * n_params * batch bytes = n_params * weight_bytes + kv_bytes Compare against ridge_point(chip) to see if you are memory- or compute-bound. """ # TODO: implement. raise NotImplementedError def ridge_point(chip: str) -> float: """peak FLOP/s / HBM bandwidth — the arithmetic intensity at which a kernel stops being memory-bound and becomes compute-bound.""" # TODO: implement using _chip(chip). raise NotImplementedError # ====================================================================================== # 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. C = n_chips * peak_flops * mfu * days * 86400 Raises ValueError if n_chips or days <= 0, or mfu not in (0, 1]. """ # TODO: implement. raise NotImplementedError 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?""" # TODO: implement. raise NotImplementedError 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)). Returns the tuple (n_params, n_tokens). """ # TODO: implement. raise NotImplementedError def lifetime_flops(n_params: int, train_tokens: float, inference_tokens: float) -> float: """6ND to train + 2N per served token, forever. This is the objective Chinchilla ignores, and Phase 02 optimizes. """ # TODO: implement. raise NotImplementedError 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. Returns {"rental_dollars", "energy_kwh", "energy_dollars", "flops", "dollars_per_1e21_flops"}. PUE (Power Usage Effectiveness) accounts for cooling and power delivery on top of the chips; 1.1-1.5 is normal. Raises ValueError if pue < 1 or any price < 0. `dollars_per_1e21_flops` is the number to carry in your head: it converts any research proposal into dollars instantly. """ # TODO: implement. raise NotImplementedError 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?" Must include a "checks" dict with these boolean keys: "memory_fits_sharded" : training_memory(N)["total"] < cluster_hbm * 0.7 "data_available" : n_tokens <= available_unique_tokens "single_chip_serving" : n_params * 2 <= per-chip HBM and the keys "total_flops", "n_params", "n_tokens", "training_memory_bytes", "cluster_hbm_bytes", "chips_to_hold_weights_bf16", plus everything cost_report returns. """ # TODO: implement. raise NotImplementedError if __name__ == "__main__": print("Fill in the TODOs, then compare against: python solution.py")