#!/usr/bin/env python3 """Detect a Gemma 4 variant on BI-150 from its config.json. The unified image serves both members of this family: - Gemma-4-31B-it (dense, 60 layers, 4 global KV heads) -> TP4, 8095, eager off, util 0.90, 4 seqs - GEV-26B-Decide (MoE, 30 layers, 128 experts) -> TP2, 8094, eager on, util 0.95, 2 seqs Both are `Gemma4ForConditionalGeneration` and share the same two vLLM patches, so only the runtime profile differs. Dense vs MoE is decided by the presence of an expert count in the config (top-level or `text_config`). On top of the static profile the script *validates* the profile against the model config, because both variants have a hard TP ceiling: * dense: `num_global_key_value_heads` must be divisible by TP (Gemma-4-31B has 4 -> TP8 would split a KV head into 0.5) * MoE: `moe_intermediate_size / TP` must be divisible by 32 (CoreX `ixfops.moe_w16a16_group_gemm` requires k % 32 == 0; GEV 704/2=352 ok, 704/4=176 -> RuntimeError) Set GEMMA4_SKIP_TP_CHECK=1 to downgrade that validation to a warning. Detected values are printed as shell assignments so run.sh / entrypoint.sh can `eval` them; `--json` is used by verify.sh. """ from __future__ import annotations import argparse import json import os import pathlib import shlex import sys PROFILES = { "gemma4-31b": { "tp": 4, "port": 8095, "served_name": "gemma-4-31B-it", "gpu_count": 4, "max_num_seqs": 4, "gpu_mem_util": "0.90", "enforce_eager": 0, "description": "Gemma 4 31B-it dense (BF16, TP4, 4 GPUs)", }, "gev-26b-decide": { "tp": 2, "port": 8094, "served_name": "GEV-26B-Decide", "gpu_count": 2, "max_num_seqs": 2, "gpu_mem_util": "0.95", "enforce_eager": 1, "description": "AutoTrust GEV-26B-Decide MoE (BF16, TP2, 2 GPUs)", }, } def load_config(model_dir: str) -> dict: path = pathlib.Path(model_dir) / "config.json" if not path.is_file(): raise ValueError(f"config.json not found: {path}") try: return json.loads(path.read_text()) except json.JSONDecodeError as exc: raise ValueError(f"invalid config.json: {exc}") from exc def text_section(cfg: dict) -> dict: """Return the text sub-config if present, else the top-level config.""" text = cfg.get("text_config") if isinstance(text, dict): return text return cfg def expert_count(text: dict, cfg: dict) -> int | None: for source in (text, cfg): for key in ("num_experts", "num_local_experts", "n_routed_experts"): value = source.get(key) if value is not None: return int(value) return None def divisibility_errors(kind: str, tp: int, text: dict, cfg: dict) -> list[str]: """Return human-readable reasons why `tp` cannot be used for this model.""" errors: list[str] = [] for key in ("num_attention_heads", "num_key_value_heads", "num_global_key_value_heads"): value = text.get(key, cfg.get(key)) if value is None: continue value = int(value) if value % tp != 0: errors.append( f"{key}={value} is not divisible by TP={tp} " f"({value} / {tp} = {value / tp:g})" ) if kind == "gev-26b-decide": moe = text.get("moe_intermediate_size", cfg.get("moe_intermediate_size")) if moe is not None: moe = int(moe) if moe % tp != 0: errors.append(f"moe_intermediate_size={moe} is not divisible by TP={tp}") elif (moe // tp) % 32 != 0: errors.append( f"moe_intermediate_size/TP = {moe}/{tp} = {moe // tp} is not a " f"multiple of 32; CoreX ixfops.moe_w16a16_group_gemm requires " f"k % 32 == 0" ) return errors def detect(model_dir: str, tp_override: int | None = None) -> dict: cfg = load_config(model_dir) archs = [str(a) for a in (cfg.get("architectures") or [])] arch = archs[0] if archs else "" if not arch.startswith("Gemma4"): raise ValueError( f"unsupported architecture={arch!r}; expected something like " f"'Gemma4ForConditionalGeneration'" ) text = text_section(cfg) experts = expert_count(text, cfg) if experts is not None and experts > 1: kind = "gev-26b-decide" detection = f"num_experts={experts}" elif experts is not None: # num_experts=1 is how some dense configs spell "no experts". kind = "gemma4-31b" detection = f"num_experts={experts} (dense)" else: name = pathlib.Path(model_dir).name.lower() if "gev" in name: kind = "gev-26b-decide" elif "gemma" in name: kind = "gemma4-31b" else: raise ValueError( "cannot determine Gemma 4 variant: no num_experts in the config " "and the directory name contains neither 'gemma' nor 'gev'" ) detection = "directory name" profile = dict(PROFILES[kind]) tp = int(profile["tp"]) if tp_override is None else int(tp_override) errors = divisibility_errors(kind, tp, text, cfg) if errors and os.environ.get("GEMMA4_SKIP_TP_CHECK") != "1": raise ValueError( f"{kind} cannot run with TP={tp} on this model config:\n " + "\n ".join(errors) + "\nThe base image or the model variant may have changed; inspect " "config.json. If you have already verified another TP manually, " "set GEMMA4_SKIP_TP_CHECK=1 to bypass this check." ) if errors: print( "WARN: TP check skipped by GEMMA4_SKIP_TP_CHECK=1:\n " + "\n ".join(errors), file=sys.stderr, ) if tp_override is not None: profile["tp"] = tp profile["gpu_count"] = tp profile.update( { "kind": kind, "architecture": arch, "num_experts": experts, "num_hidden_layers": text.get("num_hidden_layers", cfg.get("num_hidden_layers")), "hidden_size": text.get("hidden_size", cfg.get("hidden_size")), "detection": detection, "model_dir": str(pathlib.Path(model_dir).resolve()), } ) return profile def main() -> int: parser = argparse.ArgumentParser( description="Detect a Gemma 4 variant (Gemma-4-31B / GEV-26B-Decide) from config.json" ) parser.add_argument("model_dir", help="directory containing config.json") parser.add_argument( "--json", action="store_true", help="print JSON instead of shell assignments", ) parser.add_argument( "--tp", type=int, default=None, help="validate against this TP instead of the profile default", ) args = parser.parse_args() try: profile = detect(args.model_dir, tp_override=args.tp) except ValueError as exc: print(f"ERROR: {exc}", file=sys.stderr) return 2 if args.json: print(json.dumps(profile, ensure_ascii=False, indent=2)) return 0 values = { "DETECTED_MODEL_KIND": profile["kind"], "DETECTED_TP": profile["tp"], "DETECTED_PORT": profile["port"], "DETECTED_SERVED_NAME": profile["served_name"], "DETECTED_GPU_COUNT": profile["gpu_count"], "DETECTED_MAX_NUM_SEQS": profile["max_num_seqs"], "DETECTED_GPU_MEM_UTIL": profile["gpu_mem_util"], "DETECTED_ENFORCE_EAGER": profile["enforce_eager"], "DETECTED_DESCRIPTION": profile["description"], } for key, value in values.items(): print(f"{key}={shlex.quote(str(value))}") return 0 if __name__ == "__main__": raise SystemExit(main())