Fix router experts silently ignoring --hidden-dim/--n-blocks

expert_hidden_dim/expert_n_blocks were hardcoded to 128/3 in
DEFAULT_CONFIG, independent of model.hidden_dim/n_blocks, so a routed
run always got fixed 128/3-wide experts no matter what --hidden-dim/
--n-blocks was passed. They now default to 0 ("unset"), which
resolve_expert_dims() resolves by inheriting the model dims; an
explicit override still works and now warns when it diverges from
model.hidden_dim/n_blocks, since the checkpoint dir name won't
reflect it.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

giant/cli.py

@@ -87,8 +87,9 @@ def _batch_size_estimate_dims(model_cfg: dict, training: bool) -> tuple[int, int
     """
     router_cfg = model_cfg.get("router")
     if router_cfg and router_cfg.get("enabled"):
-        hidden_dim = model_cfg.get("expert_hidden_dim", 128)
-        n_blocks = model_cfg.get("expert_n_blocks", 3)
+        hidden_dim, n_blocks = gconfig.resolve_expert_dims(
+            router_cfg, model_cfg["hidden_dim"], model_cfg["n_blocks"]
+        )
         if training:
             n_blocks *= _router_total_experts(router_cfg)
         return hidden_dim, n_blocks

giant/config.py

@@ -67,8 +67,13 @@ DEFAULT_CONFIG: dict = {
             "enabled": False,
             "type": "energy",  # selects the Router impl from ROUTER_REGISTRY
             "n_experts": 4,
-            "expert_hidden_dim": 128,
-            "expert_n_blocks": 3,
+            # 0 means "inherit model.hidden_dim/n_blocks" (see
+            # resolve_expert_dims below) — not a fixed 128/3, which silently
+            # ignored --hidden-dim/--n-blocks whenever routing was enabled.
+            # TOML has no null literal to round-trip (same pattern as
+            # critic_lr/wandb_run_name above), hence 0 rather than None.
+            "expert_hidden_dim": 0,
+            "expert_n_blocks": 0,
             "temperature": 0.5,  # energy/pdg-router kwarg
             "learn_centers": True,  # energy/pdg-router kwarg
             "lambda_balance": 0.0,  # optional load-balance aux loss weight
@@ -263,6 +268,22 @@ def merge_cli_overrides(
     return cfg
 
 
+def resolve_expert_dims(
+    router_cfg: dict, hidden_dim: int, n_blocks: int
+) -> tuple[int, int]:
+    """Resolve a router's expert hidden_dim/n_blocks, inheriting from the
+    monolith's when left at the 0 ("unset") sentinel.
+
+    Used by both `giant.pipeline` (to build the checkpoint's `model_config`)
+    and `giant.cli`'s batch-size auto-estimate, so `--hidden-dim`/`--n-blocks`
+    size the experts the same way in both places unless
+    `router.expert_hidden_dim`/`expert_n_blocks` are explicitly overridden.
+    """
+    expert_hidden_dim = router_cfg.get("expert_hidden_dim") or hidden_dim
+    expert_n_blocks = router_cfg.get("expert_n_blocks") or n_blocks
+    return expert_hidden_dim, expert_n_blocks
+
+
 def seed_everything(seed: int) -> None:
     random.seed(seed)
     np.random.seed(seed)

giant/model/network.py

@@ -1130,8 +1130,10 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
         shared = dict(
             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),
+            expert_hidden_dim=model_config.get("expert_hidden_dim")
+            or model_config.get("hidden_dim", 128),
+            expert_n_blocks=model_config.get("expert_n_blocks")
+            or model_config.get("n_blocks", 3),
             emb_dim=model_config.get("emb_dim", EMB_DIM),
             dropout=model_config.get("dropout", 0.1),
             conditioning=model_config.get("conditioning", "embedding"),

giant/pipeline.py

@@ -176,6 +176,25 @@ def run_train_job(
     )
 
     emb_dim = m.get("emb_dim", EMB_DIM)
+    expert_hidden_dim, expert_n_blocks = config.resolve_expert_dims(
+        router_cfg, m["hidden_dim"], m["n_blocks"]
+    )
+    if router_cfg.get("enabled") and (expert_hidden_dim, expert_n_blocks) != (
+        m["hidden_dim"],
+        m["n_blocks"],
+    ):
+        # Only reachable via an explicit router.expert_hidden_dim/n_blocks
+        # override (the 0/"unset" sentinel always resolves to m["hidden_dim"]/
+        # ["n_blocks"] — see resolve_expert_dims), so this is never a false
+        # positive from inheritance, only a deliberate narrow/wide-experts
+        # config the checkpoint dir name (_h{hidden_dim}_b{n_blocks}) won't
+        # reflect.
+        echo(
+            f"  warning: experts are {expert_hidden_dim}x{expert_n_blocks}, "
+            f"different from model.hidden_dim/n_blocks ({m['hidden_dim']}x"
+            f"{m['n_blocks']}) — the checkpoint dir name reflects the latter, "
+            "not the experts actually being trained"
+        )
 
     model_config = {
         "pdg_vocab": len(pdg_map),
@@ -188,8 +207,8 @@ def run_train_job(
         "sec_slot_dim": SEC_SLOT_DIM,
         "conditioning": conditioning,
         "router": dict(router_cfg),
-        "expert_hidden_dim": router_cfg["expert_hidden_dim"],
-        "expert_n_blocks": router_cfg["expert_n_blocks"],
+        "expert_hidden_dim": expert_hidden_dim,
+        "expert_n_blocks": expert_n_blocks,
         # Read by `predict`/`rollout` (which never receive their own --mode
         # flag) to auto-detect which sampler a checkpoint needs.
         "mode": t["mode"],