Unify the two v0.2->v0.3 migration surfaces (issues.md Issue 6)
giant/config.py:migrate_config (config.toml) and giant/model/network.py:_migrate_legacy_model_config (checkpoint model_config) independently hand-maintained the same v0.2 facts and an identical router expert-sizing rejection. Extract the shared knowledge into a new leaf module, giant/_migration.py (V02_MODEL_KEY_TO_STAGES, V02_FIXED_FACTS, reject_legacy_router_expert_sizing), consumed by both. Also replace NSecConfig's legacy-only, nullable legacy_owner sentinel (living in an extra: dict catch-all) with a normal, always-set owner: str = "stage2" field, so build_models reads one concrete two-valued key instead of branching on a legacy marker. Record in CLAUDE.md that v0.2 checkpoint-loading support has no expiry decided yet, since /ceph still holds pre-v0.3.0 checkpoints. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
@@ -91,6 +91,6 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
|
||||
|
||||
A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
|
||||
|
||||
**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time.
|
||||
**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. This config break is why v0.2-shaped configs/checkpoints need migrating at all (`config.migrate_config`, `model.network._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`) — v0.2 checkpoint-loading support has **no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.
|
||||
|
||||
**Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet.
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Shared v0.2 -> v0.3 migration knowledge.
|
||||
|
||||
v0.3.0 broke the config format (single `[train]` + `[model]` -> `[conditioning]`/
|
||||
`[stage1_model]`/`[stage2_model]`/`[train]`), and that break has to be absorbed by two
|
||||
independent migration surfaces: `giant.config.migrate_config` (a v0.2 `config.toml`) and
|
||||
`giant.model.network._migrate_legacy_model_config` (a v0.2 checkpoint's flat
|
||||
`model_config` dict). Both translate the same v0.2 facts into the same v0.3 shape, so
|
||||
the facts live here once rather than as two hand-maintained copies — see issues.md
|
||||
Issue 6.
|
||||
|
||||
A dependency-free leaf module so neither `config.py` nor `network.py` has to import the
|
||||
other to share this.
|
||||
"""
|
||||
|
||||
# v0.2 model-shaped keys (config.toml's [model] table, or a checkpoint's flat
|
||||
# model_config dict — same key names in both) applied identically to both v0.3 stage
|
||||
# blocks, because v0.2 had only one trunk shape shared by both stages.
|
||||
V02_MODEL_KEY_TO_STAGES: tuple[tuple[str, str], ...] = (
|
||||
("hidden_dim", "hidden_dim"),
|
||||
("n_blocks", "n_res_blocks"),
|
||||
("dropout", "dropout"),
|
||||
)
|
||||
|
||||
# v0.2 architectural facts that had no corresponding config key at all — always true of
|
||||
# a v0.2 model, so both migration surfaces inject them unconditionally. Keyed by dotted
|
||||
# path relative to the migrated dict's root. NOTE: conditioning.*.n_layers (2) differs
|
||||
# from the v0.3 *default* (1) — not a typo, v0.2's conditioning MLP was always 2 layers
|
||||
# deep.
|
||||
V02_FIXED_FACTS: dict[str, object] = {
|
||||
"conditioning.out_dim": 128,
|
||||
"conditioning.particle.n_layers": 2,
|
||||
"conditioning.material.n_layers": 2,
|
||||
"stage1_model.active": True,
|
||||
"stage1_model.flow.time_dim": 64,
|
||||
"stage1_model.ddpm.time_dim": 64,
|
||||
"stage2_model.active": True,
|
||||
"stage2_model.flow.time_dim": 64,
|
||||
"stage2_model.ddpm.time_dim": 64,
|
||||
"stage2_model.context_dim": 64,
|
||||
"stage2_model.decoder": "one_shot",
|
||||
"stage2_model.particle_type.target": "physical",
|
||||
}
|
||||
|
||||
|
||||
def reject_legacy_router_expert_sizing(router_cfg: dict, *, source: str) -> None:
|
||||
"""Pop and validate v0.2's per-expert width/depth override, in place.
|
||||
|
||||
v0.3.0 removed per-expert sizing — experts always inherit the stage's
|
||||
hidden_dim/n_res_blocks — so a v0.2 router config/checkpoint that set a non-default
|
||||
`expert_hidden_dim`/`expert_n_blocks` describes experts with a different width/depth
|
||||
than the monolith, and can only be reproduced by v0.2 code. Silently dropping these
|
||||
keys (a router builder's kwarg filtering would do this for free) would resize the
|
||||
experts instead of refusing, so this raises loudly.
|
||||
|
||||
Always pops both keys, whether or not they were non-default, so callers can go on
|
||||
to use the (now-cleaned) `router_cfg` unconditionally. `source` names what's being
|
||||
migrated (e.g. "v0.2 config's model.router" or "this checkpoint's
|
||||
model_config.router") for the error message.
|
||||
"""
|
||||
expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0)
|
||||
expert_n_blocks = router_cfg.pop("expert_n_blocks", 0)
|
||||
if not (expert_hidden_dim or expert_n_blocks):
|
||||
return
|
||||
raise ValueError(
|
||||
f"{source} sets expert_hidden_dim/expert_n_blocks to a non-default value "
|
||||
f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed per-expert "
|
||||
"sizing (experts always inherit the stage's hidden_dim/n_res_blocks), so "
|
||||
"this router's experts have a different width/depth than the monolith. "
|
||||
"This checkpoint/config can only be loaded by v0.2 code."
|
||||
)
|
||||
+20
-53
@@ -14,6 +14,8 @@ from pathlib import Path
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing
|
||||
|
||||
|
||||
class Conditioning(str, Enum):
|
||||
"""`conditioning.particle.type` / `conditioning.material.type` choices —
|
||||
@@ -47,13 +49,11 @@ CONFIG_VERSION = 3
|
||||
# `lambda` is a Python keyword, so dict key "lambda" is always exposed as the
|
||||
# field `lambda_weight`.
|
||||
#
|
||||
# Two sub-blocks — router and n_sec — carry genuinely dynamic keys that don't
|
||||
# fit a fixed schema: composed-router `axis{i}_{field}` flags (see
|
||||
# giant.model.network._parse_composed_axes) and pipeline.py's runtime-seeded
|
||||
# `centers_init`, plus n_sec's `legacy_owner` (injected only by
|
||||
# _migrate_legacy_model_config for v0.2 checkpoints). Both dataclasses carry
|
||||
# an `extra: dict` catch-all so these keys round-trip losslessly without
|
||||
# becoming named fields that would leak into every new run's config.toml.
|
||||
# The router sub-block carries genuinely dynamic keys that don't fit a fixed schema:
|
||||
# composed-router `axis{i}_{field}` flags (see giant.model.network._parse_composed_axes)
|
||||
# and pipeline.py's runtime-seeded `centers_init`. It carries an `extra: dict` catch-all
|
||||
# so these keys round-trip losslessly without becoming named fields that would leak into
|
||||
# every new run's config.toml.
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -384,26 +384,24 @@ class NSecConfig:
|
||||
# only, never for rollout.
|
||||
mode: str = "head"
|
||||
lambda_weight: float = 0.1 # dict key "lambda" — cross-entropy weight for the head
|
||||
# Holds "legacy_owner" when injected by _migrate_legacy_model_config
|
||||
# (v0.2 checkpoints only) — not a user-facing config.toml key.
|
||||
extra: dict = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def legacy_owner(self) -> str | None:
|
||||
return self.extra.get("legacy_owner")
|
||||
# Which stage's module physically owns the n_sec_head weights: "stage2" (default,
|
||||
# fresh v0.3.0 runs — Stage2OneShot/Stage2Autoregressive builds it) or "stage1"
|
||||
# (a migrated v0.2 checkpoint — see network._migrate_legacy_model_config, whose
|
||||
# n_sec head was trained against Stage 1's own ConditionEncoder output and so has
|
||||
# to stay attached there, not just be labeled as such).
|
||||
owner: str = "stage2"
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict | None) -> "NSecConfig":
|
||||
d = d or {}
|
||||
known = {"mode", "lambda"}
|
||||
return cls(
|
||||
mode=d.get("mode", "head"),
|
||||
lambda_weight=d.get("lambda", 0.1),
|
||||
extra={k: v for k, v in d.items() if k not in known},
|
||||
owner=d.get("owner", "stage2"),
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {"mode": self.mode, "lambda": self.lambda_weight, **self.extra}
|
||||
return {"mode": self.mode, "lambda": self.lambda_weight, "owner": self.owner}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -1024,14 +1022,6 @@ _V02_TRAIN_PASSTHROUGH = (
|
||||
"wandb_log_every",
|
||||
)
|
||||
|
||||
# v0.2 model.hidden_dim/n_blocks/dropout applied identically to both stages
|
||||
# (there was only ever one trunk shape) -> copied to both stage{1,2}_model.
|
||||
_V02_MODEL_TO_BOTH_STAGES = (
|
||||
("hidden_dim", "hidden_dim"),
|
||||
("n_blocks", "n_res_blocks"),
|
||||
("dropout", "dropout"),
|
||||
)
|
||||
|
||||
# v0.2 train.{n_critic,gp_weight,critic_lr} applied identically to both
|
||||
# stages' wgan sub-table (there was only ever one wgan objective, shared).
|
||||
_V02_TRAIN_TO_BOTH_STAGES_WGAN = (
|
||||
@@ -1091,7 +1081,7 @@ def migrate_config(cfg: dict) -> dict:
|
||||
_set_path(new, f"stage1_model.wgan.{new_key}", old_train[old_key])
|
||||
_set_path(new, f"stage2_model.wgan.{new_key}", old_train[old_key])
|
||||
|
||||
for old_key, new_key in _V02_MODEL_TO_BOTH_STAGES:
|
||||
for old_key, new_key in V02_MODEL_KEY_TO_STAGES:
|
||||
if old_key in old_model:
|
||||
_set_path(new, f"stage1_model.{new_key}", old_model[old_key])
|
||||
_set_path(new, f"stage2_model.{new_key}", old_model[old_key])
|
||||
@@ -1108,18 +1098,7 @@ def migrate_config(cfg: dict) -> dict:
|
||||
_set_path(new, "stage2_model.k_max", old_model["k_max"])
|
||||
|
||||
if old_router:
|
||||
expert_hidden_dim = old_router.pop("expert_hidden_dim", 0)
|
||||
expert_n_blocks = old_router.pop("expert_n_blocks", 0)
|
||||
if expert_hidden_dim or expert_n_blocks:
|
||||
raise ValueError(
|
||||
"v0.2 config sets model.router.expert_hidden_dim/"
|
||||
f"expert_n_blocks to a non-default value "
|
||||
f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed "
|
||||
"per-expert sizing (experts always inherit the stage's "
|
||||
"hidden_dim/n_res_blocks), so this config's routed experts "
|
||||
"have a different width/depth than the monolith and its "
|
||||
"checkpoint can only be loaded by v0.2 code."
|
||||
)
|
||||
reject_legacy_router_expert_sizing(old_router, source="v0.2 config's model.router")
|
||||
_set_path(new, "stage1_model.router", dict(old_router))
|
||||
stage2_router = dict(old_router)
|
||||
stage2_router["tie_to_stage1"] = False
|
||||
@@ -1127,21 +1106,9 @@ def migrate_config(cfg: dict) -> dict:
|
||||
|
||||
# v0.2 architectural facts with no corresponding config key at all —
|
||||
# always set once we've determined we're migrating a v0.2 dict,
|
||||
# independent of what the file did/didn't specify. NOTE: n_layers here
|
||||
# (2) differs from the v0.3 *default* (1) — this is not a typo, see the
|
||||
# docstring above.
|
||||
_set_path(new, "conditioning.out_dim", 128)
|
||||
_set_path(new, "conditioning.particle.n_layers", 2)
|
||||
_set_path(new, "conditioning.material.n_layers", 2)
|
||||
_set_path(new, "stage1_model.active", True)
|
||||
_set_path(new, "stage1_model.flow.time_dim", 64)
|
||||
_set_path(new, "stage1_model.ddpm.time_dim", 64)
|
||||
_set_path(new, "stage2_model.active", True)
|
||||
_set_path(new, "stage2_model.flow.time_dim", 64)
|
||||
_set_path(new, "stage2_model.ddpm.time_dim", 64)
|
||||
_set_path(new, "stage2_model.context_dim", 64)
|
||||
_set_path(new, "stage2_model.decoder", "one_shot")
|
||||
_set_path(new, "stage2_model.particle_type.target", "physical")
|
||||
# independent of what the file did/didn't specify (see giant._migration).
|
||||
for path, value in V02_FIXED_FACTS.items():
|
||||
_set_path(new, path, value)
|
||||
|
||||
new_meta = dict(cfg.pop("meta", {}))
|
||||
new_meta["config_version"] = CONFIG_VERSION
|
||||
|
||||
+31
-42
@@ -7,6 +7,7 @@ import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from giant._migration import V02_FIXED_FACTS, reject_legacy_router_expert_sizing
|
||||
from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig
|
||||
from giant.constants import (
|
||||
COND_DIM,
|
||||
@@ -983,7 +984,7 @@ class Stage1Model(nn.Module):
|
||||
raise RuntimeError(
|
||||
"this Stage1Model has no n_sec_head — n_sec now lives on "
|
||||
"stage 2 by default; this method only exists "
|
||||
"for a migrated v0.2 checkpoint (legacy_owner='stage1')"
|
||||
"for a migrated v0.2 checkpoint (n_sec.owner='stage1')"
|
||||
)
|
||||
c_emb = self.cond_enc(cond_cont, cond_cat)
|
||||
return self.n_sec_head(c_emb)
|
||||
@@ -1103,7 +1104,7 @@ class Stage2OneShot(nn.Module):
|
||||
if self.n_sec_head is None:
|
||||
raise RuntimeError(
|
||||
"this Stage2OneShot has no n_sec_head — it belongs to a "
|
||||
"migrated v0.2 checkpoint (legacy_owner='stage1'); call "
|
||||
"migrated v0.2 checkpoint (n_sec.owner='stage1'); call "
|
||||
"stage1.predict_n_sec(cond_cont, cond_cat) instead"
|
||||
)
|
||||
c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
|
||||
@@ -1345,7 +1346,7 @@ class Stage2Autoregressive(nn.Module):
|
||||
if self.n_sec_head is None:
|
||||
raise RuntimeError(
|
||||
"this Stage2Autoregressive has no n_sec_head — it belongs to "
|
||||
"a migrated v0.2 checkpoint (legacy_owner='stage1'); call "
|
||||
"a migrated v0.2 checkpoint (n_sec.owner='stage1'); call "
|
||||
"stage1.predict_n_sec(cond_cont, cond_cat) instead"
|
||||
)
|
||||
return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out))
|
||||
@@ -1451,12 +1452,11 @@ def _migrate_legacy_model_config(model_config: dict) -> dict:
|
||||
`{"pdg_vocab", "mat_vocab", "conditioning", "stage1_model",
|
||||
"stage2_model"}` shape `build_models` expects.
|
||||
|
||||
Sets `stage2_model.n_sec.legacy_owner = "stage1"` so the n_sec_head
|
||||
weights a v0.2 checkpoint carries on its Stage-1 module keep loading
|
||||
there instead of the new default location
|
||||
(`Stage2OneShot`) — the n_sec head was trained against Stage 1's own
|
||||
`ConditionEncoder` output, so it has to stay attached to Stage 1's
|
||||
module, not just be labeled as such.
|
||||
Sets `stage2_model.n_sec.owner = "stage1"` so the n_sec_head weights a v0.2
|
||||
checkpoint carries on its Stage-1 module keep loading there instead of the new
|
||||
default location (`Stage2OneShot`) — the n_sec head was trained against Stage 1's
|
||||
own `ConditionEncoder` output, so it has to stay attached to Stage 1's module, not
|
||||
just be labeled as such.
|
||||
|
||||
Only the monolithic (non-routed) trunk shape is exercised by the step-2
|
||||
migration test; a routed v0.2 checkpoint still builds correctly here
|
||||
@@ -1473,55 +1473,44 @@ def _migrate_legacy_model_config(model_config: dict) -> dict:
|
||||
k_max = m.get("k_max", K_MAX)
|
||||
noise_dim = m.get("noise_dim", 64)
|
||||
router_cfg = dict(m.get("router") or {})
|
||||
expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0)
|
||||
expert_n_blocks = router_cfg.pop("expert_n_blocks", 0)
|
||||
if expert_hidden_dim or expert_n_blocks:
|
||||
raise ValueError(
|
||||
"this checkpoint's model_config.router sets expert_hidden_dim/"
|
||||
f"expert_n_blocks to a non-default value "
|
||||
f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed "
|
||||
"per-expert sizing (experts always inherit the stage's "
|
||||
"hidden_dim/n_res_blocks), so this checkpoint's routed experts "
|
||||
"have a different width/depth than the monolith — silently "
|
||||
"dropping these keys would resize the experts instead of "
|
||||
"refusing. This checkpoint can "
|
||||
"only be loaded by v0.2 code."
|
||||
)
|
||||
reject_legacy_router_expert_sizing(router_cfg, source="this checkpoint's model_config.router")
|
||||
router_cfg.setdefault("enabled", False)
|
||||
|
||||
F = V02_FIXED_FACTS
|
||||
cond_n_layers = F["conditioning.particle.n_layers"] # same fact for both axes
|
||||
return {
|
||||
"pdg_vocab": m["pdg_vocab"],
|
||||
"mat_vocab": m["mat_vocab"],
|
||||
"conditioning": {
|
||||
"out_dim": 128,
|
||||
"out_dim": F["conditioning.out_dim"],
|
||||
"share_stages": False,
|
||||
"particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": 2},
|
||||
"material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": 2},
|
||||
"particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers},
|
||||
"material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers},
|
||||
},
|
||||
"stage1_model": {
|
||||
"active": True,
|
||||
"active": F["stage1_model.active"],
|
||||
"generator": generator,
|
||||
"hidden_dim": hidden_dim,
|
||||
"n_res_blocks": n_blocks,
|
||||
"dropout": dropout,
|
||||
"flow": {"time_dim": 64},
|
||||
"ddpm": {"time_dim": 64},
|
||||
"flow": {"time_dim": F["stage1_model.flow.time_dim"]},
|
||||
"ddpm": {"time_dim": F["stage1_model.ddpm.time_dim"]},
|
||||
"wgan": {"noise_dim": noise_dim},
|
||||
"router": dict(router_cfg),
|
||||
},
|
||||
"stage2_model": {
|
||||
"active": True,
|
||||
"decoder": "one_shot",
|
||||
"active": F["stage2_model.active"],
|
||||
"decoder": F["stage2_model.decoder"],
|
||||
"generator": generator,
|
||||
"hidden_dim": hidden_dim,
|
||||
"n_res_blocks": n_blocks,
|
||||
"dropout": dropout,
|
||||
"k_max": k_max,
|
||||
"context_dim": 64,
|
||||
"n_sec": {"mode": "head", "legacy_owner": "stage1"},
|
||||
"particle_type": {"target": "physical"},
|
||||
"flow": {"time_dim": 64},
|
||||
"ddpm": {"time_dim": 64},
|
||||
"context_dim": F["stage2_model.context_dim"],
|
||||
"n_sec": {"mode": "head", "owner": "stage1"},
|
||||
"particle_type": {"target": F["stage2_model.particle_type.target"]},
|
||||
"flow": {"time_dim": F["stage2_model.flow.time_dim"]},
|
||||
"ddpm": {"time_dim": F["stage2_model.ddpm.time_dim"]},
|
||||
"wgan": {"noise_dim": noise_dim},
|
||||
"router": {**router_cfg, "tie_to_stage1": False},
|
||||
},
|
||||
@@ -1545,7 +1534,7 @@ def migrate_legacy_state_dict(old_stage1_sd: dict, old_stage2_sd: dict) -> tuple
|
||||
new_stage1 = {}
|
||||
for k, v in old_stage1_sd.items():
|
||||
if k.startswith("n_sec_head."):
|
||||
new_stage1[k] = v # stays top-level (legacy_owner="stage1")
|
||||
new_stage1[k] = v # stays top-level (n_sec.owner="stage1")
|
||||
else:
|
||||
new_stage1[_trunk_prefix(k)] = v
|
||||
|
||||
@@ -1614,8 +1603,8 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
|
||||
# matches the pre-dataclass .get("time_dim", 64) fallback, which
|
||||
# always hit its default for a wgan sub-block too.
|
||||
time_dim = getattr(s1_spec, generator).time_dim if generator != "wgan" else 64
|
||||
legacy_owner = s2_spec.n_sec.legacy_owner
|
||||
n_sec_head_k_max = s2_spec.k_max if legacy_owner == "stage1" else None
|
||||
n_sec_owner = s2_spec.n_sec.owner
|
||||
n_sec_head_k_max = s2_spec.k_max if n_sec_owner == "stage1" else None
|
||||
result["stage1"] = Stage1Model(
|
||||
pdg_vocab=pdg_vocab,
|
||||
mat_vocab=mat_vocab,
|
||||
@@ -1646,7 +1635,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
|
||||
# wgan has no time_dim concept — see the matching comment in stage 1
|
||||
# above.
|
||||
time_dim = getattr(s2_spec, generator).time_dim if generator != "wgan" else 64
|
||||
legacy_owner = s2_spec.n_sec.legacy_owner
|
||||
n_sec_owner = s2_spec.n_sec.owner
|
||||
k_max = s2_spec.k_max
|
||||
particle_type_cfg = s2_spec.particle_type.to_dict()
|
||||
|
||||
@@ -1667,7 +1656,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
|
||||
noise_dim=s2_spec.wgan.noise_dim,
|
||||
k_max=k_max,
|
||||
router=stage2_router,
|
||||
build_n_sec_head=legacy_owner != "stage1",
|
||||
build_n_sec_head=n_sec_owner != "stage1",
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
history=ar_cfg.history,
|
||||
attn_n_heads=ar_cfg.attn_n_heads,
|
||||
@@ -1692,7 +1681,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
|
||||
noise_dim=s2_spec.wgan.noise_dim,
|
||||
k_max=k_max,
|
||||
router=stage2_router,
|
||||
build_n_sec_head=legacy_owner != "stage1",
|
||||
build_n_sec_head=n_sec_owner != "stage1",
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
cond_enc=shared_cond_enc,
|
||||
)
|
||||
|
||||
@@ -45,7 +45,7 @@ architecture matrix grows, not about rot or breakage.
|
||||
| 3 | `cli.py:train()` is a 58-parameter, 510-line fat controller | **High** | Medium | **Fixed** (`2bfb1ab`) |
|
||||
| 4 | `cli.py` is at 35.8 % coverage and holds untested override-precedence logic | **High** | Medium | **Fixed** (`2bfb1ab`, partial — see status note) |
|
||||
| 5 | Inference bootstrap is duplicated verbatim between `predict` and `rollout` | **High** | Small | **Fixed** |
|
||||
| 6 | Two independent v0.2→v0.3 migration surfaces encode the same knowledge | Medium | Medium | Open |
|
||||
| 6 | Two independent v0.2→v0.3 migration surfaces encode the same knowledge | Medium | Medium | **Fixed** |
|
||||
| 7 | Positional tuple contracts (9-tuple, 7-tuple) between data, model and training layers | Medium | Small | Open |
|
||||
| 8 | `network.py` is 1745 lines holding three distinct modules | Medium | Small | Open |
|
||||
| 9 | `scripts` is published as a top-level distribution package | Medium | Small | Open |
|
||||
@@ -770,6 +770,34 @@ imported at `analysis/` module scope.
|
||||
|
||||
## Issue 6 — Two independent v0.2→v0.3 migration surfaces encode the same knowledge
|
||||
|
||||
> **Status: Fixed.** A new leaf module, `giant/_migration.py`, now holds the v0.2 facts
|
||||
> both `config.migrate_config` and `network._migrate_legacy_model_config` need:
|
||||
> `V02_MODEL_KEY_TO_STAGES` (the `hidden_dim`/`n_blocks`/`dropout` → both-stages
|
||||
> mapping), `V02_FIXED_FACTS` (the ~12 "architectural facts with no config key" —
|
||||
> conditioning `out_dim`/`n_layers`, both stages' `active`/`flow.time_dim`/
|
||||
> `ddpm.time_dim`, stage 2's `context_dim`/`decoder`/`particle_type.target`), and
|
||||
> `reject_legacy_router_expert_sizing` (one canonical error message for a non-default
|
||||
> `router.expert_hidden_dim`/`expert_n_blocks`, replacing the two hand-written copies).
|
||||
> Both migration functions now consume this module instead of hardcoding their own
|
||||
> copies, so a future correction needs one edit, not two. Separately (fix item 2), the
|
||||
> checkpoint-only `n_sec.legacy_owner` nullable sentinel (living in `NSecConfig.extra`,
|
||||
> `None` unless a v0.2 checkpoint was migrated) is replaced by a normal, always-set
|
||||
> `NSecConfig.owner: str = "stage2"` field — the current schema now carries
|
||||
> `stage2_model.n_sec.owner` for every run (`"stage2"` by default, `"stage1"` only for a
|
||||
> migrated v0.2 checkpoint), so `build_models` reads one concrete two-valued key instead
|
||||
> of branching on a nullable legacy marker; `NSecConfig.extra` was dropped entirely since
|
||||
> `legacy_owner` was its only consumer. Fix item 3 (retention policy) is recorded in
|
||||
> `CLAUDE.md`'s v0.3.0 roadmap paragraph as **explicitly undecided** rather than a
|
||||
> concrete expiry — `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs
|
||||
> referencing them, and no criterion for dropping the shim (or
|
||||
> `tests/legacy/network_v02_snapshot.py`) has been agreed yet. `tests/legacy/
|
||||
> network_v02_snapshot.py` itself is untouched, per the scope guard below. Existing
|
||||
> regression tests updated for the rename
|
||||
> (`tests/test_config.py`, `tests/test_migration_v02_v03.py`); the four bit-identical
|
||||
> `test_migration_*` cases and the full suite (804 tests) pass unchanged, confirming no
|
||||
> migrated value changed. Everything below this point describes the pre-fix state and is
|
||||
> kept for historical context.
|
||||
|
||||
**Severity: Medium. Effort: Medium.**
|
||||
|
||||
**Location:** `giant/config.py:514` (`migrate_config`) and
|
||||
|
||||
@@ -95,10 +95,15 @@ def test_stage1_router_config_has_no_tie_to_stage1_key():
|
||||
assert "tie_to_stage1" not in gconfig.RouterConfig().to_dict()
|
||||
|
||||
|
||||
def test_n_sec_config_extra_round_trips_legacy_owner():
|
||||
n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "legacy_owner": "stage1"})
|
||||
assert n_sec.legacy_owner == "stage1"
|
||||
assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "legacy_owner": "stage1"}
|
||||
def test_n_sec_config_owner_defaults_to_stage2():
|
||||
n_sec = gconfig.NSecConfig()
|
||||
assert n_sec.owner == "stage2"
|
||||
|
||||
|
||||
def test_n_sec_config_owner_round_trips():
|
||||
n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"})
|
||||
assert n_sec.owner == "stage1"
|
||||
assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -130,7 +130,7 @@ def _run_migration_check(mode: str, conditioning: str) -> None:
|
||||
new_stage1, new_stage2 = new_models["stage1"], new_models["stage2"]
|
||||
assert isinstance(new_stage1, net.Stage1Model)
|
||||
assert isinstance(new_stage2, net.Stage2OneShot)
|
||||
# legacy_owner="stage1": n_sec lives on stage1, not stage2, for a
|
||||
# n_sec.owner="stage1": n_sec lives on stage1, not stage2, for a
|
||||
# migrated v0.2 checkpoint.
|
||||
assert new_stage1.n_sec_head is not None
|
||||
assert new_stage2.n_sec_head is None
|
||||
@@ -177,7 +177,7 @@ def test_migration_wgan_physical():
|
||||
|
||||
def test_migrate_legacy_model_config_shape():
|
||||
"""_migrate_legacy_model_config produces the nested shape build_models
|
||||
expects, with the legacy_owner marker set so build_models routes the
|
||||
expects, with the n_sec.owner marker set so build_models routes the
|
||||
n_sec head back onto stage 1."""
|
||||
legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
|
||||
migrated = net._migrate_legacy_model_config(legacy_cfg)
|
||||
@@ -187,7 +187,7 @@ def test_migrate_legacy_model_config_shape():
|
||||
assert migrated["conditioning"]["particle"]["n_layers"] == 2
|
||||
assert migrated["conditioning"]["material"]["n_layers"] == 2
|
||||
assert migrated["stage1_model"]["hidden_dim"] == HIDDEN_DIM
|
||||
assert migrated["stage2_model"]["n_sec"]["legacy_owner"] == "stage1"
|
||||
assert migrated["stage2_model"]["n_sec"]["owner"] == "stage1"
|
||||
assert migrated["stage2_model"]["decoder"] == "one_shot"
|
||||
|
||||
|
||||
@@ -278,6 +278,6 @@ def test_build_models_accepts_new_nested_shape_unchanged():
|
||||
models = net.build_models(cfg)
|
||||
assert isinstance(models["stage1"], net.Stage1Model)
|
||||
assert isinstance(models["stage2"], net.Stage2OneShot)
|
||||
# Fresh v0.3.0 config, no legacy_owner: n_sec lives on stage 2.
|
||||
# Fresh v0.3.0 config, n_sec.owner defaults to "stage2": n_sec lives on stage 2.
|
||||
assert models["stage1"].n_sec_head is None
|
||||
assert models["stage2"].n_sec_head is not None
|
||||
|
||||
Reference in New Issue
Block a user