lars
55332db67a
CI / Lint (ruff check) (push) Successful in 31s
CI / Format (ruff format) (push) Successful in 32s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 26s
CI / Type check (ty) (push) Successful in 29s
CI / Format (ruff format) (pull_request) Successful in 33s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 35s
CI / Tests (pull_request) Successful in 3m47s
CI / Tests (push) Successful in 3m55s
Rejoins lines that only wrapped because they exceeded the old 88-char limit; ruff check and the full test suite (725 passed) are unaffected.
284 lines
10 KiB
Python
284 lines
10 KiB
Python
"""Migration acceptance test for v0.3.0 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 — a separate
|
|
portal-machine follow-up with a real checkpoint is planned instead. 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.
|
|
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_migrate_legacy_model_config_nonzero_expert_dims_raises():
|
|
"""Regression: a v0.2 checkpoint's model_config carrying a non-default
|
|
expert_hidden_dim/expert_n_blocks must fail loudly through this path too
|
|
— not just giant.config.migrate_config's parallel TOML-load path.
|
|
Silently dropping these keys (build_router's kwarg filtering) would
|
|
resize the experts instead of refusing."""
|
|
legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
|
|
legacy_cfg["router"] = {
|
|
"enabled": True,
|
|
"expert_hidden_dim": 128,
|
|
"expert_n_blocks": 0,
|
|
}
|
|
try:
|
|
net._migrate_legacy_model_config(legacy_cfg)
|
|
assert False, "expected ValueError"
|
|
except ValueError as e:
|
|
assert "expert_hidden_dim" in str(e)
|
|
|
|
|
|
def test_migrate_legacy_model_config_zero_expert_dims_dropped_silently():
|
|
legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
|
|
legacy_cfg["router"] = {
|
|
"enabled": True,
|
|
"expert_hidden_dim": 0,
|
|
"expert_n_blocks": 0,
|
|
}
|
|
migrated = net._migrate_legacy_model_config(legacy_cfg)
|
|
assert "expert_hidden_dim" not in migrated["stage1_model"]["router"]
|
|
assert "expert_n_blocks" not in migrated["stage1_model"]["router"]
|
|
assert "expert_hidden_dim" not in migrated["stage2_model"]["router"]
|
|
assert "expert_n_blocks" not in migrated["stage2_model"]["router"]
|
|
|
|
|
|
def test_build_models_with_legacy_config_nonzero_expert_dims_raises():
|
|
"""The same check must also fire through the actual caller,
|
|
build_models, not just the internal helper directly."""
|
|
legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
|
|
legacy_cfg["router"] = {
|
|
"enabled": True,
|
|
"expert_hidden_dim": 128,
|
|
"expert_n_blocks": 0,
|
|
}
|
|
try:
|
|
net.build_models(legacy_cfg)
|
|
assert False, "expected ValueError"
|
|
except ValueError as e:
|
|
assert "expert_hidden_dim" in str(e)
|
|
|
|
|
|
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
|