lars
9ce55e5013
Decomposes the ten permutation classes in giant/model/network.py into the reusable parts from docs/v0.3.0-design.md §5: ConditionEncoder (now independently configurable per particle/material axis), ContextAdapter, Trunk/MonolithicTrunk/RoutedTrunk/ExpertTrunk, and the stage classes Stage1Model/Stage2OneShot/CriticModel (Stage2Autoregressive stubbed, raises NotImplementedError until step 4/5). build_models/build_critics now return a dict keyed by stage and accept the new nested config shape, with routed WGAN reachable for the first time (the old --mode wgan --router rejection is gone) and stage2_model.router.tie_to_stage1 sharing a literal Router instance. A v0.2 checkpoint's flat model_config auto-migrates via _migrate_legacy_model_config + migrate_legacy_state_dict, preserving the n_sec_head's attachment to Stage1Model (legacy_owner="stage1", design doc §4.1). tests/test_migration_v02_v03.py proves this bit-identical against a frozen v0.2 snapshot (tests/legacy/network_v02_snapshot.py) for both flow and wgan, both conditioning modes. scripts/check_migration_v02_v03.py is the real-checkpoint counterpart for a portal machine with /ceph access. giant/model/schedule.py's flow-matching/DDPM loss helpers are updated to the new model-call convention (t as a keyword). giant/sample.py, giant/rollout.py, and giant/validate.py are not yet updated (deferred to design doc step 6) — their exercising tests are marked xfail with that reasoning rather than silently broken. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
242 lines
8.7 KiB
Python
242 lines
8.7 KiB
Python
"""Migration acceptance test for v0.3.0 step 2 (docs/v0.3.0-design.md §4.3,
|
|
§12 step 2): "load a v0.2 checkpoint through migrate_config + the new
|
|
build_models, and diff its outputs against v0.2 code on the same input
|
|
batch — bit-identical, or the refactor has changed something it should not
|
|
have."
|
|
|
|
No `/ceph` access on this machine (see CLAUDE.md's Compute environment
|
|
section), so a real trained checkpoint can't be used here — see
|
|
docs/v0.3.0-design.md's plan for the separate portal-machine follow-up with a
|
|
real checkpoint. This test is the synthetic stand-in: build a v0.2-shaped
|
|
model from the frozen `tests/legacy/network_v02_snapshot.py` classes with
|
|
fixed-seed random weights (playing the role of "a v0.2 checkpoint"), migrate
|
|
its config and remap its state dict onto the new `build_models` output, and
|
|
assert the two produce bit-identical output on the same random input batch.
|
|
"""
|
|
|
|
import torch
|
|
|
|
from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
|
|
from giant.model import network as net
|
|
from tests.legacy import network_v02_snapshot as legacy
|
|
|
|
PDG_VOCAB = 12
|
|
MAT_VOCAB = 4
|
|
HIDDEN_DIM = 32
|
|
N_BLOCKS = 2
|
|
EMB_DIM = 8
|
|
K = 6 # small k_max for a fast test
|
|
BATCH = 5
|
|
|
|
|
|
def _legacy_model_config(mode: str, conditioning: str) -> dict:
|
|
return {
|
|
"pdg_vocab": PDG_VOCAB,
|
|
"mat_vocab": MAT_VOCAB,
|
|
"hidden_dim": HIDDEN_DIM,
|
|
"n_blocks": N_BLOCKS,
|
|
"emb_dim": EMB_DIM,
|
|
"dropout": 0.0,
|
|
"k_max": K,
|
|
"conditioning": conditioning,
|
|
"router": {"enabled": False},
|
|
"mode": mode,
|
|
"noise_dim": 16,
|
|
}
|
|
|
|
|
|
def _random_batch(seed: int):
|
|
g = torch.Generator().manual_seed(seed)
|
|
cond_cont = torch.randn(BATCH, COND_DIM, generator=g)
|
|
cond_cat = torch.randint(0, min(PDG_VOCAB, MAT_VOCAB), (BATCH, 2), generator=g)
|
|
x1 = torch.randn(BATCH, X_DIM, generator=g)
|
|
x2 = torch.randn(BATCH, K * SEC_SLOT_DIM, generator=g)
|
|
t = torch.rand(BATCH, generator=g)
|
|
return cond_cont, cond_cat, x1, x2, t
|
|
|
|
|
|
def _assert_bit_identical(a: torch.Tensor, b: torch.Tensor, label: str) -> None:
|
|
assert a.shape == b.shape, f"{label}: shape mismatch {a.shape} vs {b.shape}"
|
|
assert torch.equal(a, b), (
|
|
f"{label}: outputs diverged, max abs diff = {(a - b).abs().max().item()}"
|
|
)
|
|
|
|
|
|
def _run_migration_check(mode: str, conditioning: str) -> None:
|
|
torch.manual_seed(0)
|
|
legacy_cfg = _legacy_model_config(mode, conditioning)
|
|
|
|
if mode == "wgan":
|
|
old_stage1 = legacy.WGANGenerator(
|
|
pdg_vocab=PDG_VOCAB,
|
|
mat_vocab=MAT_VOCAB,
|
|
hidden_dim=HIDDEN_DIM,
|
|
n_blocks=N_BLOCKS,
|
|
emb_dim=EMB_DIM,
|
|
noise_dim=16,
|
|
dropout=0.0,
|
|
k_max=K,
|
|
conditioning=conditioning,
|
|
)
|
|
old_stage2 = legacy.WGANSecondaryGenerator(
|
|
pdg_vocab=PDG_VOCAB,
|
|
mat_vocab=MAT_VOCAB,
|
|
hidden_dim=HIDDEN_DIM,
|
|
n_blocks=N_BLOCKS,
|
|
emb_dim=EMB_DIM,
|
|
sec_dim=K * SEC_SLOT_DIM,
|
|
noise_dim=16,
|
|
dropout=0.0,
|
|
conditioning=conditioning,
|
|
)
|
|
else:
|
|
old_stage1 = legacy.DenoisingMLP(
|
|
pdg_vocab=PDG_VOCAB,
|
|
mat_vocab=MAT_VOCAB,
|
|
hidden_dim=HIDDEN_DIM,
|
|
n_blocks=N_BLOCKS,
|
|
emb_dim=EMB_DIM,
|
|
dropout=0.0,
|
|
k_max=K,
|
|
conditioning=conditioning,
|
|
)
|
|
old_stage2 = legacy.SecondaryDecoder(
|
|
pdg_vocab=PDG_VOCAB,
|
|
mat_vocab=MAT_VOCAB,
|
|
hidden_dim=HIDDEN_DIM,
|
|
n_blocks=N_BLOCKS,
|
|
emb_dim=EMB_DIM,
|
|
sec_dim=K * SEC_SLOT_DIM,
|
|
dropout=0.0,
|
|
conditioning=conditioning,
|
|
)
|
|
old_stage1.eval()
|
|
old_stage2.eval()
|
|
|
|
cond_cont, cond_cat, x1, x2, t = _random_batch(seed=123)
|
|
z1 = torch.randn(BATCH, 16, generator=torch.Generator().manual_seed(456))
|
|
z2 = torch.randn(BATCH, 16, generator=torch.Generator().manual_seed(789))
|
|
|
|
with torch.no_grad():
|
|
if mode == "wgan":
|
|
old_out1 = old_stage1(z1, cond_cont, cond_cat)
|
|
else:
|
|
old_out1 = old_stage1(x1, t, cond_cont, cond_cat)
|
|
old_n_sec = old_stage1.predict_n_sec(cond_cont, cond_cat)
|
|
if mode == "wgan":
|
|
old_out2 = old_stage2(z2, cond_cont, cond_cat, old_out1)
|
|
else:
|
|
old_out2 = old_stage2(x2, t, cond_cont, cond_cat, old_out1)
|
|
|
|
# --- migrate: config + state dict, through the new build_models ---
|
|
new_models = net.build_models(legacy_cfg)
|
|
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
|
|
# migrated v0.2 checkpoint (design doc §4.1).
|
|
assert new_stage1.n_sec_head is not None
|
|
assert new_stage2.n_sec_head is None
|
|
|
|
remapped1, remapped2 = net.migrate_legacy_state_dict(
|
|
old_stage1.state_dict(), old_stage2.state_dict()
|
|
)
|
|
missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
|
|
missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
|
|
assert not missing1 and not unexpected1
|
|
assert not missing2 and not unexpected2
|
|
new_stage1.eval()
|
|
new_stage2.eval()
|
|
|
|
with torch.no_grad():
|
|
if mode == "wgan":
|
|
new_out1 = new_stage1(z1, cond_cont, cond_cat)
|
|
else:
|
|
new_out1 = new_stage1(x1, cond_cont, cond_cat, t=t)
|
|
new_n_sec = new_stage1.predict_n_sec(cond_cont, cond_cat)
|
|
if mode == "wgan":
|
|
new_out2 = new_stage2(z2, cond_cont, cond_cat, new_out1)
|
|
else:
|
|
new_out2 = new_stage2(x2, cond_cont, cond_cat, new_out1, t=t)
|
|
|
|
_assert_bit_identical(old_out1, new_out1, f"stage1 output ({mode}, {conditioning})")
|
|
_assert_bit_identical(
|
|
old_n_sec, new_n_sec, f"n_sec logits ({mode}, {conditioning})"
|
|
)
|
|
_assert_bit_identical(old_out2, new_out2, f"stage2 output ({mode}, {conditioning})")
|
|
|
|
|
|
def test_migration_flow_embedding():
|
|
_run_migration_check(mode="flow", conditioning="embedding")
|
|
|
|
|
|
def test_migration_flow_physical():
|
|
_run_migration_check(mode="flow", conditioning="physical")
|
|
|
|
|
|
def test_migration_wgan_embedding():
|
|
_run_migration_check(mode="wgan", conditioning="embedding")
|
|
|
|
|
|
def test_migration_wgan_physical():
|
|
_run_migration_check(mode="wgan", conditioning="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
|
|
n_sec head back onto stage 1."""
|
|
legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
|
|
migrated = net._migrate_legacy_model_config(legacy_cfg)
|
|
assert migrated["pdg_vocab"] == PDG_VOCAB
|
|
assert migrated["mat_vocab"] == MAT_VOCAB
|
|
assert migrated["conditioning"]["particle"]["type"] == "physical"
|
|
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"]["decoder"] == "one_shot"
|
|
|
|
|
|
def test_build_models_accepts_new_nested_shape_unchanged():
|
|
"""A dict that already has a 'stage1_model' key (the new shape) is
|
|
passed through build_models without going through the legacy migration
|
|
path at all."""
|
|
cfg = {
|
|
"pdg_vocab": PDG_VOCAB,
|
|
"mat_vocab": MAT_VOCAB,
|
|
"conditioning": {
|
|
"out_dim": 32,
|
|
"particle": {"type": "physical", "emb_dim": EMB_DIM, "n_layers": 1},
|
|
"material": {"type": "physical", "emb_dim": EMB_DIM, "n_layers": 1},
|
|
},
|
|
"stage1_model": {
|
|
"active": True,
|
|
"generator": "flow",
|
|
"hidden_dim": HIDDEN_DIM,
|
|
"n_res_blocks": N_BLOCKS,
|
|
"dropout": 0.0,
|
|
"flow": {"time_dim": 16},
|
|
"router": {"enabled": False},
|
|
},
|
|
"stage2_model": {
|
|
"active": True,
|
|
"decoder": "one_shot",
|
|
"generator": "flow",
|
|
"hidden_dim": HIDDEN_DIM,
|
|
"n_res_blocks": N_BLOCKS,
|
|
"dropout": 0.0,
|
|
"k_max": K,
|
|
"context_dim": 16,
|
|
"n_sec": {"mode": "head"},
|
|
"flow": {"time_dim": 16},
|
|
"router": {"enabled": False, "tie_to_stage1": False},
|
|
},
|
|
}
|
|
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.
|
|
assert models["stage1"].n_sec_head is None
|
|
assert models["stage2"].n_sec_head is not None
|