From f3fec8bcb38f7c12717abeffc5ecc3a51927b71a Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 15 Jul 2026 13:53:21 +0200 Subject: [PATCH] Add ComposedRouter for multi-axis MoE gating Route on several independent axes at once (e.g. energy x pdg), each with its own expert count and hyperparameters. The joint gate is the outer product of per-axis softmax gates, so it stays a partition of unity and top1/balance_loss factor per-axis. Config uses flat axis{i}_{field} keys in model.router (TOML/CLI friendly), also settable via repeatable --router-axis flags. Co-Authored-By: Claude Opus 4.8 --- giant/cli.py | 51 +++++++++- giant/config.py | 7 ++ giant/model/network.py | 147 +++++++++++++++++++++++---- tests/test_router.py | 220 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 405 insertions(+), 20 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index f36e074..b7a3a5f 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -63,6 +63,43 @@ def _batch_size_estimate_dims(model_cfg: dict) -> tuple[int, int]: return model_cfg["hidden_dim"], model_cfg["n_blocks"] +def _coerce_scalar(value: str) -> object: + """Best-effort str -> bool/int/float, else leave as str. + + CLI flag values always arrive as strings; router kwargs like + `n_experts` (int) or `temperature` (float) need to come out typed the + same way a TOML file's native types would, since they're merged into + the same `model.router` dict as file-sourced config. + """ + if value.lower() in ("true", "false"): + return value.lower() == "true" + try: + return int(value) + except ValueError: + pass + try: + return float(value) + except ValueError: + pass + return value + + +def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]: + """Parse repeated `--router-axis "type:key=val,key=val"` flags into + `axis{i}_{field}` flat keys (see `_parse_composed_axes` in + giant.model.network), indexed by flag order — the Nth `--router-axis` + becomes axis N. + """ + out: dict[str, object] = {} + for i, spec in enumerate(specs): + axis_type, _, rest = spec.partition(":") + out[f"axis{i}_type"] = axis_type + for pair in filter(None, rest.split(",")): + key, _, val = pair.partition("=") + out[f"axis{i}_{key}"] = _coerce_scalar(val) + return out + + _CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions") @@ -175,6 +212,16 @@ def train( n_experts: Annotated[ Optional[int], typer.Option("--n-experts", help="Number of routed experts") ] = None, + router_axis: Annotated[ + Optional[list[str]], + typer.Option( + "--router-axis", + help="Composed-router axis spec 'type:key=val,key=val' (repeatable; " + "Nth flag = axis N). Use with --router-type composed instead of " + "--n-experts, e.g. --router-axis 'energy:n_experts=4' " + "--router-axis 'pdg:n_experts=3,emb_dim=8'", + ), + ] = None, val_fraction: Annotated[ Optional[float], typer.Option("--val-fraction", "-f") ] = None, @@ -264,7 +311,7 @@ def train( }.items() if v is not None } - cli_router = { + cli_router: dict[str, object] = { k: v for k, v in { "enabled": router, @@ -273,6 +320,8 @@ def train( }.items() if v is not None } + if router_axis: + cli_router.update(_parse_router_axis_flags(router_axis)) if cli_router: cli_model["router"] = cli_router cfg = gconfig.merge_cli_overrides( diff --git a/giant/config.py b/giant/config.py index 7e4f8e7..8ce67f5 100644 --- a/giant/config.py +++ b/giant/config.py @@ -43,6 +43,13 @@ DEFAULT_CONFIG: dict = { # (0.0 still trains a working router — the gate gets gradient # through the downstream flow loss like EnergyRouter's centers — # but only lambda_proc > 0 grounds it in the true `process` label) + # type = "composed" routes on multiple axes at once (e.g. energy x + # pdg), each with its own expert count/hyperparameters. Axes are + # NOT in these defaults (there's no meaningful default axis list) + # — set them as flat axis{i}_{field} keys instead of "n_experts", + # e.g. axis0_type = "energy", axis0_n_experts = 4, axis1_type = + # "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. See + # giant.model.network._parse_composed_axes / `--router-axis`. }, }, } diff --git a/giant/model/network.py b/giant/model/network.py index 7a6e8ca..57d25df 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -1,5 +1,6 @@ import inspect import math +import re import torch import torch.nn as nn @@ -444,6 +445,79 @@ class ProcessRouter(Router): return F.cross_entropy(self.logits(cond_cont, cond_cat), labels) +class ComposedRouter(Router): + """Joint router over independent axes (e.g. energy x pdg), outer-product gated. + + Wraps N already-built sub-routers, each free to have its own + `n_experts` and hyperparameters (an `EnergyRouter(n_experts=4, ...)` + composed with a `PdgRouter(n_experts=3, ...)` needs no axis to match + the other's expert count). The joint gate is the outer product of the + per-axis softmax gates, flattened to `(B, prod(n_experts_i))` — still a + partition of unity, since each factor is one. Because the axes are + routed independently, the joint argmax factors into the per-axis + argmaxes, so `top1` (inherited from `Router`) costs no more than + routing each axis alone despite the multiplicative expert count; the + same is true of `balance_loss` (inherited, computed on the flattened + joint gate — now one importance term per *joint* expert cell). + + Not registered in `ROUTER_REGISTRY` / buildable via `build_router`, + since those assume one `n_experts` int shared by a single router type; + use `build_composed_router` instead, which resolves a list of per-axis + specs (each independently typed and sized) through `build_router`. + """ + + def __init__(self, routers: list[Router]) -> None: + if not routers: + raise ValueError("ComposedRouter needs at least one sub-router") + n_experts = 1 + for r in routers: + n_experts *= r.n_experts + super().__init__(n_experts) + self.routers = nn.ModuleList(routers) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + joint = self.routers[0].gate(cond_cont, cond_cat) # (B, n_0) + for router in self.routers[1:]: + g = router.gate(cond_cont, cond_cat) # (B, n_i) + joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten( + 1 + ) # (B, prod so far) + return joint + + def classify_loss( + self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor + ) -> torch.Tensor: + """Sum of each sub-router's own classify_loss (0 for unsupervised axes).""" + total = torch.zeros((), device=cond_cont.device) + for router in self.routers: + total = total + router.classify_loss(cond_cont, cond_cat, labels) + return total + + +def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter: + """Build a `ComposedRouter` from a list of per-axis router specs. + + Each spec is a `{"type": ..., "n_experts": ..., ...per-axis kwargs}` + dict resolved through `build_router` exactly like a single-axis router + config, so axes can differ in both expert count and hyperparameters + (e.g. an energy axis's `temperature` vs a pdg axis's `emb_dim`). + `shared_kwargs` (`pdg_vocab`, `mat_vocab`, ...) are merged under each + spec, with the spec's own keys taking precedence. + """ + routers = [ + build_router( + spec["type"], + spec["n_experts"], + **{ + **shared_kwargs, + **{k: v for k, v in spec.items() if k not in ("type", "n_experts")}, + }, + ) + for spec in specs + ] + return ComposedRouter(routers) + + class ExpertTrunk(nn.Module): """One small expert: `input_proj -> ResBlock stack -> out_proj`. @@ -669,6 +743,54 @@ _SEC_DECODER_MODEL_KEYS = { } +_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$") + + +def _parse_composed_axes(router_cfg: dict) -> list[dict]: + """Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts. + + Flat keys (rather than a nested list-of-dicts) keep composed-router + config expressible in the same one-level-of-nesting TOML/CLI shape as + every other router option (`model.router` stays a flat table of + scalars) — e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`, + `axis1_type = "pdg"`, `axis1_n_experts = 3`, `axis1_emb_dim = 8`. + Axis indices must be contiguous from 0; order follows the index, not + dict insertion order (TOML/CLI merging doesn't preserve it reliably). + """ + axes: dict[int, dict] = {} + for key, value in router_cfg.items(): + m = _AXIS_KEY_RE.match(key) + if m is None: + continue + idx, field = int(m.group(1)), m.group(2) + axes.setdefault(idx, {})[field] = value + missing = set(range(len(axes))) - axes.keys() + if missing: + raise ValueError(f"composed router config has gaps at axis indices {missing}") + return [axes[i] for i in range(len(axes))] + + +def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> Router: + """Resolve one `model.router` config into a `Router`, single-axis or composed. + + `router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys + (see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair. + """ + shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab) + if router_cfg["type"] == "composed": + return build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab) + router_kwargs = { + k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts") + } + # Not every router needs these (EnergyRouter doesn't declare them, so + # build_router's kwarg filtering drops them silently) but ProcessRouter + # needs its own pdg/material embeddings sized to match the checkpoint's + # vocab, same as the trunk's ConditionEncoder. + router_kwargs.setdefault("pdg_vocab", pdg_vocab) + router_kwargs.setdefault("mat_vocab", mat_vocab) + return build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs) + + def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: """Construct (stage1, sec_decoder) from a persisted/CLI model_config dict. @@ -679,36 +801,23 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: """ router_cfg = model_config.get("router") if router_cfg and router_cfg.get("enabled"): - router_kwargs = { - k: v - for k, v in router_cfg.items() - if k not in ("enabled", "type", "n_experts") - } - # Not every router needs these (EnergyRouter doesn't declare them, so - # build_router's kwarg filtering drops them silently) but - # ProcessRouter needs its own pdg/material embeddings sized to match - # the checkpoint's vocab, same as the trunk's ConditionEncoder. - router_kwargs.setdefault("pdg_vocab", model_config["pdg_vocab"]) - router_kwargs.setdefault("mat_vocab", model_config["mat_vocab"]) + pdg_vocab = model_config["pdg_vocab"] + mat_vocab = model_config["mat_vocab"] shared = dict( - pdg_vocab=model_config["pdg_vocab"], - mat_vocab=model_config["mat_vocab"], + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, expert_hidden_dim=model_config.get("expert_hidden_dim", 128), expert_n_blocks=model_config.get("expert_n_blocks", 3), emb_dim=model_config.get("emb_dim", EMB_DIM), dropout=model_config.get("dropout", 0.1), ) stage1 = RoutedDenoisingMLP( - router=build_router( - router_cfg["type"], router_cfg["n_experts"], **router_kwargs - ), + router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab), k_max=model_config.get("k_max", K_MAX), **shared, ) sec_decoder = RoutedSecondaryDecoder( - router=build_router( - router_cfg["type"], router_cfg["n_experts"], **router_kwargs - ), + router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab), **shared, ) return stage1, sec_decoder diff --git a/tests/test_router.py b/tests/test_router.py index eef1c3f..31f6b54 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -4,6 +4,7 @@ import torch from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( + ComposedRouter, DenoisingMLP, EnergyRouter, PdgRouter, @@ -12,6 +13,7 @@ from giant.model.network import ( RoutedDenoisingMLP, RoutedSecondaryDecoder, SecondaryDecoder, + build_composed_router, build_models, build_router, ) @@ -289,6 +291,224 @@ def test_build_models_routed_with_process_router(): 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 ───────────────────────────────────────────────────────