"""Tests for the FLOPs / memory / budget calculator. pytest test_lab.py -v # your lab.py LAB_MODULE=solution pytest test_lab.py -v # the reference (must be green) """ import importlib import math import os import pytest lab = importlib.import_module(os.environ.get("LAB_MODULE", "lab")) # ====================================================================================== # 1. The matmul primitive # ====================================================================================== def test_matmul_flops_basic(): assert lab.matmul_flops(1, 1, 1) == 2 # one multiply, one add assert lab.matmul_flops(2, 3, 4) == 2 * 2 * 3 * 4 assert lab.matmul_flops(1, 4096, 16384) == 2 * 4096 * 16384 def test_matmul_flops_is_2x_weight_count(): """A linear layer costs 2 x (its parameter count) FLOPs per token. This single identity is the root of C = 6ND.""" d_in, d_out = 4096, 11008 weights = d_in * d_out assert lab.matmul_flops(1, d_in, d_out) == 2 * weights @pytest.mark.parametrize("m,k,n", [(0, 4, 4), (4, 0, 4), (4, 4, 0), (-1, 4, 4)]) def test_matmul_flops_rejects_nonpositive(m, k, n): with pytest.raises(ValueError): lab.matmul_flops(m, k, n) def test_padded_matmul_never_less_than_useful(): for m, k, n in [(1000, 4000, 4000), (1, 1, 1), (128, 128, 128)]: assert lab.padded_matmul_flops(m, k, n) >= lab.matmul_flops(m, k, n) def test_padded_matmul_exact_on_tile_multiples(): """Dimensions already on the tile boundary must incur zero padding.""" assert lab.padded_matmul_flops(128, 4096, 8192) == lab.matmul_flops(128, 4096, 8192) def test_padded_matmul_rejects_bad_tile(): with pytest.raises(ValueError): lab.padded_matmul_flops(4, 4, 4, tile=0) # ====================================================================================== # 2. Parameter counting # ====================================================================================== def test_params_per_layer_known_values(): p = lab.params_per_layer(d_model=4096, d_ff=11008, n_heads=32, n_kv_heads=8, d_head=128, gated=True) assert p["attn"] == 41_943_040 assert p["mlp"] == 135_266_304 assert p["norms"] == 8_192 assert p["total"] == p["attn"] + p["mlp"] + p["norms"] def test_gated_mlp_is_1_5x_ungated(): """SwiGLU needs three matrices, not two.""" kw = dict(d_model=1024, d_ff=4096, n_heads=8, n_kv_heads=8, d_head=128) assert (lab.params_per_layer(gated=True, **kw)["mlp"] == 1.5 * lab.params_per_layer(gated=False, **kw)["mlp"]) def test_gqa_shrinks_only_kv_projections(): """GQA reduces W_k and W_v; W_q and W_o are untouched.""" mha = lab.params_per_layer(4096, 11008, 32, 32, 128)["attn"] gqa = lab.params_per_layer(4096, 11008, 32, 8, 128)["attn"] mqa = lab.params_per_layer(4096, 11008, 32, 1, 128)["attn"] assert mha > gqa > mqa # W_q + W_o are half of MHA attention and are unchanged. q_plus_o = 2 * 4096 * 32 * 128 assert gqa == q_plus_o + 2 * 4096 * 8 * 128 assert mqa == q_plus_o + 2 * 4096 * 1 * 128 def test_params_per_layer_rejects_bad_gqa(): with pytest.raises(ValueError): lab.params_per_layer(4096, 11008, n_heads=8, n_kv_heads=32, d_head=128) with pytest.raises(ValueError): lab.params_per_layer(4096, 11008, n_heads=32, n_kv_heads=7, d_head=128) def test_total_params_llama2_7b_shape(): m = lab.total_params(n_layers=32, d_model=4096, d_ff=11008, n_heads=32, n_kv_heads=32, d_head=128, vocab=32000) assert 6.7e9 < m["total"] < 6.8e9 assert m["non_embedding"] == m["body"] + m["final_norm"] assert m["total"] == m["non_embedding"] + m["embed"] + m["unembed"] def test_tied_embeddings_removes_the_unembedding(): kw = dict(n_layers=8, d_model=512, d_ff=2048, n_heads=4, n_kv_heads=4, d_head=128, vocab=256000) untied = lab.total_params(tied_embeddings=False, **kw) tied = lab.total_params(tied_embeddings=True, **kw) assert tied["unembed"] == 0 assert untied["total"] - tied["total"] == 256000 * 512 def test_embedding_trap_small_models_are_mostly_embeddings(): """The reason scaling work reports NON-EMBEDDING parameters.""" tiny = lab.total_params(8, 512, 2048, 4, 4, 128, vocab=256000) frac = (tiny["embed"] + tiny["unembed"]) / tiny["total"] assert frac > 0.85 # ---- MoE ----------------------------------------------------------------------------- def test_moe_total_exceeds_active(): m = lab.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) assert m["total"] > m["active"] assert m["sparsity_ratio"] == pytest.approx(m["total"] / m["active"]) assert m["sparsity_ratio"] > 10 def test_moe_degenerates_to_dense_when_all_experts_active(): """BOUNDARY: top_k == n_experts means every token uses every expert.""" kw = dict(n_layers=4, d_model=512, d_ff_expert=1024, n_experts=8, n_heads=8, n_kv_heads=8, d_head=64) m = lab.moe_parameter_split(top_k=8, **kw) assert m["total"] == m["active"] assert m["sparsity_ratio"] == pytest.approx(1.0) def test_moe_shared_experts_count_in_both(): """A shared expert runs for EVERY token, so it lands in total AND active.""" kw = dict(n_layers=2, d_model=256, d_ff_expert=512, n_experts=8, top_k=2, n_heads=4, n_kv_heads=4, d_head=64) a = lab.moe_parameter_split(shared_experts=0, **kw) b = lab.moe_parameter_split(shared_experts=1, **kw) one_expert = 3 * 256 * 512 assert b["total"] - a["total"] == 2 * one_expert assert b["active"] - a["active"] == 2 * one_expert def test_moe_rejects_topk_over_n_experts(): with pytest.raises(ValueError): lab.moe_parameter_split(4, 512, 1024, n_experts=8, top_k=9, n_heads=8, n_kv_heads=8, d_head=64) # ====================================================================================== # 3. FLOPs # ====================================================================================== def test_6nd(): assert lab.training_flops(1e9, 1e12) == 6e21 assert lab.inference_flops(1e9, 1e12) == 2e21 def test_training_is_3x_inference(): """INVARIANT: forward 2N, backward 4N. Training is exactly 3x a forward pass.""" N, D = 7e9, 1e11 assert lab.training_flops(N, D) == pytest.approx(3 * lab.inference_flops(N, D)) def test_flops_scale_linearly_in_both_arguments(): """DIMENSIONAL ANALYSIS: doubling either N or D doubles C.""" base = lab.training_flops(1e9, 1e12) assert lab.training_flops(2e9, 1e12) == pytest.approx(2 * base) assert lab.training_flops(1e9, 2e12) == pytest.approx(2 * base) @pytest.mark.parametrize("n,d", [(0, 1e12), (1e9, 0), (-1, 1e12)]) def test_training_flops_rejects_nonpositive(n, d): with pytest.raises(ValueError): lab.training_flops(n, d) def test_exact_step_reproduces_the_slide_identity(): """18*B*T*D*F + 24*B*T*D*N*H, for MHA + gated MLP. This is Feinberg's Princeton slide, written out. If this passes you have the per-shape accounting right. """ B, T, D, F, N, H, L = 4, 2048, 4096, 11008, 32, 128, 32 r = lab.training_flops_exact(batch=B, seq_len=T, d_model=D, d_ff=F, n_heads=N, d_head=H, n_layers=L, include_attention_matmuls=False) assert r["mlp"] == 18 * B * T * D * F * L assert r["attn_proj"] == 24 * B * T * D * N * H * L # ... and the factored form on the slide. assert r["total"] == 6 * B * T * (3 * D * F + 4 * D * N * H) * L def test_exact_step_total_is_sum_of_parts(): r = lab.training_flops_exact(2, 512, 512, 2048, 8, 64, 4) assert r["total"] == r["mlp"] + r["attn_proj"] + r["attn_seq"] def test_attention_term_is_quadratic_in_sequence_length(): """Doubling T quadruples the sequence-dependent attention FLOPs.""" kw = dict(batch=1, d_model=512, d_ff=2048, n_heads=8, d_head=64, n_layers=2) a = lab.training_flops_exact(seq_len=128, **kw)["attn_seq"] b = lab.training_flops_exact(seq_len=256, **kw)["attn_seq"] assert b == pytest.approx(4 * a) def test_causal_halving_halves_only_the_sequence_term(): kw = dict(batch=1, seq_len=256, d_model=512, d_ff=2048, n_heads=8, d_head=64, n_layers=2) full = lab.training_flops_exact(causal_halving=False, **kw) half = lab.training_flops_exact(causal_halving=True, **kw) assert half["attn_seq"] == full["attn_seq"] // 2 assert half["mlp"] == full["mlp"] assert half["attn_proj"] == full["attn_proj"] def test_exclude_attention_matmuls_zeroes_that_term(): kw = dict(batch=1, seq_len=256, d_model=512, d_ff=2048, n_heads=8, d_head=64, n_layers=2) assert lab.training_flops_exact(include_attention_matmuls=False, **kw)["attn_seq"] == 0 def test_gqa_reduces_projection_flops(): kw = dict(batch=1, seq_len=256, d_model=512, d_ff=2048, n_heads=8, d_head=64, n_layers=2) mha = lab.training_flops_exact(n_kv_heads=8, **kw)["attn_proj"] gqa = lab.training_flops_exact(n_kv_heads=2, **kw)["attn_proj"] assert gqa < mha def test_attention_fraction_grows_with_context_and_is_bounded(): fracs = [lab.attention_flop_fraction(70e9, 80, 8192, T) for T in (2048, 8192, 32768, 131072, 1048576)] assert fracs == sorted(fracs) # monotone increasing assert all(0.0 < f < 1.0 for f in fracs) # it is a fraction assert fracs[0] < 0.10 # negligible at 2k assert fracs[-1] > 0.90 # dominant at 1M def test_attention_fraction_rejects_bad_seq_len(): with pytest.raises(ValueError): lab.attention_flop_fraction(70e9, 80, 8192, 0) # ====================================================================================== # 4. Memory # ====================================================================================== def test_adam_mixed_precision_is_16_bytes_per_param(): """2 (bf16 w) + 2 (bf16 g) + 12 (fp32 master + m + v) = 16.""" m = lab.training_memory(1e9) assert m["total"] == pytest.approx(16e9) def test_optimizer_choice_changes_only_the_optimizer_term(): adam = lab.training_memory(1e9, optimizer="adam") ada = lab.training_memory(1e9, optimizer="adafactor") assert adam["weights"] == ada["weights"] assert adam["grads"] == ada["grads"] assert adam["optimizer"] > ada["optimizer"] def test_zero_stages_shard_progressively_more(): totals = [lab.training_memory(70e9, zero_stage=s, dp_degree=64)["total"] for s in (0, 1, 2, 3)] assert totals == sorted(totals, reverse=True) # strictly decreasing def test_zero_with_dp_degree_1_is_a_noop(): """BOUNDARY: sharding across one replica shards nothing.""" a = lab.training_memory(1e9, zero_stage=0, dp_degree=1) b = lab.training_memory(1e9, zero_stage=3, dp_degree=1) assert a["total"] == pytest.approx(b["total"]) @pytest.mark.parametrize("kwargs", [ {"optimizer": "nonexistent"}, {"zero_stage": 4}, {"zero_stage": -1}, {"dp_degree": 0}, ]) def test_training_memory_validates(kwargs): with pytest.raises(ValueError): lab.training_memory(1e9, **kwargs) def test_checkpointing_reduces_activation_memory(): kw = dict(batch=8, seq_len=8192, d_model=8192, n_layers=80) none = lab.activation_bytes(checkpointing=None, **kw) sel = lab.activation_bytes(checkpointing="selective", **kw) full = lab.activation_bytes(checkpointing="full", **kw) assert none > sel > full assert full == pytest.approx(none / 16) # default multiplier is 16 def test_activation_bytes_rejects_unknown_mode(): with pytest.raises(ValueError): lab.activation_bytes(1, 1, 1, 1, checkpointing="magic") # ---- KV cache ------------------------------------------------------------------------ def test_kv_cache_known_value(): """80 layers, 8 kv heads, d_head 128, 8k context, batch 1, bf16.""" b = lab.kv_cache_bytes(80, 8, 128, 8192, 1) assert b == 2 * 80 * 8 * 128 * 8192 * 1 * 2 def test_kv_cache_linear_in_every_dimension(): base = lab.kv_cache_bytes(80, 8, 128, 8192, 1) assert lab.kv_cache_bytes(160, 8, 128, 8192, 1) == 2 * base assert lab.kv_cache_bytes(80, 16, 128, 8192, 1) == 2 * base assert lab.kv_cache_bytes(80, 8, 128, 16384, 1) == 2 * base assert lab.kv_cache_bytes(80, 8, 128, 8192, 2) == 2 * base def test_gqa_group_factor_is_the_kv_saving(): """THE inference co-design lever: 64 kv heads -> 8 is exactly 8x less cache.""" mha = lab.kv_cache_bytes(80, 64, 128, 8192, 32) gqa = lab.kv_cache_bytes(80, 8, 128, 8192, 32) assert mha == 8 * gqa def test_kv_cache_rejects_nonpositive(): with pytest.raises(ValueError): lab.kv_cache_bytes(80, 8, 128, 0, 1) def test_max_concurrent_zero_when_model_does_not_fit(): """BOUNDARY: a 70B model in bf16 is 140 GB and cannot fit on one 80 GB H100.""" assert lab.max_concurrent_requests(80e9, 70e9, 1e9) == 0 def test_max_concurrent_increases_when_kv_shrinks(): mha = lab.kv_cache_bytes(80, 64, 128, 8192, 1) gqa = lab.kv_cache_bytes(80, 8, 128, 8192, 1) assert (lab.max_concurrent_requests(640e9, 70e9, gqa) > lab.max_concurrent_requests(640e9, 70e9, mha)) def test_decode_is_memory_bound_at_realistic_batch_sizes(): """The single fact that explains batching, GQA, quantization and speculation.""" ridge = lab.ridge_point("H100") for B in (1, 8, 64, 256): kv = lab.kv_cache_bytes(80, 8, 128, 8192, B) assert lab.decode_arithmetic_intensity(70e9, kv, B) < ridge def test_arithmetic_intensity_improves_with_batch(): """Batching amortizes the weight read — which is WHY batching works.""" ais = [] for B in (1, 8, 64, 256): kv = lab.kv_cache_bytes(80, 8, 128, 8192, B) ais.append(lab.decode_arithmetic_intensity(70e9, kv, B)) assert ais == sorted(ais) def test_ridge_point_matches_the_spec_sheet(): peak, _hbm, bw, _w = lab.HARDWARE["H100"] assert lab.ridge_point("H100") == pytest.approx(peak / bw) def test_ridge_point_rejects_unknown_chip(): with pytest.raises(ValueError): lab.ridge_point("Xeon Phi") # ====================================================================================== # 5. Budgets # ====================================================================================== def test_budget_flops_roundtrip(): C = lab.budget_to_flops("H100", 1000, 30, mfu=0.4) assert lab.flops_to_days(C, "H100", 1000, mfu=0.4) == pytest.approx(30) def test_budget_scales_linearly_in_chips_days_and_mfu(): base = lab.budget_to_flops("H100", 100, 10, 0.4) assert lab.budget_to_flops("H100", 200, 10, 0.4) == pytest.approx(2 * base) assert lab.budget_to_flops("H100", 100, 20, 0.4) == pytest.approx(2 * base) assert lab.budget_to_flops("H100", 100, 10, 0.8) == pytest.approx(2 * base) @pytest.mark.parametrize("kwargs", [ {"n_chips": 0}, {"days": 0}, {"mfu": 0}, {"mfu": 1.5}, {"mfu": -0.1}, ]) def test_budget_to_flops_validates(kwargs): args = dict(chip="H100", n_chips=100, days=10, mfu=0.4) args.update(kwargs) with pytest.raises(ValueError): lab.budget_to_flops(**args) def test_budget_to_flops_rejects_unknown_chip(): with pytest.raises(ValueError): lab.budget_to_flops("Babbage Engine", 10, 1) def test_chinchilla_split_satisfies_6nd(): """INVARIANT: the returned (N, D) must consume exactly the budget.""" C = 1e24 N, D = lab.chinchilla_split(C) assert 6 * N * D == pytest.approx(C, rel=1e-9) def test_chinchilla_split_honours_the_token_ratio(): N, D = lab.chinchilla_split(1e24, tokens_per_param=20.0) assert D / N == pytest.approx(20.0) N2, D2 = lab.chinchilla_split(1e24, tokens_per_param=100.0) assert D2 / N2 == pytest.approx(100.0) assert N2 < N # more tokens per param => a smaller model def test_chinchilla_split_known_answer(): """1000 H100 x 30 days at 40% MFU is ~92B params on ~1.85T tokens.""" C = lab.budget_to_flops("H100", 1000, 30, 0.4) N, D = lab.chinchilla_split(C) assert 90e9 < N < 95e9 assert 1.8e12 < D < 1.9e12 def test_chinchilla_split_validates(): with pytest.raises(ValueError): lab.chinchilla_split(0) with pytest.raises(ValueError): lab.chinchilla_split(1e24, tokens_per_param=0) def test_lifetime_flops_crossover(): """The whole point of inference-aware scaling: an overtrained SMALL model wins once you serve enough tokens.""" big = (70e9, 1.4e12) # Chinchilla-optimal small = (20e9, 8.0e12) # deliberately overtrained assert lab.lifetime_flops(*big, 1e12) < lab.lifetime_flops(*small, 1e12) assert lab.lifetime_flops(*big, 1e14) > lab.lifetime_flops(*small, 1e14) def test_lifetime_flops_is_train_plus_serve(): assert lab.lifetime_flops(1e9, 1e12, 1e12) == pytest.approx(6e21 + 2e21) def test_cost_report_dollars_per_flop_consistent(): r = lab.cost_report("H100", 1000, 30, price_per_chip_hour=2.50) assert (r["dollars_per_1e21_flops"] == pytest.approx(r["rental_dollars"] / r["flops"] * 1e21)) assert r["energy_dollars"] == pytest.approx(r["energy_kwh"] * 0.12) def test_cost_report_pue_multiplies_energy(): a = lab.cost_report("H100", 10, 1, 2.5, pue=1.0) b = lab.cost_report("H100", 10, 1, 2.5, pue=2.0) assert b["energy_kwh"] == pytest.approx(2 * a["energy_kwh"]) assert b["rental_dollars"] == pytest.approx(a["rental_dollars"]) def test_cost_report_rejects_pue_below_one(): with pytest.raises(ValueError): lab.cost_report("H100", 10, 1, 2.5, pue=0.5) def test_budget_report_is_internally_consistent(): r = lab.budget_report("H100", 1000, 30) assert 6 * r["n_params"] * r["n_tokens"] == pytest.approx(r["total_flops"], rel=1e-9) assert set(r["checks"]) == {"memory_fits_sharded", "data_available", "single_chip_serving"} assert r["chips_to_hold_weights_bf16"] == math.ceil(r["n_params"] * 2 / 80e9) def test_budget_report_flags_impossible_data_requirement(): """A huge budget on a small corpus must fail the data check.""" r = lab.budget_report("H100", 100000, 365, available_unique_tokens=1e12) assert r["checks"]["data_available"] is False def test_budget_report_flags_unservable_model(): """92B params in bf16 is 185 GB — it does not fit on one 80 GB H100.""" r = lab.budget_report("H100", 1000, 30) assert r["checks"]["single_chip_serving"] is False def test_budget_report_is_deterministic(): """DETERMINISM: same inputs, byte-identical output. No clocks, no randomness.""" a = lab.budget_report("H100", 1000, 30) b = lab.budget_report("H100", 1000, 30) assert repr(a) == repr(b)