lars
5b63dfd588
The whole point of conditioning="physical" is generalizing to a species/material outside the training menu, but two independent code paths still hard-required training-vocab membership: - giant/data/transforms.py: build_cond_features unconditionally raised KeyError on an out-of-vocab pdg/material. _vectorized_map_lookup gains a strict=False mode (dummy index instead of raising), used only under conditioning="physical" where ConditionEncoder never reads cond_cat anyway; "embedding" mode is untouched and still raises, since cond_cat IS the conditioning signal there. - giant/rollout.py: the known_pdg termination gate still killed a track on step 1 for any pdg outside pdg_map, regardless of conditioning mode. Now skipped entirely under conditioning="physical". - giant/model/network.py: PdgRouter/ProcessRouter always build their own training-vocab nn.Embedding independent of conditioning, silently reintroducing the same limitation at the routing layer. build_models now raises loudly if conditioning="physical" is paired with either router type, rather than silently building a model that can't generalize the way it claims to. This unblocks the held-out-species/material generalization experiment against the multi-material dataset (see CLAUDE.md roadmap). Each fix has a regression test, including an end-to-end rollout test seeded with a resolvable-but-out-of-vocab PDG code. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1032 lines
34 KiB
Python
1032 lines
34 KiB
Python
"""Tests for the mixture-of-experts routing prototype (giant/model/network.py)."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
|
|
from giant.model.network import (
|
|
ComposedRouter,
|
|
DenoisingMLP,
|
|
EnergyRouter,
|
|
PdgRouter,
|
|
ProcessRouter,
|
|
ROUTER_REGISTRY,
|
|
RoutedDenoisingMLP,
|
|
RoutedSecondaryDecoder,
|
|
SecondaryDecoder,
|
|
build_composed_router,
|
|
build_models,
|
|
build_router,
|
|
)
|
|
|
|
|
|
def _cond(B=8, pdg=3, mat=2):
|
|
cond_cont = torch.randn(B, COND_DIM)
|
|
cond_cat = torch.stack(
|
|
[torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
|
|
)
|
|
return cond_cont, cond_cat
|
|
|
|
|
|
def _routed_stage1(n_experts=4, pdg=3, mat=2, **router_kwargs):
|
|
router = build_router("energy", n_experts, **router_kwargs)
|
|
return RoutedDenoisingMLP(
|
|
pdg_vocab=pdg,
|
|
mat_vocab=mat,
|
|
router=router,
|
|
expert_hidden_dim=16,
|
|
expert_n_blocks=2,
|
|
)
|
|
|
|
|
|
def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs):
|
|
router = build_router("energy", n_experts, **router_kwargs)
|
|
return RoutedSecondaryDecoder(
|
|
pdg_vocab=pdg,
|
|
mat_vocab=mat,
|
|
router=router,
|
|
expert_hidden_dim=16,
|
|
expert_n_blocks=2,
|
|
)
|
|
|
|
|
|
# ── Router / EnergyRouter contract ──────────────────────────────────────────
|
|
|
|
|
|
def test_energy_router_registered():
|
|
assert ROUTER_REGISTRY["energy"] is EnergyRouter
|
|
|
|
|
|
def test_energy_router_gate_partition_of_unity():
|
|
router = EnergyRouter(n_experts=4)
|
|
cond_cont, cond_cat = _cond(16)
|
|
g = router.gate(cond_cont, cond_cat)
|
|
assert g.shape == (16, 4)
|
|
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
|
|
|
|
|
def test_energy_router_top1_matches_gate_argmax():
|
|
router = EnergyRouter(n_experts=4)
|
|
cond_cont, cond_cat = _cond(16)
|
|
assert torch.equal(
|
|
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
|
|
)
|
|
|
|
|
|
def test_energy_router_hardens_as_temperature_shrinks():
|
|
"""As tau -> 0 the soft gate should converge to a one-hot at the argmax."""
|
|
router = EnergyRouter(n_experts=4, temperature=1e-4)
|
|
cond_cont, cond_cat = _cond(16)
|
|
g = router.gate(cond_cont, cond_cat)
|
|
top1 = router.top1(cond_cont, cond_cat)
|
|
onehot = torch.nn.functional.one_hot(top1, num_classes=4).float()
|
|
torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
|
|
|
|
|
|
def test_energy_router_balance_loss_is_nonnegative_scalar():
|
|
router = EnergyRouter(n_experts=4)
|
|
cond_cont, cond_cat = _cond(16)
|
|
loss = router.balance_loss(cond_cont, cond_cat)
|
|
assert loss.shape == ()
|
|
assert loss.item() >= 0.0
|
|
|
|
|
|
def test_build_router_ignores_unrecognized_kwargs():
|
|
# lambda_balance is a model_config.router key but not an EnergyRouter kwarg
|
|
router = build_router("energy", 4, temperature=0.3, lambda_balance=0.5)
|
|
assert isinstance(router, EnergyRouter)
|
|
assert router.temperature == 0.3
|
|
|
|
|
|
def test_energy_router_default_centers_are_linspace():
|
|
router = EnergyRouter(n_experts=4)
|
|
torch.testing.assert_close(router.centers, torch.linspace(-2.0, 2.0, 4))
|
|
|
|
|
|
def test_energy_router_centers_init_overrides_default():
|
|
centers_init = [-1.0, 0.0, 0.5, 3.0]
|
|
router = EnergyRouter(n_experts=4, centers_init=centers_init)
|
|
torch.testing.assert_close(router.centers, torch.tensor(centers_init))
|
|
|
|
|
|
def test_energy_router_centers_init_wrong_length_raises():
|
|
try:
|
|
EnergyRouter(n_experts=4, centers_init=[0.0, 1.0])
|
|
except ValueError:
|
|
return
|
|
raise AssertionError("expected ValueError for centers_init length mismatch")
|
|
|
|
|
|
def test_energy_router_centers_init_respects_learn_centers_flag():
|
|
learned = EnergyRouter(
|
|
n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True
|
|
)
|
|
fixed = EnergyRouter(
|
|
n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False
|
|
)
|
|
assert isinstance(learned.centers, torch.nn.Parameter)
|
|
assert not isinstance(fixed.centers, torch.nn.Parameter)
|
|
|
|
|
|
def test_build_router_threads_centers_init_through_energy_router():
|
|
centers_init = [-1.5, -0.5, 0.5, 1.5]
|
|
router = build_router("energy", 4, centers_init=centers_init)
|
|
torch.testing.assert_close(router.centers, torch.tensor(centers_init))
|
|
|
|
|
|
def test_build_router_unknown_type_raises():
|
|
try:
|
|
build_router("nonexistent", 4)
|
|
except ValueError:
|
|
return
|
|
raise AssertionError("expected ValueError for unknown router type")
|
|
|
|
|
|
# ── EnergyRouter learn_width / learn_temperature ────────────────────────────
|
|
|
|
|
|
def test_energy_router_learn_width_matches_fixed_temperature_at_init():
|
|
"""Enabling learn_width should be a no-op at init — the warm-started
|
|
per-expert width must reproduce the fixed-temperature gate exactly."""
|
|
centers_init = [-1.0, 0.0, 0.5, 1.5]
|
|
fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
|
|
learned = EnergyRouter(
|
|
n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True
|
|
)
|
|
cond_cont, cond_cat = _cond(16)
|
|
torch.testing.assert_close(
|
|
learned.gate(cond_cont, cond_cat),
|
|
fixed.gate(cond_cont, cond_cat),
|
|
atol=1e-5,
|
|
rtol=0,
|
|
)
|
|
|
|
|
|
def test_energy_router_learn_temperature_matches_fixed_temperature_at_init():
|
|
centers_init = [-1.0, 0.0, 0.5, 1.5]
|
|
fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
|
|
learned = EnergyRouter(
|
|
n_experts=4,
|
|
temperature=0.3,
|
|
centers_init=centers_init,
|
|
learn_temperature=True,
|
|
)
|
|
cond_cont, cond_cat = _cond(16)
|
|
torch.testing.assert_close(
|
|
learned.gate(cond_cont, cond_cat),
|
|
fixed.gate(cond_cont, cond_cat),
|
|
atol=1e-5,
|
|
rtol=0,
|
|
)
|
|
|
|
|
|
def test_energy_router_learn_width_and_temperature_mutually_exclusive_raises():
|
|
try:
|
|
EnergyRouter(n_experts=4, learn_width=True, learn_temperature=True)
|
|
except ValueError:
|
|
return
|
|
raise AssertionError(
|
|
"expected ValueError for learn_width and learn_temperature both set"
|
|
)
|
|
|
|
|
|
def test_energy_router_width_ratio_bounds_must_bracket_one_raises():
|
|
try:
|
|
EnergyRouter(
|
|
n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0
|
|
)
|
|
except ValueError:
|
|
return
|
|
raise AssertionError(
|
|
"expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0"
|
|
)
|
|
|
|
|
|
def test_energy_router_effective_width_stays_within_bounds():
|
|
router = EnergyRouter(
|
|
n_experts=4,
|
|
temperature=0.5,
|
|
learn_width=True,
|
|
width_min_ratio=0.1,
|
|
width_max_ratio=10.0,
|
|
)
|
|
lo, hi = 0.1 * 0.5, 10.0 * 0.5
|
|
with torch.no_grad():
|
|
router.raw_width.fill_(1e6)
|
|
width = router.effective_width()
|
|
assert isinstance(width, torch.Tensor)
|
|
assert torch.all(width <= hi + 1e-4)
|
|
with torch.no_grad():
|
|
router.raw_width.fill_(-1e6)
|
|
width = router.effective_width()
|
|
assert isinstance(width, torch.Tensor)
|
|
assert torch.all(width >= lo - 1e-4)
|
|
|
|
|
|
def test_energy_router_learn_width_gate_still_partition_of_unity():
|
|
router = EnergyRouter(n_experts=4, learn_width=True)
|
|
with torch.no_grad():
|
|
router.raw_width.copy_(torch.randn(4) * 3)
|
|
cond_cont, cond_cat = _cond(16)
|
|
g = router.gate(cond_cont, cond_cat)
|
|
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
|
|
|
|
|
def test_energy_router_learn_width_hardens_when_pushed_to_floor():
|
|
"""Pushing every expert's width toward the (tiny) floor should harden the
|
|
gate to a one-hot at the nearest center, generalizing the fixed-
|
|
temperature->0 hardening test to the per-expert path."""
|
|
router = EnergyRouter(
|
|
n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0
|
|
)
|
|
with torch.no_grad():
|
|
router.raw_width.fill_(-1e6)
|
|
cond_cont, cond_cat = _cond(16)
|
|
g = router.gate(cond_cont, cond_cat)
|
|
e = cond_cont[:, router.energy_idx].unsqueeze(-1)
|
|
d2 = (e - router.centers.unsqueeze(0)) ** 2
|
|
onehot = torch.nn.functional.one_hot(d2.argmin(dim=-1), num_classes=4).float()
|
|
torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
|
|
|
|
|
|
def test_energy_router_own_width_controls_own_coverage_independent_of_others():
|
|
"""Widening one expert's width should monotonically grow only that
|
|
expert's own gate share, without needing to touch any other expert's
|
|
width — the "each expert learns its own coverage independently" property
|
|
this feature is meant to add."""
|
|
router = EnergyRouter(
|
|
n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0]
|
|
)
|
|
cond_cont, cond_cat = _cond(4)
|
|
cond_cont[:, 3] = 3.0 # fixed energy, unequal distance to each center
|
|
|
|
shares = []
|
|
with torch.no_grad():
|
|
for raw in torch.linspace(-8.0, 8.0, 9):
|
|
router.raw_width[0] = raw
|
|
shares.append(router.gate(cond_cont, cond_cat)[0, 0].item())
|
|
|
|
assert all(a <= b + 1e-6 for a, b in zip(shares, shares[1:]))
|
|
|
|
|
|
def test_build_router_threads_learn_width_kwargs_through():
|
|
router = build_router(
|
|
"energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0
|
|
)
|
|
assert isinstance(router, EnergyRouter)
|
|
assert router.learn_width is True
|
|
assert isinstance(router.raw_width, torch.nn.Parameter)
|
|
assert router.raw_width.shape == (4,)
|
|
|
|
|
|
def test_router_entropy_loss_is_nonnegative_bounded_scalar():
|
|
router = EnergyRouter(n_experts=4)
|
|
cond_cont, cond_cat = _cond(16)
|
|
loss = router.entropy_loss(cond_cont, cond_cat)
|
|
assert loss.shape == ()
|
|
assert 0.0 <= loss.item() <= 1.0
|
|
|
|
|
|
# ── Router.combine_weights (straight-through Gumbel-softmax) ───────────────
|
|
|
|
|
|
def test_combine_weights_defaults_to_gate():
|
|
"""gumbel=False (the default) must be a pure pass-through to gate()."""
|
|
router = EnergyRouter(n_experts=4)
|
|
cond_cont, cond_cat = _cond(16)
|
|
torch.testing.assert_close(
|
|
router.combine_weights(cond_cont, cond_cat),
|
|
router.gate(cond_cont, cond_cat),
|
|
)
|
|
|
|
|
|
def test_combine_weights_gumbel_train_mode_is_hard_one_hot():
|
|
router = EnergyRouter(n_experts=4)
|
|
router.gumbel = True
|
|
router.gumbel_tau = 0.5
|
|
router.train()
|
|
cond_cont, cond_cat = _cond(16)
|
|
weights = router.combine_weights(cond_cont, cond_cat)
|
|
assert weights.shape == (16, 4)
|
|
torch.testing.assert_close(weights.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
|
assert torch.all((weights.max(dim=-1).values - 1.0).abs() < 1e-5)
|
|
|
|
|
|
def test_combine_weights_gumbel_eval_mode_falls_back_to_gate():
|
|
"""No Gumbel noise at eval — combine_weights must match gate() exactly,
|
|
same as the gumbel=False path, once the router is in eval mode."""
|
|
router = EnergyRouter(n_experts=4)
|
|
router.gumbel = True
|
|
router.eval()
|
|
cond_cont, cond_cat = _cond(16)
|
|
torch.testing.assert_close(
|
|
router.combine_weights(cond_cont, cond_cat),
|
|
router.gate(cond_cont, cond_cat),
|
|
)
|
|
|
|
|
|
def test_combine_weights_gumbel_straight_through_gradient_reaches_centers():
|
|
router = EnergyRouter(n_experts=4, learn_centers=True)
|
|
router.gumbel = True
|
|
router.gumbel_tau = 0.5
|
|
router.train()
|
|
cond_cont, cond_cat = _cond(16)
|
|
weights = router.combine_weights(cond_cont, cond_cat)
|
|
weights.sum().backward()
|
|
assert router.centers.grad is not None
|
|
assert torch.any(router.centers.grad != 0.0)
|
|
|
|
|
|
def test_build_router_from_cfg_sets_gumbel_from_config():
|
|
from giant.model.network import _build_router_from_cfg
|
|
|
|
router = _build_router_from_cfg(
|
|
{"enabled": True, "type": "energy", "n_experts": 4, "gumbel": True},
|
|
pdg_vocab=3,
|
|
mat_vocab=2,
|
|
)
|
|
assert router.gumbel is True
|
|
|
|
router_off = _build_router_from_cfg(
|
|
{"enabled": True, "type": "energy", "n_experts": 4},
|
|
pdg_vocab=3,
|
|
mat_vocab=2,
|
|
)
|
|
assert router_off.gumbel is False
|
|
|
|
|
|
def test_build_router_from_cfg_sets_gumbel_for_composed_router():
|
|
from giant.model.network import _build_router_from_cfg
|
|
|
|
router = _build_router_from_cfg(
|
|
{
|
|
"enabled": True,
|
|
"type": "composed",
|
|
"gumbel": True,
|
|
"axis0_type": "energy",
|
|
"axis0_n_experts": 4,
|
|
"axis1_type": "pdg",
|
|
"axis1_n_experts": 3,
|
|
},
|
|
pdg_vocab=5,
|
|
mat_vocab=2,
|
|
)
|
|
assert isinstance(router, ComposedRouter)
|
|
assert router.gumbel is True
|
|
|
|
|
|
def test_routed_denoising_mlp_forward_runs_with_gumbel_enabled():
|
|
"""End-to-end forward through _route_forward's train branch with
|
|
straight-through Gumbel-softmax combine weights enabled."""
|
|
B = 8
|
|
model = _routed_stage1(n_experts=3)
|
|
model.router.gumbel = True
|
|
model.router.gumbel_tau = 0.5
|
|
model.train()
|
|
x_t = torch.randn(B, X_DIM)
|
|
t = torch.rand(B)
|
|
cond_cont, cond_cat = _cond(B)
|
|
out = model(x_t, t, cond_cont, cond_cat)
|
|
assert out.shape == (B, X_DIM)
|
|
assert torch.isfinite(out).all()
|
|
|
|
|
|
# ── PdgRouter ────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_pdg_router_registered():
|
|
assert ROUTER_REGISTRY["pdg"] is PdgRouter
|
|
|
|
|
|
def test_pdg_router_gate_partition_of_unity():
|
|
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
|
cond_cont, cond_cat = _cond(16)
|
|
g = router.gate(cond_cont, cond_cat)
|
|
assert g.shape == (16, 4)
|
|
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
|
|
|
|
|
def test_pdg_router_top1_matches_gate_argmax():
|
|
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
|
cond_cont, cond_cat = _cond(16)
|
|
assert torch.equal(
|
|
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
|
|
)
|
|
|
|
|
|
def test_pdg_router_hardens_as_temperature_shrinks():
|
|
"""As tau -> 0 the soft gate should converge to a one-hot at the argmax."""
|
|
router = PdgRouter(n_experts=4, pdg_vocab=3, temperature=1e-4)
|
|
cond_cont, cond_cat = _cond(16)
|
|
g = router.gate(cond_cont, cond_cat)
|
|
top1 = router.top1(cond_cont, cond_cat)
|
|
onehot = torch.nn.functional.one_hot(top1, num_classes=4).float()
|
|
torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
|
|
|
|
|
|
def test_pdg_router_balance_loss_is_nonnegative_scalar():
|
|
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
|
cond_cont, cond_cat = _cond(16)
|
|
loss = router.balance_loss(cond_cont, cond_cat)
|
|
assert loss.shape == ()
|
|
assert loss.item() >= 0.0
|
|
|
|
|
|
def test_pdg_router_classify_loss_defaults_to_zero():
|
|
"""PDG is already known at gate time (unlike ProcessRouter's process
|
|
label), so no supervision is needed — falls back to Router's default."""
|
|
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
|
cond_cont, cond_cat = _cond(16)
|
|
labels = torch.randint(0, 4, (16,))
|
|
loss = router.classify_loss(cond_cont, cond_cat, labels)
|
|
assert loss.shape == ()
|
|
assert loss.item() == 0.0
|
|
|
|
|
|
def test_pdg_router_only_reads_pdg_column():
|
|
"""Gate must depend on cond_cat[:, 0] (pdg) only, not cond_cont or material."""
|
|
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
|
cond_cont, cond_cat = _cond(16)
|
|
g_before = router.gate(cond_cont, cond_cat)
|
|
|
|
cond_cont_perturbed = torch.randn_like(cond_cont)
|
|
cond_cat_diff_mat = cond_cat.clone()
|
|
cond_cat_diff_mat[:, 1] = (cond_cat_diff_mat[:, 1] + 1) % 2
|
|
g_after = router.gate(cond_cont_perturbed, cond_cat_diff_mat)
|
|
|
|
torch.testing.assert_close(g_before, g_after, atol=1e-6, rtol=0)
|
|
|
|
|
|
def test_build_router_pdg_type_uses_pdg_vocab():
|
|
router = build_router("pdg", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8)
|
|
assert isinstance(router, PdgRouter)
|
|
assert router.pdg_emb.num_embeddings == 5
|
|
|
|
|
|
def test_build_models_routed_with_pdg_router():
|
|
model_config = dict(
|
|
pdg_vocab=4,
|
|
mat_vocab=2,
|
|
emb_dim=16,
|
|
dropout=0.1,
|
|
k_max=K_MAX,
|
|
expert_hidden_dim=16,
|
|
expert_n_blocks=2,
|
|
router={
|
|
"enabled": True,
|
|
"type": "pdg",
|
|
"n_experts": 3,
|
|
},
|
|
)
|
|
stage1, sec_decoder = build_models(model_config)
|
|
assert isinstance(stage1, RoutedDenoisingMLP)
|
|
assert isinstance(stage1.router, PdgRouter)
|
|
assert len(stage1.experts) == 3
|
|
assert stage1.router.pdg_emb.num_embeddings == 4
|
|
|
|
|
|
def test_build_models_rejects_pdg_router_with_physical_conditioning():
|
|
"""conditioning="physical" is meant to generalize beyond the training PDG
|
|
vocab; PdgRouter always uses a training-vocab nn.Embedding regardless of
|
|
conditioning, so the combination must raise rather than silently building
|
|
a model that can't actually generalize the way it claims to."""
|
|
model_config = dict(
|
|
pdg_vocab=4,
|
|
mat_vocab=2,
|
|
emb_dim=16,
|
|
dropout=0.1,
|
|
k_max=K_MAX,
|
|
expert_hidden_dim=16,
|
|
expert_n_blocks=2,
|
|
conditioning="physical",
|
|
router={"enabled": True, "type": "pdg", "n_experts": 3},
|
|
)
|
|
with pytest.raises(ValueError, match="physical"):
|
|
build_models(model_config)
|
|
|
|
|
|
def test_build_models_rejects_composed_router_with_pdg_axis_and_physical_conditioning():
|
|
model_config = dict(
|
|
pdg_vocab=4,
|
|
mat_vocab=2,
|
|
emb_dim=16,
|
|
dropout=0.1,
|
|
k_max=K_MAX,
|
|
expert_hidden_dim=16,
|
|
expert_n_blocks=2,
|
|
conditioning="physical",
|
|
router={
|
|
"enabled": True,
|
|
"type": "composed",
|
|
"axis0_type": "energy",
|
|
"axis0_n_experts": 2,
|
|
"axis1_type": "pdg",
|
|
"axis1_n_experts": 3,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="physical"):
|
|
build_models(model_config)
|
|
|
|
|
|
# ── ProcessRouter ────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_process_router_registered():
|
|
assert ROUTER_REGISTRY["process"] is ProcessRouter
|
|
|
|
|
|
def test_process_router_gate_partition_of_unity():
|
|
router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
|
|
cond_cont, cond_cat = _cond(16)
|
|
g = router.gate(cond_cont, cond_cat)
|
|
assert g.shape == (16, 4)
|
|
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
|
|
|
|
|
def test_process_router_top1_matches_gate_argmax():
|
|
router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
|
|
cond_cont, cond_cat = _cond(16)
|
|
assert torch.equal(
|
|
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
|
|
)
|
|
|
|
|
|
def test_process_router_balance_loss_is_nonnegative_scalar():
|
|
router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
|
|
cond_cont, cond_cat = _cond(16)
|
|
loss = router.balance_loss(cond_cont, cond_cat)
|
|
assert loss.shape == ()
|
|
assert loss.item() >= 0.0
|
|
|
|
|
|
def test_process_router_classify_loss_decreases_with_training():
|
|
"""The classifier should be able to fit an arbitrary label assignment —
|
|
a sanity check that gradients actually flow to the process classifier."""
|
|
torch.manual_seed(0)
|
|
router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
|
|
cond_cont, cond_cat = _cond(32)
|
|
labels = torch.randint(0, 4, (32,))
|
|
|
|
opt = torch.optim.Adam(router.parameters(), lr=0.05)
|
|
first = router.classify_loss(cond_cont, cond_cat, labels).item()
|
|
for _ in range(50):
|
|
opt.zero_grad()
|
|
loss = router.classify_loss(cond_cont, cond_cat, labels)
|
|
loss.backward()
|
|
opt.step()
|
|
last = loss.item()
|
|
assert last < first
|
|
|
|
|
|
def test_energy_router_classify_loss_defaults_to_zero():
|
|
"""Routers with no supervised signal (EnergyRouter) fall back to the
|
|
Router base class's zero-loss default."""
|
|
router = EnergyRouter(n_experts=4)
|
|
cond_cont, cond_cat = _cond(16)
|
|
labels = torch.randint(0, 4, (16,))
|
|
loss = router.classify_loss(cond_cont, cond_cat, labels)
|
|
assert loss.shape == ()
|
|
assert loss.item() == 0.0
|
|
|
|
|
|
def test_build_router_process_type_uses_pdg_mat_vocab():
|
|
router = build_router("process", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8)
|
|
assert isinstance(router, ProcessRouter)
|
|
assert router.pdg_emb.num_embeddings == 5
|
|
assert router.mat_emb.num_embeddings == 3
|
|
|
|
|
|
def test_build_models_routed_with_process_router():
|
|
model_config = dict(
|
|
pdg_vocab=4,
|
|
mat_vocab=2,
|
|
emb_dim=16,
|
|
dropout=0.1,
|
|
k_max=K_MAX,
|
|
expert_hidden_dim=16,
|
|
expert_n_blocks=2,
|
|
router={
|
|
"enabled": True,
|
|
"type": "process",
|
|
"n_experts": 3,
|
|
"lambda_proc": 1.0,
|
|
},
|
|
)
|
|
stage1, sec_decoder = build_models(model_config)
|
|
assert isinstance(stage1, RoutedDenoisingMLP)
|
|
assert isinstance(stage1.router, ProcessRouter)
|
|
assert len(stage1.experts) == 3
|
|
assert stage1.router.pdg_emb.num_embeddings == 4
|
|
assert stage1.router.mat_emb.num_embeddings == 2
|
|
|
|
|
|
# ── ComposedRouter ───────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_composed_router_n_experts_is_product():
|
|
router = ComposedRouter(
|
|
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
|
)
|
|
assert router.n_experts == 12
|
|
|
|
|
|
def test_composed_router_gate_partition_of_unity():
|
|
router = ComposedRouter(
|
|
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
|
)
|
|
cond_cont, cond_cat = _cond(16, pdg=5)
|
|
g = router.gate(cond_cont, cond_cat)
|
|
assert g.shape == (16, 12)
|
|
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
|
|
|
|
|
def test_composed_router_gate_is_outer_product_of_sub_gates():
|
|
energy_router = EnergyRouter(n_experts=4)
|
|
pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
|
|
router = ComposedRouter([energy_router, pdg_router])
|
|
cond_cont, cond_cat = _cond(16, pdg=5)
|
|
|
|
g_energy = energy_router.gate(cond_cont, cond_cat) # (16, 4)
|
|
g_pdg = pdg_router.gate(cond_cont, cond_cat) # (16, 3)
|
|
expected = (g_energy.unsqueeze(-1) * g_pdg.unsqueeze(1)).flatten(1) # (16, 12)
|
|
|
|
torch.testing.assert_close(router.gate(cond_cont, cond_cat), expected)
|
|
|
|
|
|
def test_composed_router_top1_factors_into_per_axis_argmax():
|
|
"""Joint argmax over the outer product must equal the pair of per-axis
|
|
argmaxes, flattened with the same row-major index convention as gate()."""
|
|
energy_router = EnergyRouter(n_experts=4)
|
|
pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
|
|
router = ComposedRouter([energy_router, pdg_router])
|
|
cond_cont, cond_cat = _cond(16, pdg=5)
|
|
|
|
joint_idx = router.top1(cond_cont, cond_cat)
|
|
energy_idx = energy_router.top1(cond_cont, cond_cat)
|
|
pdg_idx = pdg_router.top1(cond_cont, cond_cat)
|
|
expected = energy_idx * pdg_router.n_experts + pdg_idx
|
|
|
|
assert torch.equal(joint_idx, expected)
|
|
|
|
|
|
def test_composed_router_supports_different_expert_counts_per_axis():
|
|
router = ComposedRouter(
|
|
[EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)]
|
|
)
|
|
assert router.n_experts == 10
|
|
cond_cont, cond_cat = _cond(8, pdg=5)
|
|
assert router.gate(cond_cont, cond_cat).shape == (8, 10)
|
|
|
|
|
|
def test_composed_router_classify_loss_sums_sub_router_losses():
|
|
"""energy/pdg both default to zero, so the composed loss should too."""
|
|
router = ComposedRouter(
|
|
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
|
)
|
|
cond_cont, cond_cat = _cond(16, pdg=5)
|
|
labels = torch.randint(0, 4, (16,))
|
|
loss = router.classify_loss(cond_cont, cond_cat, labels)
|
|
assert loss.shape == ()
|
|
assert loss.item() == 0.0
|
|
|
|
|
|
def test_composed_router_rejects_empty_router_list():
|
|
try:
|
|
ComposedRouter([])
|
|
except ValueError:
|
|
return
|
|
raise AssertionError("expected ValueError for empty router list")
|
|
|
|
|
|
def test_composed_router_not_in_registry():
|
|
assert "composed" not in ROUTER_REGISTRY
|
|
|
|
|
|
# ── _parse_composed_axes (axis{i}_{field} flat-key config convention) ───────
|
|
|
|
|
|
def test_parse_composed_axes_groups_indexed_keys():
|
|
from giant.model.network import _parse_composed_axes
|
|
|
|
router_cfg = {
|
|
"enabled": True,
|
|
"type": "composed",
|
|
"axis0_type": "energy",
|
|
"axis0_n_experts": 4,
|
|
"axis0_temperature": 0.3,
|
|
"axis1_type": "pdg",
|
|
"axis1_n_experts": 3,
|
|
"axis1_emb_dim": 6,
|
|
}
|
|
axes = _parse_composed_axes(router_cfg)
|
|
assert axes == [
|
|
{"type": "energy", "n_experts": 4, "temperature": 0.3},
|
|
{"type": "pdg", "n_experts": 3, "emb_dim": 6},
|
|
]
|
|
|
|
|
|
def test_parse_composed_axes_ignores_unrelated_keys():
|
|
from giant.model.network import _parse_composed_axes
|
|
|
|
router_cfg = {
|
|
"enabled": True,
|
|
"type": "composed",
|
|
"lambda_balance": 0.0,
|
|
"axis0_type": "energy",
|
|
"axis0_n_experts": 4,
|
|
}
|
|
axes = _parse_composed_axes(router_cfg)
|
|
assert axes == [{"type": "energy", "n_experts": 4}]
|
|
|
|
|
|
def test_parse_composed_axes_raises_on_index_gap():
|
|
from giant.model.network import _parse_composed_axes
|
|
|
|
router_cfg = {
|
|
"type": "composed",
|
|
"axis0_type": "energy",
|
|
"axis0_n_experts": 4,
|
|
# axis1 missing entirely
|
|
"axis2_type": "pdg",
|
|
"axis2_n_experts": 3,
|
|
}
|
|
try:
|
|
_parse_composed_axes(router_cfg)
|
|
except ValueError:
|
|
return
|
|
raise AssertionError("expected ValueError for a gap in axis indices")
|
|
|
|
|
|
def test_build_composed_router_resolves_per_axis_specs():
|
|
router = build_composed_router(
|
|
[
|
|
{"type": "energy", "n_experts": 4, "temperature": 0.3},
|
|
{"type": "pdg", "n_experts": 3, "emb_dim": 6},
|
|
],
|
|
pdg_vocab=5,
|
|
mat_vocab=2,
|
|
)
|
|
assert isinstance(router, ComposedRouter)
|
|
assert router.n_experts == 12
|
|
energy_router, pdg_router = router.routers
|
|
assert isinstance(energy_router, EnergyRouter)
|
|
assert energy_router.temperature == 0.3
|
|
assert isinstance(pdg_router, PdgRouter)
|
|
assert pdg_router.pdg_emb.num_embeddings == 5
|
|
assert pdg_router.pdg_emb.embedding_dim == 6
|
|
|
|
|
|
def test_build_models_routed_with_composed_router():
|
|
model_config = dict(
|
|
pdg_vocab=5,
|
|
mat_vocab=2,
|
|
emb_dim=16,
|
|
dropout=0.1,
|
|
k_max=K_MAX,
|
|
expert_hidden_dim=16,
|
|
expert_n_blocks=2,
|
|
router={
|
|
"enabled": True,
|
|
"type": "composed",
|
|
"axis0_type": "energy",
|
|
"axis0_n_experts": 4,
|
|
"axis1_type": "pdg",
|
|
"axis1_n_experts": 3,
|
|
},
|
|
)
|
|
stage1, sec_decoder = build_models(model_config)
|
|
assert isinstance(stage1, RoutedDenoisingMLP)
|
|
assert isinstance(stage1.router, ComposedRouter)
|
|
assert len(stage1.experts) == 12
|
|
assert len(sec_decoder.experts) == 12
|
|
# stage1 and sec_decoder must not share router weights (same convention
|
|
# as the single-axis routers built by build_models).
|
|
assert stage1.router is not sec_decoder.router
|
|
|
|
|
|
def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
|
|
from giant.sample import sample_flow, sample_secondaries
|
|
|
|
model_config = dict(
|
|
pdg_vocab=3,
|
|
mat_vocab=2,
|
|
emb_dim=16,
|
|
dropout=0.1,
|
|
k_max=K_MAX,
|
|
expert_hidden_dim=8,
|
|
expert_n_blocks=1,
|
|
router={
|
|
"enabled": True,
|
|
"type": "composed",
|
|
"axis0_type": "energy",
|
|
"axis0_n_experts": 2,
|
|
"axis1_type": "pdg",
|
|
"axis1_n_experts": 2,
|
|
},
|
|
)
|
|
stage1, sec_decoder = build_models(model_config)
|
|
B = 5
|
|
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
|
|
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
|
|
assert stage1_norm.shape == (B, X_DIM)
|
|
assert n_sec_pred.shape == (B,)
|
|
|
|
sec_cont, sec_type_emb, sec_valid = sample_secondaries(
|
|
sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
|
|
)
|
|
assert sec_cont.shape == (B, K_MAX, 4)
|
|
assert sec_valid.shape == (B, K_MAX)
|
|
|
|
|
|
# ── RoutedDenoisingMLP ───────────────────────────────────────────────────────
|
|
|
|
|
|
def test_routed_denoising_mlp_output_shape_train_and_eval():
|
|
B = 8
|
|
model = _routed_stage1()
|
|
x_t = torch.randn(B, X_DIM)
|
|
t = torch.rand(B)
|
|
cond_cont, cond_cat = _cond(B)
|
|
|
|
model.train()
|
|
out_train = model(x_t, t, cond_cont, cond_cat)
|
|
assert out_train.shape == (B, X_DIM)
|
|
|
|
model.eval()
|
|
with torch.no_grad():
|
|
out_eval = model(x_t, t, cond_cont, cond_cat)
|
|
assert out_eval.shape == (B, X_DIM)
|
|
|
|
|
|
def test_routed_denoising_mlp_gradients_flow_in_train_mode():
|
|
"""Soft mixture in train mode should touch every expert's parameters."""
|
|
B = 8
|
|
model = _routed_stage1(n_experts=3)
|
|
x_t = torch.randn(B, X_DIM)
|
|
t = torch.rand(B)
|
|
cond_cont, cond_cat = _cond(B)
|
|
model.train()
|
|
flow_loss = model(x_t, t, cond_cont, cond_cat).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_routed_denoising_mlp_eval_dispatch_matches_manual_grouping():
|
|
"""Eval-mode grouped top-1 dispatch must equal running each row through
|
|
its assigned expert individually (batch order shouldn't matter)."""
|
|
B = 12
|
|
model = _routed_stage1(n_experts=4)
|
|
model.eval()
|
|
x_t = torch.randn(B, X_DIM)
|
|
t = torch.rand(B)
|
|
cond_cont, cond_cat = _cond(B)
|
|
|
|
with torch.no_grad():
|
|
batched = model(x_t, t, cond_cont, cond_cat)
|
|
|
|
t_emb = model.time_emb(t)
|
|
c_emb = model.cond_enc(cond_cont, cond_cat)
|
|
cond = torch.cat([t_emb, c_emb], dim=-1)
|
|
idx = model.router.top1(cond_cont, cond_cat)
|
|
manual = torch.zeros_like(x_t)
|
|
for i in range(B):
|
|
manual[i] = model.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0]
|
|
|
|
torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4)
|
|
|
|
|
|
def test_routed_denoising_mlp_predict_n_sec_shape():
|
|
B = 6
|
|
model = _routed_stage1()
|
|
cond_cont, cond_cat = _cond(B)
|
|
logits = model.predict_n_sec(cond_cont, cond_cat)
|
|
assert logits.shape == (B, K_MAX + 1)
|
|
|
|
|
|
def test_routed_denoising_mlp_has_no_pdg_embedding_weight_method():
|
|
model = _routed_stage1(pdg=5, mat=2)
|
|
assert not hasattr(model, "pdg_embedding_weight")
|
|
|
|
|
|
# ── RoutedSecondaryDecoder ───────────────────────────────────────────────────
|
|
|
|
|
|
def test_routed_secondary_decoder_output_shape_train_and_eval():
|
|
B = 8
|
|
decoder = _routed_sec_decoder()
|
|
x_t = torch.randn(B, SEC_DIM)
|
|
t = torch.rand(B)
|
|
cond_cont, cond_cat = _cond(B)
|
|
stage1_out = torch.randn(B, X_DIM)
|
|
|
|
decoder.train()
|
|
out_train = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
|
|
assert out_train.shape == (B, SEC_DIM)
|
|
|
|
decoder.eval()
|
|
with torch.no_grad():
|
|
out_eval = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
|
|
assert out_eval.shape == (B, SEC_DIM)
|
|
|
|
|
|
def test_routed_secondary_decoder_gradients_flow():
|
|
B = 4
|
|
decoder = _routed_sec_decoder(n_experts=3)
|
|
x_t = torch.randn(B, SEC_DIM)
|
|
t = torch.rand(B)
|
|
cond_cont, cond_cat = _cond(B)
|
|
stage1_out = torch.randn(B, X_DIM)
|
|
decoder.train()
|
|
decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward()
|
|
for name, p in decoder.named_parameters():
|
|
assert p.grad is not None, f"no grad for {name}"
|
|
|
|
|
|
# ── build_models dispatch ────────────────────────────────────────────────────
|
|
|
|
|
|
def test_build_models_monolith_when_router_absent():
|
|
model_config = dict(
|
|
pdg_vocab=4,
|
|
mat_vocab=2,
|
|
hidden_dim=32,
|
|
n_blocks=2,
|
|
emb_dim=16,
|
|
dropout=0.1,
|
|
k_max=K_MAX,
|
|
)
|
|
stage1, sec_decoder = build_models(model_config)
|
|
assert isinstance(stage1, DenoisingMLP)
|
|
assert isinstance(sec_decoder, SecondaryDecoder)
|
|
|
|
|
|
def test_build_models_monolith_when_router_disabled():
|
|
model_config = dict(
|
|
pdg_vocab=4,
|
|
mat_vocab=2,
|
|
hidden_dim=32,
|
|
n_blocks=2,
|
|
emb_dim=16,
|
|
dropout=0.1,
|
|
k_max=K_MAX,
|
|
router={"enabled": False, "type": "energy", "n_experts": 4},
|
|
)
|
|
stage1, sec_decoder = build_models(model_config)
|
|
assert isinstance(stage1, DenoisingMLP)
|
|
assert isinstance(sec_decoder, SecondaryDecoder)
|
|
|
|
|
|
def test_build_models_routed_when_enabled():
|
|
model_config = dict(
|
|
pdg_vocab=4,
|
|
mat_vocab=2,
|
|
emb_dim=16,
|
|
dropout=0.1,
|
|
k_max=K_MAX,
|
|
expert_hidden_dim=16,
|
|
expert_n_blocks=2,
|
|
router={
|
|
"enabled": True,
|
|
"type": "energy",
|
|
"n_experts": 4,
|
|
"temperature": 0.5,
|
|
"learn_centers": True,
|
|
"lambda_balance": 0.0,
|
|
},
|
|
)
|
|
stage1, sec_decoder = build_models(model_config)
|
|
assert isinstance(stage1, RoutedDenoisingMLP)
|
|
assert isinstance(sec_decoder, RoutedSecondaryDecoder)
|
|
assert len(stage1.experts) == 4
|
|
assert len(sec_decoder.experts) == 4
|
|
|
|
|
|
def test_build_models_routed_pair_is_drop_in_for_sample_flow():
|
|
"""Exercise the exact calling convention giant/sample.py uses."""
|
|
from giant.sample import sample_flow, sample_secondaries
|
|
|
|
model_config = dict(
|
|
pdg_vocab=3,
|
|
mat_vocab=2,
|
|
emb_dim=16,
|
|
dropout=0.1,
|
|
k_max=K_MAX,
|
|
expert_hidden_dim=8,
|
|
expert_n_blocks=1,
|
|
router={"enabled": True, "type": "energy", "n_experts": 2},
|
|
)
|
|
stage1, sec_decoder = build_models(model_config)
|
|
B = 5
|
|
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
|
|
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
|
|
assert stage1_norm.shape == (B, X_DIM)
|
|
assert n_sec_pred.shape == (B,)
|
|
|
|
sec_cont, sec_type_emb, sec_valid = sample_secondaries(
|
|
sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
|
|
)
|
|
assert sec_cont.shape == (B, K_MAX, 4)
|
|
assert sec_valid.shape == (B, K_MAX)
|