Upload params.py

params.py

@@ -0,0 +1,165 @@
+"""Exact parameter accounting.
+
+The headline number is trainable weights PLUS frozen buffer values.
+"""
+from __future__ import annotations
+
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+# Buffers that hold only derived constants would be excluded and listed here.
+# The design deliberately has none; RoPE tables are plain Python state.
+DERIVED_BUFFERS: tuple[str, ...] = ()
+
+
+def _tied_targets(model: nn.Module) -> set[str]:
+    """Parameter names that are copies of another parameter, so count once.
+
+    Object identity is not enough: `from_pretrained` constructs the model before
+    tying the weights, so during construction `lm_head.weight` is still a
+    distinct tensor. Counting it would double the embedding and trip the budget
+    check on an otherwise valid model.
+    """
+    if not getattr(model.config, "tie_word_embeddings", False):
+        return set()
+    keys = getattr(model, "_tied_weights_keys", None) or []
+    # transformers 4.x: a list of target names. 5.x: {target: source}.
+    return set(keys.keys()) if isinstance(keys, dict) else set(keys)
+
+
+def account(model: nn.Module) -> dict:
+    """Walk named_parameters() and named_buffers(). Returns a count breakdown."""
+    trainable, frozen, excluded = 0, 0, 0
+    by_param: "OrderedDict[str, int]" = OrderedDict()
+    by_buffer: "OrderedDict[str, int]" = OrderedDict()
+    excluded_names: list[str] = []
+
+    tied = _tied_targets(model)
+    seen: set[int] = set()
+    for name, p in model.named_parameters():
+        if id(p) in seen:          # tied weights counted once
+            continue
+        if name in tied or any(name.endswith("." + t) for t in tied):
+            continue
+        seen.add(id(p))
+        n = p.numel()
+        by_param[name] = n
+        trainable += n
+
+    for name, b in model.named_buffers():
+        if b is None or id(b) in seen:
+            continue
+        seen.add(id(b))
+        n = b.numel()
+        if any(name.endswith(suffix) for suffix in DERIVED_BUFFERS):
+            excluded += n
+            excluded_names.append(name)
+            continue
+        by_buffer[name] = n
+        frozen += n
+
+    return {
+        "trainable": trainable,
+        "frozen": frozen,
+        "total": trainable + frozen,
+        "excluded": excluded,
+        "excluded_names": excluded_names,
+        "by_param": by_param,
+        "by_buffer": by_buffer,
+    }
+
+
+def _group(model: nn.Module) -> "OrderedDict[str, int]":
+    """Spec-shaped component table, computed from the live model."""
+    cfg = model.config
+    base = model.model if hasattr(model, "model") else model
+    rows: "OrderedDict[str, int]" = OrderedDict()
+
+    def count(mods) -> int:
+        seen, total = set(), 0
+        for m in mods:
+            if m is None:
+                continue
+            for p in m.parameters(recurse=True):
+                if id(p) not in seen:
+                    seen.add(id(p))
+                    total += p.numel()
+        return total
+
+    layer0 = base.layers[0]
+    n_layer = len(base.layers)
+
+    attn = count([layer0.attn.q_proj, layer0.attn.k_proj, layer0.attn.v_proj, layer0.attn.o_proj])
+    qk = count([layer0.attn.q_norm, layer0.attn.k_norm])
+    gate = count([layer0.attn.gate_proj])
+    mlp = count([layer0.mlp])
+    norms = count([layer0.attn_in_norm, layer0.attn_out_norm, layer0.mlp_in_norm, layer0.mlp_out_norm])
+    canon = count([layer0.canon_a, layer0.canon_c])
+    per_layer = attn + qk + gate + mlp + norms + canon
+
+    rows["Attention q, k, v, o"] = attn
+    rows["QK-norm"] = qk
+    rows[f"{cfg.attn_gate.capitalize()} gate (W + b)"] = gate
+    rows[f"SwiGLU MLP (3 x {cfg.d_model} x {cfg.ffn_hidden})"] = mlp
+    rows[f"Peri-LN, 4 norms x {cfg.d_model}"] = norms
+    rows["Canon A + C"] = canon
+    rows["--- Per layer"] = per_layer
+    rows[f"{n_layer} layers"] = per_layer * n_layer
+
+    emb = base.embed_tokens.weight.numel()
+    rows[f"Tied embedding ({cfg.vocab_size} x {cfg.d_model})"] = emb
+    rows["Final norm"] = base.final_norm.weight.numel()
+
+    info = account(model)
+    rows["--- Trainable"] = info["trainable"]
+    if info["frozen"]:
+        rows["Frozen buffers"] = info["frozen"]
+    rows["--- Total"] = info["total"]
+    return rows
+
+
+def table(model: nn.Module) -> str:
+    rows = _group(model)
+    info = account(model)
+    width = max(len(k) for k in rows) + 2
+    lines = [f"{'Component'.ljust(width)}{'Count':>12}", "-" * (width + 12)]
+    for key, val in rows.items():
+        if key.startswith("--- "):
+            lines.append("-" * (width + 12))
+            key = key[4:]
+        lines.append(f"{key.ljust(width)}{val:>12,}")
+    lines.append("-" * (width + 12))
+
+    budget = getattr(model.config, "max_total_params", None)
+    warn = getattr(model.config, "warn_total_params", None)
+    if budget:
+        ok = "OK" if info["total"] <= budget else "OVER BUDGET"
+        lines.append(f"{'Budget (max_total_params)'.ljust(width)}{budget:>12,}  [{ok}]")
+    if warn:
+        flag = "ok" if info["total"] <= warn else "WARN"
+        lines.append(f"{'Warn threshold'.ljust(width)}{warn:>12,}  [{flag}]")
+    if info["excluded"]:
+        lines.append(
+            f"Excluded derived buffers ({info['excluded']:,}): "
+            + ", ".join(info["excluded_names"])
+        )
+    return "\n".join(lines)
+
+
+def count_from_config(cfg_or_model) -> dict:
+    """Build the model if given a config dict, then account for it.
+
+    Construction itself raises when the budget is exceeded, so this doubles as
+    the budget check.
+    """
+    if isinstance(cfg_or_model, nn.Module):
+        return account(cfg_or_model)
+    from ethoric.model.configuration_ethoric import EthoricConfig
+    from ethoric.model.modeling_ethoric import EthoricForCausalLM
+
+    model_cfg = EthoricConfig.from_pipeline(cfg_or_model)
+    with torch.device("meta"):
+        model = EthoricForCausalLM(model_cfg)
+    return account(model)