add config/param_count.py (exact param accounting, 90-110M window)

config/param_count.py

@@ -0,0 +1,115 @@
+"""Exact parameter accounting for a Llama-style decoder-only model.
+
+The parameter budget (§2: 90-110M *including* embeddings) cannot be met by eyeballing hidden size
+and layer count: at a 50k vocab the embedding alone is 25-50% of the model, and whether the
+embeddings are tied changes the total by the whole embedding matrix. So count it arithmetically,
+then confirm against the real object.
+
+Two independent counts are printed deliberately:
+  formula  -- closed form, so the number is explainable without instantiating the model
+  torch    -- sum of p.numel() over a constructed LlamaForCausalLM, which is the number a third
+              party can reproduce from the published config
+
+They must agree. If they do not, the formula is wrong (usually GQA head packing or tying), and the
+torch count is authoritative because that is what `AutoModelForCausalLM` will report.
+
+Run:  python code/config/param_count.py            # candidate table
+      python code/config/param_count.py --verify   # also build each model (needs transformers)
+"""
+
+import argparse
+import json
+
+# Measured on Kaggle (docs/00-platform-notes.md section 9): this shape is 112.9M, i.e. OVER budget.
+PROBE_SHAPE = dict(hidden_size=768, num_hidden_layers=12, num_attention_heads=12,
+                   num_key_value_heads=3, intermediate_size=2048, vocab_size=50257,
+                   tie_word_embeddings=True)
+
+CANDIDATES = {
+    # Shrinking the probe shape toward the window from several directions, so the choice is between
+    # real options rather than a single lucky guess.
+    "A_768_h_11l_gqa3_2048": dict(PROBE_SHAPE, num_hidden_layers=11),
+    "B_768_h_12l_gqa3_1792": dict(PROBE_SHAPE, intermediate_size=1792),
+    "C_768_h_12l_gqa2_1824": dict(PROBE_SHAPE, intermediate_size=1824, num_key_value_heads=2),
+    "D_640_h_16l_gqa4_2048": dict(PROBE_SHAPE, hidden_size=640, num_hidden_layers=16,
+                                   num_attention_heads=10, num_key_value_heads=4,
+                                   intermediate_size=2048),
+    "E_768_h_12l_gqa3_32k": dict(PROBE_SHAPE, vocab_size=32768),
+    "F_896_h_9l_gqa2_2432": dict(PROBE_SHAPE, hidden_size=896, num_hidden_layers=9,
+                                  num_attention_heads=14, num_key_value_heads=2,
+                                  intermediate_size=2432),
+    "G_1024_h_8l_gqa2_2816": dict(PROBE_SHAPE, hidden_size=1024, num_hidden_layers=8,
+                                  num_attention_heads=8, num_key_value_heads=2,
+                                  intermediate_size=2816),
+}
+
+
+def count_formula(v, h, nl, nh, nkv, ff, tied):
+    """Closed-form Llama param count. All terms explained inline so the arithmetic is checkable."""
+    embed = v * h
+    d_head = h // nh
+    # GQA: k/v are d_head * nkv wide, not h, so the ratio nkv/nh matters.
+    attn = h * h + 2 * d_head * nkv * h + h * h            # q, k, v, o
+    mlp = 3 * h * ff                                        # SwiGLU gate + up + down
+    norms = 2 * h * nl                                      # input + post-attention RMSNorm
+    per_layer = attn + mlp
+    total = embed * (1 if tied else 2) + per_layer * nl + norms + h
+    return {
+        "embedding_one_copy": embed,
+        "per_layer_attn": attn, "per_layer_mlp": mlp, "layers": nl,
+        "transformer_blocks": per_layer * nl,
+        "norms": norms + h,
+        "tied_embeddings": tied,
+        "total": total,
+        "embedding_pct_of_total": round(100.0 * embed / total, 1),
+    }
+
+
+def count_torch(cfg):
+    import torch
+    from transformers import LlamaConfig, LlamaForCausalLM
+
+    m = LlamaForCausalLM(LlamaConfig(**cfg))
+    n = sum(p.numel() for p in m.parameters())
+    emb = m.get_input_embeddings().weight.numel()
+    tied = m.config.tie_word_embeddings and (
+        m.get_input_embeddings().weight.data_ptr() == m.lm_head.weight.data_ptr())
+    del m
+    torch.cuda.empty_cache() if torch.cuda.is_available() else None
+    return {"total": n, "embedding": emb, "actually_tied": bool(tied)}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--verify", action="store_true",
+                    help="build each model with transformers and compare (slow, needs the lib)")
+    args = ap.parse_args()
+
+    rows = []
+    for name, cfg in CANDIDATES.items():
+        f = count_formula(cfg["vocab_size"], cfg["hidden_size"], cfg["num_hidden_layers"],
+                          cfg["num_attention_heads"], cfg["num_key_value_heads"],
+                          cfg["intermediate_size"], cfg["tie_word_embeddings"])
+        row = {"name": name, "formula_total": f["total"], "emb_pct": f["embedding_pct_of_total"],
+               "in_90_110M": 90_000_000 <= f["total"] <= 110_000_000, **{
+                   k: cfg[k] for k in ("hidden_size", "num_hidden_layers", "num_attention_heads",
+                                        "num_key_value_heads", "intermediate_size", "vocab_size")}}
+        if args.verify:
+            t = count_torch(cfg)
+            row["torch_total"] = t["total"]
+            row["matches_formula"] = t["total"] == f["total"]
+            row["torch_reports_tied"] = t["actually_tied"]
+        rows.append(row)
+
+    p = count_formula(**{  # the probe shape, for the record
+        "v": PROBE_SHAPE["vocab_size"], "h": PROBE_SHAPE["hidden_size"],
+        "nl": PROBE_SHAPE["num_hidden_layers"], "nh": PROBE_SHAPE["num_attention_heads"],
+        "nkv": PROBE_SHAPE["num_key_value_heads"], "ff": PROBE_SHAPE["intermediate_size"],
+        "tied": PROBE_SHAPE["tie_word_embeddings"]})
+    print("probe shape (measured 112,934,400 in the image):", json.dumps(p))
+    for r in rows:
+        print(json.dumps(r))
+
+
+if __name__ == "__main__":
+    main()