From bac541240f580b3cf92c5f49b3f924d7d0596b8d Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 8 Jul 2026 15:17:19 +0200 Subject: [PATCH] Add mixture-of-experts routing prototype for Stage 1 and Stage 2 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 --- giant/cli.py | 87 ++++++---- giant/config.py | 55 ++++++- giant/model/network.py | 366 ++++++++++++++++++++++++++++++++++++++++- giant/pipeline.py | 47 ++---- giant/train.py | 38 ++++- tests/test_router.py | 296 +++++++++++++++++++++++++++++++++ 6 files changed, 815 insertions(+), 74 deletions(-) create mode 100644 tests/test_router.py diff --git a/giant/cli.py b/giant/cli.py index 4b61b10..42f8ad0 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -39,31 +39,30 @@ from giant.data.transforms import ( Normalizer, ) from giant.geometry import GeometryOracle -from giant.model.network import DenoisingMLP, SecondaryDecoder +from giant.model.network import build_models from giant.pipeline import run_train_job from giant.rollout import rollout as run_rollout from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx -_STAGE1_MODEL_KEYS = { - "pdg_vocab", - "mat_vocab", - "hidden_dim", - "n_blocks", - "emb_dim", - "dropout", - "k_max", -} -_SEC_DECODER_MODEL_KEYS = { - "pdg_vocab", - "mat_vocab", - "hidden_dim", - "n_blocks", - "emb_dim", - "dropout", -} - app = typer.Typer(no_args_is_help=True) + +def _batch_size_estimate_dims(model_cfg: dict) -> tuple[int, int]: + """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory. + + Routed models spend their FLOPs in the (smaller) expert trunks, not the + monolith's hidden_dim/n_blocks, so estimate_batch_size needs the expert + dims instead when routing is enabled. + """ + router_cfg = model_cfg.get("router") + if router_cfg and router_cfg.get("enabled"): + return ( + model_cfg.get("expert_hidden_dim", 128), + model_cfg.get("expert_n_blocks", 3), + ) + return model_cfg["hidden_dim"], model_cfg["n_blocks"] + + _CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions") @@ -159,6 +158,23 @@ def train( "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)" ), ] = None, + router: Annotated[ + Optional[bool], + typer.Option( + "--router/--no-router", + help="Route both stages through a mixture of small experts " + "instead of one monolithic trunk (see model.router in config.toml)", + ), + ] = None, + router_type: Annotated[ + Optional[str], + typer.Option( + "--router-type", help="Router implementation name (see ROUTER_REGISTRY)" + ), + ] = None, + n_experts: Annotated[ + Optional[int], typer.Option("--n-experts", help="Number of routed experts") + ] = None, val_fraction: Annotated[ Optional[float], typer.Option("--val-fraction", "-f") ] = None, @@ -238,7 +254,7 @@ def train( }.items() if v is not None } - cli_model = { + cli_model: dict[str, object] = { k: v for k, v in { "hidden_dim": hidden_dim, @@ -248,6 +264,17 @@ def train( }.items() if v is not None } + cli_router = { + k: v + for k, v in { + "enabled": router, + "type": router_type, + "n_experts": n_experts, + }.items() + if v is not None + } + if cli_router: + cli_model["router"] = cli_router cfg = gconfig.merge_cli_overrides( gconfig.DEFAULT_CONFIG, config, cli_train, cli_model ) @@ -256,9 +283,10 @@ def train( _device = torch.device(device) if device else gconfig.auto_device() if batch_size_auto: + est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m) try: t["batch_size"] = gconfig.estimate_batch_size( - m["hidden_dim"], m["n_blocks"], _device + est_hidden_dim, est_n_blocks, _device ) except ValueError as exc: typer.echo(f"error: {exc}", err=True) @@ -386,10 +414,11 @@ def predict( model_cfg = ckpt["model_config"] if batch_size_auto: + est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(model_cfg) try: batch_size_value = gconfig.estimate_batch_size( - model_cfg["hidden_dim"], - model_cfg["n_blocks"], + est_hidden_dim, + est_n_blocks, _device, training=False, ) @@ -408,15 +437,10 @@ def predict( cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) - model = DenoisingMLP( - **{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS} - ) + model, sec_decoder = build_models(model_cfg) model.load_state_dict(ckpt["model"]) model.to(_device).eval() - sec_decoder = SecondaryDecoder( - **{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS} - ) sec_decoder.load_state_dict(ckpt["sec_decoder"]) sec_decoder.to(_device).eval() @@ -759,12 +783,9 @@ def rollout( cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) - model = DenoisingMLP(**{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS}) + model, sec_decoder = build_models(model_cfg) model.load_state_dict(ckpt["model"]) model.to(_device).eval() - sec_decoder = SecondaryDecoder( - **{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS} - ) sec_decoder.load_state_dict(ckpt["sec_decoder"]) sec_decoder.to(_device).eval() typer.echo(f"loaded checkpoint: {checkpoint}") diff --git a/giant/config.py b/giant/config.py index 7001151..b5008a6 100644 --- a/giant/config.py +++ b/giant/config.py @@ -28,6 +28,16 @@ DEFAULT_CONFIG: dict = { "n_blocks": 6, "emb_dim": 16, "dropout": 0.1, + "router": { + "enabled": False, + "type": "energy", # selects the Router impl from ROUTER_REGISTRY + "n_experts": 4, + "expert_hidden_dim": 128, + "expert_n_blocks": 3, + "temperature": 0.5, # energy-router kwarg + "learn_centers": True, # energy-router kwarg + "lambda_balance": 0.0, # optional load-balance aux loss weight + }, }, } @@ -164,15 +174,29 @@ def merge_cli_overrides( train_overrides: dict, model_overrides: dict, ) -> dict: - """Resolve config as defaults -> TOML file -> explicit CLI flags.""" + """Resolve config as defaults -> TOML file -> explicit CLI flags. + + `model.router` is deep-merged one level (rather than replaced wholesale) + at each stage, so a TOML file or CLI flag only overriding e.g. + `router.enabled` doesn't drop the rest of the router defaults. + """ cfg = {"train": dict(defaults["train"]), "model": dict(defaults["model"])} + cfg["model"]["router"] = dict(defaults["model"]["router"]) if config_path is not None: file_cfg = load_toml(config_path) - for section in ("train", "model"): - cfg[section].update(file_cfg.get(section, {})) + cfg["train"].update(file_cfg.get("train", {})) + file_model = dict(file_cfg.get("model", {})) + file_router = file_model.pop("router", None) + cfg["model"].update(file_model) + if file_router: + cfg["model"]["router"].update(file_router) warn_if_git_hash_mismatch(file_cfg, config_path) + model_overrides = dict(model_overrides) + router_overrides = model_overrides.pop("router", None) cfg["train"].update(train_overrides) cfg["model"].update(model_overrides) + if router_overrides: + cfg["model"]["router"].update(router_overrides) return cfg @@ -184,17 +208,38 @@ def seed_everything(seed: int) -> None: torch.cuda.manual_seed_all(seed) +def _toml_value(v) -> str: + if isinstance(v, bool): + return "true" if v else "false" + if isinstance(v, str): + return repr(v) + return str(v) + + def save_config(cfg: dict, out_dir: Path, meta: dict) -> None: lines = [] + # One-level-nested dict values (e.g. model.router) are rendered as their + # own [section.subsection] table after the parent section, since TOML + # doesn't accept a bare dict as a `key = value` scalar line. + nested_sections: list[tuple[str, dict]] = [] for section, values in cfg.items(): lines.append(f"[{section}]") for k, v in values.items(): - lines.append(f"{k:<14} = {repr(v) if isinstance(v, str) else v}") + if isinstance(v, dict): + nested_sections.append((f"{section}.{k}", v)) + continue + lines.append(f"{k:<14} = {_toml_value(v)}") + lines.append("") + + for name, values in nested_sections: + lines.append(f"[{name}]") + for k, v in values.items(): + lines.append(f"{k:<14} = {_toml_value(v)}") lines.append("") lines.append("[meta]") for k, v in meta.items(): - lines.append(f"{k:<14} = {repr(v) if isinstance(v, str) else v}") + lines.append(f"{k:<14} = {_toml_value(v)}") (out_dir / "config.toml").write_text("\n".join(lines)) diff --git a/giant/model/network.py b/giant/model/network.py index b64b48e..91c06a1 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -1,9 +1,10 @@ +import inspect import math import torch import torch.nn as nn -from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM +from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM class SinusoidalEmbedding(nn.Module): @@ -236,3 +237,366 @@ class SecondaryDecoder(nn.Module): for block in self.blocks: x = block(x, cond) return self.out_proj(x) + + +class Router(nn.Module): + """Contract for a pluggable mixture-of-experts routing axis. + + Subclasses implement `gate` (soft partition-of-unity weights over + experts, used in train mode for a fully differentiable mixture); + `top1` and `balance_loss` have working defaults so a new routing axis + is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`. + """ + + def __init__(self, n_experts: int) -> None: + super().__init__() + self.n_experts = n_experts + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """(B, n_experts) soft weights, rows summing to 1.""" + raise NotImplementedError + + def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """(B,) hard expert index, used for eval-time grouped dispatch.""" + return self.gate(cond_cont, cond_cat).argmax(dim=-1) + + def balance_loss( + self, cond_cont: torch.Tensor, cond_cat: torch.Tensor + ) -> torch.Tensor: + """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017).""" + importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,) + return (importance.std() / (importance.mean() + 1e-8)) ** 2 + + +ROUTER_REGISTRY: dict[str, type[Router]] = {} + + +def register_router(name: str): + def decorator(cls: type[Router]) -> type[Router]: + ROUTER_REGISTRY[name] = cls + return cls + + return decorator + + +def build_router(name: str, n_experts: int, **kwargs) -> Router: + """Factory: look up a `Router` subclass by name from the registry. + + Every registered router type is fed the same `model.router` config + dict; kwargs not declared by that type's constructor are silently + dropped, so per-type hyperparameters (e.g. EnergyRouter's + `temperature`) can coexist in one config without special-casing. + """ + if name not in ROUTER_REGISTRY: + raise ValueError( + f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}" + ) + cls = ROUTER_REGISTRY[name] + accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"} + filtered = {k: v for k, v in kwargs.items() if k in accepted} + return cls(n_experts=n_experts, **filtered) + + +@register_router("energy") +class EnergyRouter(Router): + """Soft turn-on gate over normalized pre-step log-energy. + + Reads `cond_cont[:, energy_idx]` (ignores cond_cat). Learnable (or + fixed) 1-D centers, initialized spread across [-2, 2] — roughly the + z-normalized energy range. `gate(e) = softmax_i(-(e - c_i)^2 / tau)`, + differentiable in e; as tau -> 0 this hardens to nearest-center + (Voronoi) selection, which is exactly what `top1` uses at eval. + """ + + def __init__( + self, + n_experts: int = 4, + temperature: float = 0.5, + learn_centers: bool = True, + energy_idx: int = 3, + ) -> None: + super().__init__(n_experts) + self.temperature = temperature + self.energy_idx = energy_idx + centers = torch.linspace(-2.0, 2.0, n_experts) + if learn_centers: + self.centers = nn.Parameter(centers) + else: + self.register_buffer("centers", centers) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + e = cond_cont[:, self.energy_idx].unsqueeze(-1) # (B, 1) + d2 = (e - self.centers.unsqueeze(0)) ** 2 # (B, n_experts) + return torch.softmax(-d2 / self.temperature, dim=-1) + + +class ExpertTrunk(nn.Module): + """One small expert: `input_proj -> ResBlock stack -> out_proj`. + + Same shape as the monolithic DenoisingMLP/SecondaryDecoder trunk, but + intended to be narrower/shallower (per-call cost is the whole point). + """ + + def __init__( + self, + in_dim: int, + hidden_dim: int, + n_blocks: int, + merged_cond_dim: int, + dropout: float = 0.1, + ) -> None: + super().__init__() + self.input_proj = nn.Linear(in_dim, hidden_dim) + self.blocks = nn.ModuleList( + [ + ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) + for _ in range(n_blocks) + ] + ) + self.out_proj = nn.Linear(hidden_dim, in_dim) + + def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + x = self.input_proj(x) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + +def _route_forward( + experts: nn.ModuleList, + router: Router, + x: torch.Tensor, + cond: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + training: bool, +) -> torch.Tensor: + """Shared dispatch for both Routed* trunks. + + Train mode: full soft mixture `sum_i gate_i * expert_i(x)` — fully + differentiable, N-expert compute. Eval mode: grouped top-1 dispatch — + each row runs exactly one (small) expert, which is the actual source + of the per-call speedup this architecture is for. + """ + if training: + weights = router.gate(cond_cont, cond_cat) # (B, n_experts) + out = torch.zeros_like(x) + for i, expert in enumerate(experts): + out = out + weights[:, i : i + 1] * expert(x, cond) + return out + + idx = router.top1(cond_cont, cond_cat) # (B,) + out = torch.zeros_like(x) + for i, expert in enumerate(experts): + mask = idx == i + if mask.any(): + out[mask] = expert(x[mask], cond[mask]) + return out + + +class RoutedDenoisingMLP(nn.Module): + """Routed drop-in for `DenoisingMLP`. + + Shares the time embedding, `ConditionEncoder`, and `n_sec_head` (all + tiny) across experts and routes only the trunk (where the FLOPs are). + Same `forward`/`predict_n_sec`/`pdg_embedding_weight` signatures as + `DenoisingMLP`, so sample.py/rollout.py/validate.py need no changes. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + router: Router, + expert_hidden_dim: int = 128, + expert_n_blocks: int = 3, + emb_dim: int = EMB_DIM, + time_dim: int = 64, + cond_out_dim: int = 128, + x_dim: int = X_DIM, + dropout: float = 0.1, + k_max: int = K_MAX, + ) -> None: + super().__init__() + self.router = router + self.time_emb = SinusoidalEmbedding(time_dim) + self.cond_enc = ConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + out_dim=cond_out_dim, + ) + merged_cond_dim = time_dim + cond_out_dim + self.experts = nn.ModuleList( + [ + ExpertTrunk( + x_dim, + expert_hidden_dim, + expert_n_blocks, + merged_cond_dim, + dropout=dropout, + ) + for _ in range(router.n_experts) + ] + ) + self.n_sec_head = nn.Sequential( + nn.Linear(cond_out_dim, cond_out_dim), + nn.SiLU(), + nn.Linear(cond_out_dim, k_max + 1), + ) + + def forward( + self, + x_t: torch.Tensor, + t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + t_emb = self.time_emb(t) + c_emb = self.cond_enc(cond_cont, cond_cat) + cond = torch.cat([t_emb, c_emb], dim=-1) + return _route_forward( + self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training + ) + + def predict_n_sec( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + """Return n_sec logits (B, K_MAX+1) from conditioning alone.""" + c_emb = self.cond_enc(cond_cont, cond_cat) + return self.n_sec_head(c_emb) + + def pdg_embedding_weight(self) -> torch.Tensor: + """Return the PDG embedding table weights for secondary type targets.""" + return self.cond_enc.pdg_emb.weight + + +class RoutedSecondaryDecoder(nn.Module): + """Routed drop-in for `SecondaryDecoder`. + + Shares the time embedding and `SecondaryConditionEncoder` across + experts and routes only the trunk. Same `forward` signature as + `SecondaryDecoder`. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + router: Router, + expert_hidden_dim: int = 128, + expert_n_blocks: int = 3, + emb_dim: int = EMB_DIM, + time_dim: int = 64, + cond_out_dim: int = 128, + stage1_proj_dim: int = 64, + sec_dim: int = SEC_DIM, + dropout: float = 0.1, + ) -> None: + super().__init__() + self.router = router + self.time_emb = SinusoidalEmbedding(time_dim) + self.cond_enc = SecondaryConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + cond_out_dim=cond_out_dim, + stage1_proj_dim=stage1_proj_dim, + out_dim=cond_out_dim, + ) + merged_cond_dim = time_dim + cond_out_dim + self.experts = nn.ModuleList( + [ + ExpertTrunk( + sec_dim, + expert_hidden_dim, + expert_n_blocks, + merged_cond_dim, + dropout=dropout, + ) + for _ in range(router.n_experts) + ] + ) + + def forward( + self, + x_t: torch.Tensor, + t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + t_emb = self.time_emb(t) + c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out) + cond = torch.cat([t_emb, c_emb], dim=-1) + return _route_forward( + self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training + ) + + +_STAGE1_MODEL_KEYS = { + "pdg_vocab", + "mat_vocab", + "hidden_dim", + "n_blocks", + "emb_dim", + "dropout", + "k_max", +} +_SEC_DECODER_MODEL_KEYS = { + "pdg_vocab", + "mat_vocab", + "hidden_dim", + "n_blocks", + "emb_dim", + "dropout", +} + + +def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: + """Construct (stage1, sec_decoder) from a persisted/CLI model_config dict. + + Dispatches to the routed pair when `model_config["router"]["enabled"]` + is truthy; a missing/absent "router" key (pre-routing checkpoints) + falls back to the monolithic pair unchanged, so this is a drop-in + replacement for the ad-hoc constructions it replaces. + """ + 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") + } + shared = dict( + pdg_vocab=model_config["pdg_vocab"], + mat_vocab=model_config["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 + ), + 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 + ), + **shared, + ) + return stage1, sec_decoder + + stage1 = DenoisingMLP( + **{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS} + ) + sec_decoder = SecondaryDecoder( + **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS} + ) + return stage1, sec_decoder diff --git a/giant/pipeline.py b/giant/pipeline.py index 593455c..16d655f 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -14,7 +14,7 @@ from giant.data.loader import ( ) from giant.data.transforms import build_features, _WelfordAccumulator from giant.data.dataset import make_event_split, StreamingStepsDataset -from giant.model.network import DenoisingMLP, SecondaryDecoder +from giant.model.network import build_models from giant.train import train as run_training @@ -114,23 +114,22 @@ def run_train_job( "update giant/constants.py if emb_dim changed" ) - stage1_model = DenoisingMLP( - pdg_vocab=len(pdg_map), - mat_vocab=len(mat_map), - hidden_dim=m["hidden_dim"], - n_blocks=m["n_blocks"], - emb_dim=emb_dim, - dropout=m["dropout"], - k_max=K_MAX, - ) - sec_decoder = SecondaryDecoder( - pdg_vocab=len(pdg_map), - mat_vocab=len(mat_map), - hidden_dim=m["hidden_dim"], - n_blocks=m["n_blocks"], - emb_dim=emb_dim, - dropout=m["dropout"], - ) + router_cfg = m["router"] + model_config = { + "pdg_vocab": len(pdg_map), + "mat_vocab": len(mat_map), + "hidden_dim": m["hidden_dim"], + "n_blocks": m["n_blocks"], + "emb_dim": emb_dim, + "dropout": m["dropout"], + "k_max": K_MAX, + "sec_slot_dim": SEC_SLOT_DIM, + "router": dict(router_cfg), + "expert_hidden_dim": router_cfg["expert_hidden_dim"], + "expert_n_blocks": router_cfg["expert_n_blocks"], + } + + stage1_model, sec_decoder = build_models(model_config) echo( f"stage1: {sum(p.numel() for p in stage1_model.parameters()):,} parameters | " f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters" @@ -148,17 +147,6 @@ def run_train_job( ) config.save_config(cfg, out_dir, meta) - model_config = { - "pdg_vocab": len(pdg_map), - "mat_vocab": len(mat_map), - "hidden_dim": m["hidden_dim"], - "n_blocks": m["n_blocks"], - "emb_dim": emb_dim, - "dropout": m["dropout"], - "k_max": K_MAX, - "sec_slot_dim": SEC_SLOT_DIM, - } - run_training( stage1_model=stage1_model, sec_decoder=sec_decoder, @@ -172,6 +160,7 @@ def run_train_job( out_dir=out_dir, lambda_nsec=t.get("lambda_nsec", 0.1), lambda_s2=t.get("lambda_s2", 1.0), + lambda_balance=router_cfg.get("lambda_balance", 0.0), normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()}, pdg_map={str(k): v for k, v in pdg_map.items()}, mat_map={str(k): v for k, v in mat_map.items()}, diff --git a/giant/train.py b/giant/train.py index c21deb7..b30a616 100644 --- a/giant/train.py +++ b/giant/train.py @@ -26,10 +26,12 @@ _METRICS_FIELDS = [ "train_loss_s1", "train_loss_nsec", "train_loss_s2", + "train_loss_balance", "val_loss", "val_loss_s1", "val_loss_nsec", "val_loss_s2", + "val_loss_balance", "lr", "epoch_time_s", ] @@ -108,8 +110,9 @@ def _compute_losses( device: torch.device, lambda_nsec: float, lambda_s2: float, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Compute (total_loss, L_s1, L_nsec, L_s2) for one batch.""" + lambda_balance: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance) for one batch.""" cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) @@ -150,8 +153,21 @@ def _compute_losses( sec_mask, ) + # Optional MoE load-balance auxiliary loss: only present when both stages + # are routed (RoutedDenoisingMLP/RoutedSecondaryDecoder carry `.router`, + # the monolith models don't), computed on cond_cont alone (cheap — no + # trunk compute) so it's reported even when lambda_balance == 0. + if hasattr(stage1_model, "router") and hasattr(sec_decoder, "router"): + l_balance = stage1_model.router.balance_loss( + cond_cont, cond_cat + ) + sec_decoder.router.balance_loss(cond_cont, cond_cat) + else: + l_balance = torch.zeros((), device=device) + total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2 - return total, l_s1, l_nsec, l_s2 + if lambda_balance > 0: + total = total + lambda_balance * l_balance + return total, l_s1, l_nsec, l_s2, l_balance def train( @@ -167,6 +183,7 @@ def train( out_dir: str | Path, lambda_nsec: float = 0.1, lambda_s2: float = 1.0, + lambda_balance: float = 0.0, normalizer_dict: dict | None = None, pdg_map: dict | None = None, mat_map: dict | None = None, @@ -245,6 +262,7 @@ def train( train_s1_sum = 0.0 train_nsec_sum = 0.0 train_s2_sum = 0.0 + train_balance_sum = 0.0 train_n = 0 ema_loss = 0.0 bar = tqdm( @@ -256,7 +274,7 @@ def train( dynamic_ncols=True, ) for batch in bar: - loss, l_s1, l_nsec, l_s2 = _compute_losses( + loss, l_s1, l_nsec, l_s2, l_balance = _compute_losses( stage1_model, sec_decoder, batch, @@ -265,6 +283,7 @@ def train( device, lambda_nsec, lambda_s2, + lambda_balance, ) optimizer.zero_grad() loss.backward() @@ -277,6 +296,7 @@ def train( train_s1_sum += l_s1.item() * B train_nsec_sum += l_nsec.item() * B train_s2_sum += l_s2.item() * B + train_balance_sum += l_balance.item() * B train_n += B ema_loss = ( batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss @@ -299,10 +319,11 @@ def train( val_s1_sum = 0.0 val_nsec_sum = 0.0 val_s2_sum = 0.0 + val_balance_sum = 0.0 val_n = 0 with torch.no_grad(): for batch in val_loader: - loss, l_s1, l_nsec, l_s2 = _compute_losses( + loss, l_s1, l_nsec, l_s2, l_balance = _compute_losses( stage1_model, sec_decoder, batch, @@ -311,12 +332,14 @@ def train( device, lambda_nsec, lambda_s2, + lambda_balance, ) B = batch[0].size(0) val_loss_sum += loss.item() * B val_s1_sum += l_s1.item() * B val_nsec_sum += l_nsec.item() * B val_s2_sum += l_s2.item() * B + val_balance_sum += l_balance.item() * B val_n += B val_loss = val_loss_sum / max(val_n, 1) epoch_time = time.monotonic() - epoch_start @@ -328,7 +351,8 @@ def train( f" train {train_loss:.4f}" f" (s1={train_s1_sum / max(train_n, 1):.3f}" f" nsec={train_nsec_sum / max(train_n, 1):.3f}" - f" s2={train_s2_sum / max(train_n, 1):.3f})" + f" s2={train_s2_sum / max(train_n, 1):.3f}" + f" bal={train_balance_sum / max(train_n, 1):.3f})" f" val {val_loss:.4f}" f" lr {current_lr:.2e} {epoch_time:.1f}s{marker}" ) @@ -339,10 +363,12 @@ def train( "train_loss_s1": train_s1_sum / max(train_n, 1), "train_loss_nsec": train_nsec_sum / max(train_n, 1), "train_loss_s2": train_s2_sum / max(train_n, 1), + "train_loss_balance": train_balance_sum / max(train_n, 1), "val_loss": val_loss, "val_loss_s1": val_s1_sum / max(val_n, 1), "val_loss_nsec": val_nsec_sum / max(val_n, 1), "val_loss_s2": val_s2_sum / max(val_n, 1), + "val_loss_balance": val_balance_sum / max(val_n, 1), "lr": current_lr, "epoch_time_s": epoch_time, } diff --git a/tests/test_router.py b/tests/test_router.py new file mode 100644 index 0000000..36b7e73 --- /dev/null +++ b/tests/test_router.py @@ -0,0 +1,296 @@ +"""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)