diff --git a/giant/cli.py b/giant/cli.py index 73993fe..2e1d6c2 100644 --- a/giant/cli.py +++ b/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 diff --git a/giant/config.py b/giant/config.py index 54cbcd6..ff4c4fb 100644 --- a/giant/config.py +++ b/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) diff --git a/giant/model/network.py b/giant/model/network.py index 5345747..4f444ba 100644 --- a/giant/model/network.py +++ b/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"), diff --git a/giant/pipeline.py b/giant/pipeline.py index de49898..63a6f98 100644 --- a/giant/pipeline.py +++ b/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"],