Unified Gemma 4 family image for Gemma-4-31B and GEV-26B-Decide
All checks were successful
Docker Build and Push / docker (push) Successful in 1m11s
All checks were successful
Docker Build and Push / docker (push) Successful in 1m11s
- detect_model.py picks TP/port/served name/eager/util from config.json - validate TP against the model config (dense global KV heads, MoE k % 32) - one apply_patch.py covers both variants, idempotent - Dockerfile stays generic: no ENV/WORKDIR/EXPOSE/ENTRYPOINT - run.sh/verify.sh follow the llama4-bi150 layout
This commit is contained in:
236
detect_model.py
Executable file
236
detect_model.py
Executable file
@@ -0,0 +1,236 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user