lars/giant
1
0
Fork 0
Code Issues 3 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Add ComposedRouter for multi-axis MoE gating

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
lars
2026-07-15 13:53:21 +02:00
co-authored by Claude Opus 4.8
parent 550dc679c7
commit f3fec8bcb3
4 changed files with 405 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/model/network.py
+128 -19
View File
@@ -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