lars
bac541240f
Both stages can now route through a pluggable Router (EnergyRouter as the first implementation, a soft turn-on gate over pre-step log-energy) into several small ExpertTrunks instead of one monolithic trunk. Trains as a differentiable soft mixture and dispatches to a single expert per row at eval time, which is the source of the per-call speedup this prototype is after (issue #5's ~10x native-Geant4 budget). Disabled by default, so existing configs/checkpoints are unaffected; build_models() centralizes routed-vs-monolith construction across train/predict/rollout. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
297 lines
9.2 KiB
Python
297 lines
9.2 KiB
Python
"""Tests for the mixture-of-experts routing prototype (giant/model/network.py)."""
|
|
|
|
import torch
|
|
|
|
from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
|
|
from giant.model.network import (
|
|
DenoisingMLP,
|
|
EnergyRouter,
|
|
ROUTER_REGISTRY,
|
|
RoutedDenoisingMLP,
|
|
RoutedSecondaryDecoder,
|
|
SecondaryDecoder,
|
|
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_build_router_unknown_type_raises():
|
|
try:
|
|
build_router("nonexistent", 4)
|
|
except ValueError:
|
|
return
|
|
raise AssertionError("expected ValueError for unknown router type")
|
|
|
|
|
|
# ── 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_pdg_embedding_weight_shape():
|
|
model = _routed_stage1(pdg=5, mat=2)
|
|
from giant.constants import EMB_DIM
|
|
|
|
assert model.pdg_embedding_weight().shape == (5, EMB_DIM)
|
|
|
|
|
|
# ── 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)
|