giant/cli.py

@@ -1,7 +1,9 @@
 from collections import Counter
 from datetime import date, datetime, timezone
 from enum import Enum
+import math
 from pathlib import Path
+import re
 from typing import Optional
 import uuid as uuid_mod
 
@@ -47,19 +49,43 @@ from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx
 app = typer.Typer(no_args_is_help=True)
 
 
-def _batch_size_estimate_dims(model_cfg: dict) -> tuple[int, int]:
+def _router_total_experts(router_cfg: dict) -> int:
+    """Total expert count for a router config, single-axis or composed.
+
+    A composed router runs one expert per *joint* cell, so its count is the
+    product of the per-axis `axis{i}_n_experts` (mirrors
+    `ComposedRouter.__init__` in giant.model.network); a single-axis router
+    just reports its own `n_experts`.
+    """
+    if router_cfg.get("type") == "composed":
+        axis_counts = {
+            m.group(1): int(v)
+            for k, v in router_cfg.items()
+            if (m := re.match(r"^axis(\d+)_n_experts$", k))
+        }
+        return math.prod(axis_counts.values()) if axis_counts else 1
+    return int(router_cfg.get("n_experts", 1))
+
+
+def _batch_size_estimate_dims(model_cfg: dict, training: bool) -> 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.
+    dims instead when routing is enabled. Training runs the full soft mixture
+    (every expert on the whole batch), so its activation memory scales with
+    the expert count; inference does top-1 dispatch (each row hits one
+    expert), so the batch just partitions across experts and one expert's
+    dims already bound it. estimate_batch_size scales memory linearly with
+    hidden_dim * n_blocks, so the training multiplier folds into n_blocks.
     """
     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),
-        )
+        hidden_dim = model_cfg.get("expert_hidden_dim", 128)
+        n_blocks = model_cfg.get("expert_n_blocks", 3)
+        if training:
+            n_blocks *= _router_total_experts(router_cfg)
+        return hidden_dim, n_blocks
     return model_cfg["hidden_dim"], model_cfg["n_blocks"]
 
 
@@ -332,7 +358,7 @@ 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)
+        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m, training=True)
         try:
             t["batch_size"] = gconfig.estimate_batch_size(
                 est_hidden_dim, est_n_blocks, _device
@@ -463,7 +489,9 @@ def predict(
     model_cfg = ckpt["model_config"]
 
     if batch_size_auto:
-        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(model_cfg)
+        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
+            model_cfg, training=False
+        )
         try:
             batch_size_value = gconfig.estimate_batch_size(
                 est_hidden_dim,