gemma4-bi150/detect_model.py

#!/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())