"""Portal-machine follow-up for v0.3.0 step 2: diff a real v0.2 checkpoint's outputs against the new `build_models` on the same input batch. `tests/test_migration_v02_v03.py` already proves this bit-identical with synthetic random weights, but that test can't run where it matters (no `/ceph` on local dev machines — see CLAUDE.md's Compute environment section). This script is the real-checkpoint counterpart: run it on a portal machine against an actual trained checkpoint before merging `v0.3.0-stage2-autoregressive` to `master`. Usage (from the repo root, on a portal machine): uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --ema uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --batch 32 --seed 1 Run it once against a flow (or ddpm) checkpoint and once against a wgan checkpoint ("one flow checkpoint and one WGAN checkpoint"). A routed checkpoint (`model_config["router"]["enabled"]`) is only checked for successful construction — `giant.model.network.migrate_legacy_state_dict` doesn't yet remap routed (Expert-per-router) state dicts, so the bit-identical assertion is skipped with a clear warning in that case (see the function's own docstring for why). """ import argparse import sys from pathlib import Path import torch sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM # noqa: E402 from giant.model import network as net # noqa: E402 from tests.legacy import network_v02_snapshot as legacy # noqa: E402 def _random_batch(model_config: dict, batch: int, seed: int): g = torch.Generator().manual_seed(seed) pdg_vocab = model_config["pdg_vocab"] mat_vocab = model_config["mat_vocab"] k_max = model_config.get("k_max", 15) noise_dim = model_config.get("noise_dim", 64) cond_cont = torch.randn(batch, COND_DIM, generator=g) cond_cat = torch.stack( [ torch.randint(0, pdg_vocab, (batch,), generator=g), torch.randint(0, mat_vocab, (batch,), generator=g), ], dim=1, ) x1 = torch.randn(batch, X_DIM, generator=g) x2 = torch.randn(batch, k_max * SEC_SLOT_DIM, generator=g) t = torch.rand(batch, generator=g) z1 = torch.randn(batch, noise_dim, generator=g) z2 = torch.randn(batch, noise_dim, generator=g) return cond_cont, cond_cat, x1, x2, t, z1, z2 def _max_abs_diff(a: torch.Tensor, b: torch.Tensor) -> float: return (a - b).abs().max().item() def main() -> int: p = argparse.ArgumentParser(description=__doc__) p.add_argument("checkpoint", type=Path, help="Path to a v0.2 best.pt/last.pt") p.add_argument( "--ema", action="store_true", help="Use the checkpoint's EMA weights (model_ema/sec_decoder_ema) — " "what predict/rollout actually sample from — instead of raw weights.", ) p.add_argument("--batch", type=int, default=16) p.add_argument("--seed", type=int, default=0) args = p.parse_args() ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) if "model_config" not in ckpt: print(f"FAIL: {args.checkpoint} has no 'model_config' key — can't migrate it") return 1 model_config = ckpt["model_config"] mode = model_config.get("mode", "flow") routed = bool((model_config.get("router") or {}).get("enabled")) print(f"checkpoint: {args.checkpoint}") print( f" mode={mode!r} conditioning={model_config.get('conditioning')!r} " f"routed={routed} ema={args.ema}" ) stage1_key = "model_ema" if args.ema and "model_ema" in ckpt else "model" stage2_key = ( "sec_decoder_ema" if args.ema and "sec_decoder_ema" in ckpt else "sec_decoder" ) if args.ema and stage1_key == "model": print( " warning: --ema requested but no model_ema in checkpoint, using raw weights" ) # --- old side: the frozen v0.2 snapshot, loaded with the checkpoint's own weights --- old_stage1, old_stage2 = legacy.build_models(model_config) old_stage1.load_state_dict(ckpt[stage1_key]) old_stage2.load_state_dict(ckpt[stage2_key]) old_stage1.eval() old_stage2.eval() # --- new side: migrated config + remapped state dict, through the new build_models --- new_models = net.build_models(model_config) new_stage1, new_stage2 = new_models["stage1"], new_models["stage2"] assert new_stage1 is not None and new_stage2 is not None if routed: print( " routed checkpoint: migrate_legacy_state_dict only handles the " "monolithic trunk shape — verifying construction only, skipping " "the bit-identical weight/output comparison." ) print("PASS (construction only, routed checkpoint)") return 0 remapped1, remapped2 = net.migrate_legacy_state_dict( ckpt[stage1_key], ckpt[stage2_key] ) missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True) missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True) if missing1 or unexpected1 or missing2 or unexpected2: print("FAIL: state dict mismatch after remap") print(f" stage1 missing={missing1} unexpected={unexpected1}") print(f" stage2 missing={missing2} unexpected={unexpected2}") return 1 new_stage1.eval() new_stage2.eval() cond_cont, cond_cat, x1, x2, t, z1, z2 = _random_batch( model_config, args.batch, args.seed ) ok = True with torch.no_grad(): if mode == "wgan": old_out1 = old_stage1(z1, cond_cont, cond_cat) new_out1 = new_stage1(z1, cond_cont, cond_cat) else: old_out1 = old_stage1(x1, t, cond_cont, cond_cat) new_out1 = new_stage1(x1, cond_cont, cond_cat, t=t) old_n_sec = old_stage1.predict_n_sec(cond_cont, cond_cat) new_n_sec = new_stage1.predict_n_sec(cond_cont, cond_cat) if mode == "wgan": old_out2 = old_stage2(z2, cond_cont, cond_cat, old_out1) new_out2 = new_stage2(z2, cond_cont, cond_cat, new_out1) else: old_out2 = old_stage2(x2, t, cond_cont, cond_cat, old_out1) new_out2 = new_stage2(x2, cond_cont, cond_cat, new_out1, t=t) for label, old_out, new_out in [ ("stage1 output", old_out1, new_out1), ("n_sec logits", old_n_sec, new_n_sec), ("stage2 output", old_out2, new_out2), ]: identical = torch.equal(old_out, new_out) diff = _max_abs_diff(old_out, new_out) status = "OK" if identical else "MISMATCH" print(f" {label}: {status} (max abs diff = {diff:.3e})") ok = ok and identical print("PASS" if ok else "FAIL") return 0 if ok else 1 if __name__ == "__main__": raise SystemExit(main())