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>
74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
import torch
|
|
from giant.constants import COND_DIM
|
|
from giant.model.network import SinusoidalEmbedding, Stage1Model
|
|
|
|
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
|
|
|
|
def test_sinusoidal_embedding_shape():
|
|
emb = SinusoidalEmbedding(64)
|
|
t = torch.rand(16)
|
|
assert emb(t).shape == (16, 64)
|
|
|
|
|
|
def test_sinusoidal_embedding_batch_1():
|
|
emb = SinusoidalEmbedding(32)
|
|
t = torch.tensor([0.5])
|
|
assert emb(t).shape == (1, 32)
|
|
|
|
|
|
def test_stage1_model_output_shape():
|
|
B = 8
|
|
model = Stage1Model(
|
|
pdg_vocab=5,
|
|
mat_vocab=3,
|
|
particle_cfg=PARTICLE_CFG,
|
|
material_cfg=MATERIAL_CFG,
|
|
n_sec_head_k_max=15,
|
|
)
|
|
x_t = torch.randn(B, 9)
|
|
t = torch.rand(B)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.stack(
|
|
[
|
|
torch.randint(0, 5, (B,)),
|
|
torch.randint(0, 3, (B,)),
|
|
],
|
|
dim=1,
|
|
)
|
|
out = model(x_t, cond_cont, cond_cat, t=t)
|
|
assert out.shape == (B, 9)
|
|
|
|
|
|
def test_stage1_model_gradients_flow():
|
|
B = 4
|
|
model = Stage1Model(
|
|
pdg_vocab=3,
|
|
mat_vocab=2,
|
|
particle_cfg=PARTICLE_CFG,
|
|
material_cfg=MATERIAL_CFG,
|
|
hidden_dim=32,
|
|
n_res_blocks=2,
|
|
n_sec_head_k_max=15,
|
|
)
|
|
x_t = torch.randn(B, 9)
|
|
t = torch.rand(B)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
# Both paths must be exercised to get gradients through all parameters.
|
|
flow_loss = model(x_t, cond_cont, cond_cat, t=t).sum()
|
|
nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
|
|
(flow_loss + nsec_loss).backward()
|
|
for name, p in model.named_parameters():
|
|
assert p.grad is not None, f"no grad for {name}"
|
|
|
|
|
|
def test_stage1_model_no_n_sec_head_by_default():
|
|
"""Fresh v0.3.0 construction (no n_sec_head_k_max) has no n_sec head —
|
|
decision 1 (docs/v0.3.0-design.md §2) moves it to stage 2."""
|
|
model = Stage1Model(
|
|
pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG
|
|
)
|
|
assert model.n_sec_head is None
|