"""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 INFERENCE_OVERRIDES, _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) overridable: 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 .n_sec_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) overridable = sorted(p for p in in_scope if p in INFERENCE_OVERRIDES) 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), overridable=overridable, ) 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.overridable: lines.append("") lines.append("inference-overridable without retraining (giant predict/rollout --set):") for path in summary.overridable: lines.append(f" {path} ({INFERENCE_OVERRIDES[path].why})") 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)