lars
e6261cea03
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>
71 lines
3.4 KiB
Python
71 lines
3.4 KiB
Python
"""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."
|
|
)
|