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>
83 lines
2.5 KiB
Python
83 lines
2.5 KiB
Python
import pytest
|
|
import torch
|
|
from giant.constants import COND_DIM
|
|
from giant.model.network import Stage1Model
|
|
from giant.model.schedule import CosineSchedule, flow_matching_loss
|
|
from giant.sample import sample_flow, sample_ddim
|
|
|
|
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
|
|
_SAMPLE_XFAIL_REASON = (
|
|
"giant/sample.py isn't updated yet — its sample_flow/sample_ddim call "
|
|
"models positionally as model(x, t, cond_cont, cond_cat), which doesn't "
|
|
"match Stage1Model's new forward signature. Deferred to "
|
|
"docs/v0.3.0-design.md step 6."
|
|
)
|
|
|
|
|
|
def _small_model():
|
|
return 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,
|
|
)
|
|
|
|
|
|
def _batch(B=8):
|
|
x1 = torch.randn(B, 9)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
return x1, cond_cont, cond_cat
|
|
|
|
|
|
def test_flow_matching_loss_nonneg():
|
|
x1, cond_cont, cond_cat = _batch()
|
|
loss = flow_matching_loss(_small_model(), x1, cond_cont, cond_cat)
|
|
assert loss.item() >= 0.0
|
|
|
|
|
|
def test_flow_matching_loss_is_scalar():
|
|
x1, cond_cont, cond_cat = _batch()
|
|
loss = flow_matching_loss(_small_model(), x1, cond_cont, cond_cat)
|
|
assert loss.shape == ()
|
|
|
|
|
|
def test_flow_matching_loss_has_grad():
|
|
model = _small_model()
|
|
x1, cond_cont, cond_cat = _batch()
|
|
flow_matching_loss(model, x1, cond_cont, cond_cat).backward()
|
|
assert any(p.grad is not None for p in model.parameters())
|
|
|
|
|
|
@pytest.mark.xfail(reason=_SAMPLE_XFAIL_REASON, strict=False)
|
|
def test_sample_flow_shape():
|
|
B = 6
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
|
|
assert sample.shape == (B, 9)
|
|
assert n_sec.shape == (B,)
|
|
|
|
|
|
def test_ddpm_loss_nonneg():
|
|
schedule = CosineSchedule(T=50)
|
|
x1, cond_cont, cond_cat = _batch()
|
|
loss = schedule.loss(_small_model(), x1, cond_cont, cond_cat)
|
|
assert loss.item() >= 0.0
|
|
|
|
|
|
@pytest.mark.xfail(reason=_SAMPLE_XFAIL_REASON, strict=False)
|
|
def test_sample_ddim_shape():
|
|
B = 4
|
|
schedule = CosineSchedule(T=50)
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
|
sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
|
|
assert sample.shape == (B, 9)
|
|
assert n_sec.shape == (B,)
|