| """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 |
|
|
| |
| |
| 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 [] |
| |
| 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: |
| 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) |
|
|
