From 24a83486ecbb4936b3322f45148d2f4edfb70bbb Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 15 Jul 2026 16:22:25 +0200 Subject: [PATCH] Scale auto batch-size estimate by MoE expert count during training Training runs the full soft mixture (every expert over the whole batch), so routed activation memory scales with the expert count; the old estimate used one expert's dims and would overshoot free VRAM by a factor of n_experts. Fold the expert count into n_blocks for the training path (inference's top-1 dispatch still just partitions the batch, so one expert's dims bound it). Co-Authored-By: Claude Opus 4.8 --- giant/cli.py | 44 ++++++++++++++++++++++++++++++++++++-------- 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index b7a3a5f..82fe028 100644 --- a/giant/cli.py +++ b/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,