giant/giant/model/summary.py

"""Build-only model introspection (gitea #46): construct the resolved
Stage1/Stage2/critic graph from a config with no dataset attached, and report
per-module parameter counts, trunk widths, which heads exist, and — via
differential probing — which `conditioning`/`stage1_model`/`stage2_model`
config keys actually shape the built model. This is the runtime counterpart
to `tests/test_config_consumed_keys.py`'s static per-identifier audit: that
test asks "does any code reference this key's name at all", this module asks
"given *this* resolved config, does the key change what `build_models`/
`build_critics` (`giant/model/builders.py`) actually produces".

Differential probing, not identifier matching: build the model once from the
resolved config and take a structural fingerprint (`_fingerprint` — which
submodules exist, every parameter's/buffer's shape+dtype, every plain scalar
attribute stored on any module). Then, for each in-scope leaf key, perturb
just that one value (`_perturb`), rebuild, and re-fingerprint. A changed
fingerprint — or a rebuild that raises — means the key was consumed; an
identical fingerprint means construction never looked at it under this
particular config. A key can be genuinely inert under one config and live
under another (e.g. any `stage1_model.router.*` key when `router.enabled =
false`) — that config-dependence is exactly the "silently degenerate
combination" issue #46 is after, so it is reported per-run rather than
baked into a static table.

Keys legitimately owned by the trainer/sampler/rollout rather than by
`build_models`/`build_critics` (loss weights, WGAN-GP training
hyperparameters, teacher-forcing and stage1-context schedules, ...) are
cataloged in `_NOT_BUILD_TIME` below so the report doesn't flag them as
suspicious. One leaf is inert under every config today —
`stage2_model.autoregressive.order` — matching
`tests/test_config_consumed_keys.py`'s own `_KNOWN_UNUSED` entry; it is
deliberately *not* in `_NOT_BUILD_TIME`, since "always inert" is itself the
finding those two tests independently converge on.
"""

import copy
from dataclasses import dataclass, field

import torch.nn as nn

from giant.config import _get_path, _set_path, leaf_paths
from giant.model.builders import build_critics, build_models
from giant.model.trunks import RoutedTrunk

_IN_SCOPE_ROOTS = ("conditioning", "stage1_model", "stage2_model")

_PROBE_STR = "__giant_model_summary_probe__"

# A handful of string leaves branch on equality against one specific literal
# (e.g. `builders.py`: `stop_token = s2_spec.n_sec.mode == "stop_token"`),
# where every value other than that literal behaves identically. A single
# generic sentinel probe would then falsely read as inert whenever the
# config's *current* value is already one of those identically-behaving
# "other" values (e.g. mode="head") — it never crosses the one boundary that
# actually matters. Named here so probing tries the real alternative(s) too;
# every other string leaf is registry-validated (raises on garbage, still
# correctly detected as consumed) or genuinely value-independent, so doesn't
# need an entry.
_STRING_ALTERNATIVES: dict[str, tuple[str, ...]] = {
    "stage2_model.n_sec.owner": ("stage1", "stage2"),
    "stage2_model.n_sec.mode": ("stop_token", "head", "truth"),
    "stage2_model.particle_type.target": ("physical", "onehot", "embedding"),
}

# Verified by reading giant/training/trainers.py, giant/training/stage2_inputs.py
# and giant/rollout.py while implementing gitea #46 — not auto-derived, so a
# future reader touching these fields should re-check this table still holds.
_NOT_BUILD_TIME: dict[str, str] = {
    "stage1_model.init_from": "training/checkpoint.py's init_stages_from_checkpoints, run before build_stage_trainers (gitea #42)",
    "stage1_model.freeze": "trainers.py: StageSpec.freeze, gates StageTrainer._step_optimizer (gitea #42)",
    "stage2_model.init_from": "training/checkpoint.py's init_stages_from_checkpoints, run before build_stage_trainers (gitea #42)",
    "stage2_model.freeze": "trainers.py: StageSpec.freeze, gates StageTrainer._step_optimizer (gitea #42)",
    "stage1_model.lambda": "trainers.py: StageSpec.lambda_weight, the total-loss mix weight",
    "stage2_model.lambda": "trainers.py: StageSpec.lambda_weight, the total-loss mix weight",
    "stage2_model.n_sec.lambda": "trainers.py: StageSpec.n_sec_lambda, the n_sec-head loss weight",
    "stage2_model.particle_type.lambda": "trainers.py: Stage2Trainer.particle_type_lambda, the type-head loss weight",
    "stage2_model.particle_type.other_policy": "giant/rollout.py: resolves an 'other'-bucket secondary's PDG code at inference",
    "stage2_model.particle_type.class_weighting": "trainers.py: FlowDDPMStageTrainer.type_class_weights, shapes the type-head loss, not the built graph (gitea #44)",
    "stage2_model.autoregressive.teacher_forcing": "giant/training/stage2_inputs.py's training-time input assembly",
    "stage2_model.autoregressive.tf_p_start": "trainers.py's teacher-forcing schedule",
    "stage2_model.autoregressive.tf_p_end": "trainers.py's teacher-forcing schedule",
    "stage2_model.stage1_context": "trainers.py's stage1/stage2 boundary — StageTrainer._stage1_context",
    "stage2_model.ctx_p_start": "trainers.py's stage1-context sampling schedule",
    "stage2_model.ctx_p_end": "trainers.py's stage1-context sampling schedule",
    "stage1_model.router.lambda_balance": "trainers.py's load-balancing auxiliary loss weight",
    "stage1_model.router.lambda_entropy": "trainers.py's entropy-regularization auxiliary loss weight",
    "stage1_model.router.lambda_proc": "trainers.py's supervised process-classification auxiliary loss weight",
    "stage1_model.router.gumbel_tau_start": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
    "stage1_model.router.gumbel_tau_end": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
    "stage2_model.router.lambda_balance": "trainers.py's load-balancing auxiliary loss weight",
    "stage2_model.router.lambda_entropy": "trainers.py's entropy-regularization auxiliary loss weight",
    "stage2_model.router.lambda_proc": "trainers.py's supervised process-classification auxiliary loss weight",
    "stage2_model.router.gumbel_tau_start": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
    "stage2_model.router.gumbel_tau_end": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
    "stage1_model.wgan.n_critic": "trainers.py's WGAN-GP critic-update cadence",
    "stage1_model.wgan.gp_weight": "trainers.py's WGAN-GP gradient-penalty coefficient",
    "stage1_model.wgan.critic_lr": "trainers.py's critic optimizer learning rate",
    "stage2_model.wgan.n_critic": "trainers.py's WGAN-GP critic-update cadence",
    "stage2_model.wgan.gp_weight": "trainers.py's WGAN-GP gradient-penalty coefficient",
    "stage2_model.wgan.critic_lr": "trainers.py's critic optimizer learning rate",
    "stage2_model.wgan.gumbel_tau_start": "trainers.py's type-slice Gumbel-softmax temperature anneal (type_gumbel_tau_start)",
    "stage2_model.wgan.gumbel_tau_end": "trainers.py's type-slice Gumbel-softmax temperature anneal (type_gumbel_tau_end)",
}


@dataclass
class ModelSummary:
    modules: dict[str, nn.Module]
    consumed: list[str]
    inert: list[str]
    elsewhere: list[str]
    pdg_vocab: int
    mat_vocab: int
    vocab_caveats: list[str] = field(default_factory=list)


def _build_model_config(cfg: dict, pdg_vocab: int, mat_vocab: int) -> dict:
    return {
        "pdg_vocab": pdg_vocab,
        "mat_vocab": mat_vocab,
        "conditioning": cfg["conditioning"],
        "stage1_model": cfg["stage1_model"],
        "stage2_model": cfg["stage2_model"],
    }


def _built_modules(cfg: dict, pdg_vocab: int, mat_vocab: int) -> dict[str, nn.Module]:
    model_config = _build_model_config(cfg, pdg_vocab, mat_vocab)
    modules: dict[str, nn.Module] = {}
    for name, m in build_models(model_config).items():
        if m is not None:
            modules[name] = m
    for name, m in build_critics(model_config).items():
        if m is not None:
            modules[f"{name}_critic"] = m
    return modules


def _fingerprint(modules: dict[str, nn.Module]) -> list:
    """A config-shape fingerprint of the built graph: which submodules
    exist, every parameter's/buffer's shape+dtype (never values — those are
    randomly initialized and irrelevant to *structure*), and every plain
    scalar attribute any module stores on itself (e.g. `Stage2Autoregressive
    .stop_sampling`, `EnergyRouter.temperature`) — this is what makes a
    non-parametric key's effect on construction observable."""
    sig = []
    for stage_name, module in modules.items():
        for mod_name, m in module.named_modules():
            full = f"{stage_name}.{mod_name}" if mod_name else stage_name
            for k, v in vars(m).items():
                if k.startswith("_"):
                    continue
                if v is None or isinstance(v, (bool, int, float, str)):
                    sig.append((full, k, v))
        for pname, p in module.named_parameters():
            sig.append((stage_name, "param", pname, tuple(p.shape), str(p.dtype)))
        for bname, b in module.named_buffers():
            sig.append((stage_name, "buffer", bname, tuple(b.shape), str(b.dtype)))
    return sorted(sig, key=repr)


def _perturb_candidates(path: str, value) -> list:
    """Values to try perturbing `path`'s current `value` to, in order —
    probing stops at the first one that changes the fingerprint or raises.
    Almost always a single candidate; see `_STRING_ALTERNATIVES`."""
    if isinstance(value, bool):
        return [not value]
    if isinstance(value, int):
        return [value + 1]
    if isinstance(value, float):
        return [value + 1.0]
    if isinstance(value, str):
        alternatives = [v for v in _STRING_ALTERNATIVES.get(path, ()) if v != value]
        return [*alternatives, _PROBE_STR]
    raise TypeError(f"gitea #46 probing: unsupported leaf value type {type(value)!r} ({value!r})")


def _vocab_caveats(cfg: dict) -> list[str]:
    caveats = []
    if _get_path(cfg, "conditioning.particle.type") == "embedding":
        caveats.append(
            "conditioning.particle.type = 'embedding' -- pdg_vocab below is a "
            "placeholder (no dataset attached to derive the real training vocab size)"
        )
    if _get_path(cfg, "conditioning.material.type") == "embedding":
        caveats.append(
            "conditioning.material.type = 'embedding' -- mat_vocab below is a "
            "placeholder (no dataset attached to derive the real training vocab size)"
        )
    for stage in ("stage1_model", "stage2_model"):
        router_type = _get_path(cfg, f"{stage}.router.type")
        if _get_path(cfg, f"{stage}.router.enabled") and router_type in ("pdg", "process"):
            caveats.append(
                f"{stage}.router.type = {router_type!r} builds its own pdg_vocab-sized "
                "embedding -- the count above is a placeholder"
            )
    return caveats


def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
    """Build `cfg`'s model with no dataset attached and report its resolved
    graph, plus which `conditioning`/`stage1_model`/`stage2_model` config
    keys actually shaped it (differential probing — see module docstring).
    `cfg` must already be a fully-merged v0.3 config (`merge_cli_overrides`
    output) — this does not migrate or validate it."""
    modules = _built_modules(cfg, pdg_vocab, mat_vocab)
    baseline_fp = _fingerprint(modules)

    in_scope = [p for p in leaf_paths(cfg) if p.split(".", 1)[0] in _IN_SCOPE_ROOTS]
    consumed: list[str] = []
    inert: list[str] = []
    elsewhere: list[str] = []
    for path in in_scope:
        original = _get_path(cfg, path)
        changed = False
        for candidate in _perturb_candidates(path, original):
            probe_cfg = copy.deepcopy(
                {
                    "conditioning": cfg["conditioning"],
                    "stage1_model": cfg["stage1_model"],
                    "stage2_model": cfg["stage2_model"],
                }
            )
            _set_path(probe_cfg, path, candidate)
            try:
                changed = _fingerprint(_built_modules(probe_cfg, pdg_vocab, mat_vocab)) != baseline_fp
            except Exception:
                changed = True
            if changed:
                break
        if changed:
            consumed.append(path)
        elif path in _NOT_BUILD_TIME:
            elsewhere.append(path)
        else:
            inert.append(path)

    return ModelSummary(
        modules=modules,
        consumed=sorted(consumed),
        inert=sorted(inert),
        elsewhere=sorted(elsewhere),
        pdg_vocab=pdg_vocab,
        mat_vocab=mat_vocab,
        vocab_caveats=_vocab_caveats(cfg),
    )


def _tree_lines(module: nn.Module, name: str, indent: int = 0) -> list[str]:
    total = sum(p.numel() for p in module.parameters())
    in_dim = getattr(module, "in_dim", None)
    out_dim = getattr(module, "out_dim", None)
    widths = f"  [in={in_dim}, out={out_dim}]" if in_dim is not None and out_dim is not None else ""
    lines = [f"{'  ' * indent}{name} ({type(module).__name__}): {total:,}{widths}"]
    for child_name, child in module.named_children():
        lines.extend(_tree_lines(child, child_name, indent + 1))
    return lines


_HEAD_NAMES = ("n_sec_head", "type_head", "stop_head")


def _stage_header(name: str, module: nn.Module) -> list[str]:
    total = sum(p.numel() for p in module.parameters())
    lines = [f"{name}: {type(module).__name__} -- {total:,} parameters"]
    generator = getattr(module, "generator_kind", None)
    if generator is not None:
        lines.append(f"  generator: {generator}")
    trunk = getattr(module, "trunk", None)
    if trunk is not None:
        in_dim = getattr(trunk, "in_dim", "?")
        out_dim = getattr(trunk, "out_dim", "?")
        if isinstance(trunk, RoutedTrunk):
            detail = f"routed, n_experts={trunk.router.n_experts}, expert type={type(trunk.experts[0]).__name__}"
        else:
            detail = f"unrouted, {type(trunk).__name__}"
        lines.append(f"  trunk: {detail}, in={in_dim}, out={out_dim}")
    history_kind = getattr(module, "history_kind", None)
    if history_kind is not None:
        lines.append(f"  autoregressive history: {history_kind}")
    present = [h for h in _HEAD_NAMES if getattr(module, h, None) is not None]
    absent = [h for h in _HEAD_NAMES if hasattr(module, h) and getattr(module, h) is None]
    if present or absent:
        lines.append(f"  heads present: {', '.join(present) if present else 'none'}")
        if absent:
            lines.append(f"  heads absent: {', '.join(absent)}")
    return lines


def render_summary(summary: ModelSummary) -> str:
    lines: list[str] = []
    for name, module in summary.modules.items():
        lines.extend(_stage_header(name, module))
        lines.extend(_tree_lines(module, name, indent=1))
        lines.append("")

    lines.append(
        f"config keys read during construction: {len(summary.consumed)} / "
        f"read elsewhere (trainer/sampler/rollout): {len(summary.elsewhere)} / "
        f"inert under this config: {len(summary.inert)}"
    )
    if summary.elsewhere:
        lines.append("read elsewhere, not by construction:")
        for path in summary.elsewhere:
            lines.append(f"  {path}  ({_NOT_BUILD_TIME[path]})")
    lines.append("inert under this config (declared, parsed, but doing nothing here):")
    if summary.inert:
        for path in summary.inert:
            lines.append(f"  {path}")
    else:
        lines.append("  (none)")

    if summary.vocab_caveats:
        lines.append("")
        lines.append("vocab placeholder caveats:")
        for caveat in summary.vocab_caveats:
            lines.append(f"  {caveat}")

    return "\n".join(lines)