From bb15e3e76427aa55f14cdd8495e65e395f395d46 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Thu, 13 Aug 2026 00:48:12 +0800 Subject: [PATCH 1/4] feat(glm5): add GLM52Config identity and project resolved config Add GLM5Config.from_hf()/validate() and GLM52Config(GLM5Config) registered as glm_moe_dsa_5_2 with the GLM-5.2 deltas (1M context, nested rope, DSA indexer reuse fields index_topk_freq/indexer_types). GLM-5 defaults are unchanged. Replace the initializer's hand-parsed model config with the resolved config projected to the minimal engine ModelConfig: head_dim comes from qk_head_dim (not the default), and compressed_kv_dim + first_k_dense_replace are carried. Add a headless regression test locking GLM-5 to its prior resolved values, the head_dim projection, GLM-5.2 reading its real config.json, fail-loud on empty/null/inconsistent config, and the unlisted-variant warning. --- batchgen/models/glm/glm5/config.py | 372 ++++++++++++++++++- batchgen/models/glm/glm5/glm5_initializer.py | 42 ++- tests/test_batchgen_model_config.py | 257 +++++++++++++ 3 files changed, 662 insertions(+), 9 deletions(-) create mode 100644 tests/test_batchgen_model_config.py diff --git a/batchgen/models/glm/glm5/config.py b/batchgen/models/glm/glm5/config.py index dd346d997..53b4b51b2 100644 --- a/batchgen/models/glm/glm5/config.py +++ b/batchgen/models/glm/glm5/config.py @@ -17,12 +17,38 @@ - FP8: E4M3, dynamic activation, [128,128] block (same as V3) """ -from dataclasses import dataclass, field +import logging +from dataclasses import dataclass, field, fields from typing import Dict, Any, List, Optional from batchgen.config.model_config import BaseModelConfig from batchgen.config.model_registry import register_config +logger = logging.getLogger(__name__) + + +# Identity keys owned by the config class — deliberately NOT taken from HF so +# GLM-5.2 keeps its distinct model_type. Marked "consumed" so they are not +# reported as silently-dropped keys. +_HANDLED_IDENTITY_KEYS = {"model_type", "architectures"} + +# HF source keys that MUST be present (and non-null) in a real checkpoint +# config.json. An absent/null key here means a truncated or wrong config; we +# fail loud rather than silently backfilling the GLM-5 dataclass defaults (which +# describe the 744B GLM-5 base model and would mis-size the engine). +# num_local_experts is satisfied by either "n_routed_experts" or +# "num_local_experts"; compressed_kv_dim is derived from kv_lora_rank + +# qk_rope_head_dim (both listed here). +_REQUIRED_HF_KEYS = ( + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "qk_head_dim", + "kv_lora_rank", + "qk_rope_head_dim", + "first_k_dense_replace", +) + @register_config("glm_moe_dsa") @dataclass @@ -119,3 +145,347 @@ class GLM5Config(BaseModelConfig): def num_kv_heads(self) -> int: return self.num_key_value_heads + # ------------------------------------------------------------------ # + # Checkpoint resolution (HF config.json -> rich config) + # ------------------------------------------------------------------ # + @staticmethod + def _rope_theta_from_hf(hf: Dict[str, Any]) -> Optional[float]: + """Read rope_theta, handling GLM-5's flat key and GLM-5.2's nested + ``rope_parameters`` block.""" + rope_params = hf.get("rope_parameters") + if isinstance(rope_params, dict) and rope_params.get("rope_theta") is not None: + return rope_params["rope_theta"] + return hf.get("rope_theta") + + @staticmethod + def _common_hf_kwargs(hf: Dict[str, Any]) -> Dict[str, Any]: + """Explicit HF-key -> field mapping shared by GLM-5 and GLM-5.2. + + Only keys actually present in ``hf`` are emitted; everything else falls + back to the dataclass default. Identity fields (model_type, + architectures) are deliberately NOT taken from HF — they are owned by + the config class so GLM-5.2 keeps its distinct model_type. + + Every HF key this mapping *references* (whether present or not) is added + to ``hf["__consumed_keys__"]`` (a scratch set the callers pop back out), + so ``from_hf`` can report checkpoint keys it saw but never mapped. + """ + kwargs: Dict[str, Any] = {} + consumed = hf.setdefault("__consumed_keys__", set()) + consumed.update(_HANDLED_IDENTITY_KEYS) + # mlp_layer_types is not a scalar field but IS load-bearing — it is + # cross-checked against first_k_dense_replace in _assert_mlp_layer_types, + # so treat it as consumed rather than a silently-dropped key. + consumed.add("mlp_layer_types") + + def put(field_name: str, hf_key: str) -> None: + consumed.add(hf_key) + if hf_key in hf and hf[hf_key] is not None: + kwargs[field_name] = hf[hf_key] + + # Core architecture + put("vocab_size", "vocab_size") + put("hidden_size", "hidden_size") + put("intermediate_size", "intermediate_size") + put("num_hidden_layers", "num_hidden_layers") + + # Attention (MLA geometry). NOTE: head_dim and qk_head_dim are distinct + # in HF (192 vs 256); we carry both faithfully. The engine projection + # deliberately uses qk_head_dim, not head_dim. + put("num_attention_heads", "num_attention_heads") + put("num_key_value_heads", "num_key_value_heads") + put("head_dim", "head_dim") + put("qk_head_dim", "qk_head_dim") + put("qk_nope_head_dim", "qk_nope_head_dim") + put("qk_rope_head_dim", "qk_rope_head_dim") + put("v_head_dim", "v_head_dim") + put("kv_lora_rank", "kv_lora_rank") + put("q_lora_rank", "q_lora_rank") + put("attention_bias", "attention_bias") + put("attention_dropout", "attention_dropout") + + # compressed_kv_dim is not stored in HF config.json — derive it from the + # MLA dims (kv_lora_rank + qk_rope_head_dim) when both are available. + kv_lora_rank = hf.get("kv_lora_rank") + qk_rope_head_dim = hf.get("qk_rope_head_dim") + if kv_lora_rank is not None and qk_rope_head_dim is not None: + kwargs["compressed_kv_dim"] = kv_lora_rank + qk_rope_head_dim + + # DSA indexer + put("index_n_heads", "index_n_heads") + put("index_head_dim", "index_head_dim") + put("index_topk", "index_topk") + put("rope_interleave", "rope_interleave") + put("indexer_rope_interleave", "indexer_rope_interleave") + + # MoE. HF stores n_routed_experts; the engine's num_local_experts mirrors + # it (there is no separate num_local_experts key in GLM config.json). + put("n_routed_experts", "n_routed_experts") + consumed.add("num_local_experts") + if hf.get("num_local_experts") is not None: + kwargs["num_local_experts"] = hf["num_local_experts"] + elif hf.get("n_routed_experts") is not None: + kwargs["num_local_experts"] = hf["n_routed_experts"] + put("n_shared_experts", "n_shared_experts") + put("num_experts_per_tok", "num_experts_per_tok") + put("first_k_dense_replace", "first_k_dense_replace") + put("moe_intermediate_size", "moe_intermediate_size") + put("moe_layer_freq", "moe_layer_freq") + + # Router / gating + put("topk_method", "topk_method") + put("n_group", "n_group") + put("topk_group", "topk_group") + put("routed_scaling_factor", "routed_scaling_factor") + put("norm_topk_prob", "norm_topk_prob") + put("scoring_func", "scoring_func") + + # Position encoding (rope_theta may be nested under rope_parameters). + put("max_position_embeddings", "max_position_embeddings") + consumed.update({"rope_theta", "rope_parameters"}) + rope_theta = GLM5Config._rope_theta_from_hf(hf) + if rope_theta is not None: + kwargs["rope_theta"] = rope_theta + + # Normalization & activation + put("rms_norm_eps", "rms_norm_eps") + put("hidden_act", "hidden_act") + + # Quantization + consumed.add("quantization_config") + if hf.get("quantization_config") is not None: + kwargs["quantization_config"] = hf["quantization_config"] + quant_method = hf["quantization_config"].get("quant_method") + if quant_method: + kwargs["quantization"] = quant_method + + # Tokenizer / embeddings + put("bos_token_id", "bos_token_id") + put("eos_token_id", "eos_token_id") + put("pad_token_id", "pad_token_id") + put("tie_word_embeddings", "tie_word_embeddings") + + # Other GLM-specific + put("num_nextn_predict_layers", "num_nextn_predict_layers") + put("ep_size", "ep_size") + + return kwargs + + # ------------------------------------------------------------------ # + # Missing-required / silently-dropped diagnostics (shared by GLM-5.x) + # ------------------------------------------------------------------ # + @staticmethod + def _assert_required_hf_keys(cls_name: str, hf: Dict[str, Any]) -> None: + """Fail loud when a checkpoint config.json omits a required HF key. + + Guards against a truncated / wrong config.json being silently backfilled + with the 744B GLM-5 dataclass defaults. ``num_local_experts`` counts as + present if either ``n_routed_experts`` or ``num_local_experts`` is set. + """ + missing = [ + key for key in _REQUIRED_HF_KEYS + if hf.get(key) is None + ] + if hf.get("n_routed_experts") is None and hf.get("num_local_experts") is None: + missing.append("n_routed_experts|num_local_experts") + if missing: + raise ValueError( + f"{cls_name}.from_hf: checkpoint config.json is missing required " + f"fields {missing}. Refusing to silently backfill GLM-5 defaults " + f"for a truncated/incompatible config." + ) + + @staticmethod + def _warn_dropped_hf_keys(cls_name: str, hf: Dict[str, Any], consumed: set) -> None: + """WARN loudly for checkpoint keys the mapping never referenced. + + These are silently ignored today; surfacing them catches a checkpoint + that carries structurally load-bearing state the resolver does not model + (e.g. a non-contiguous ``mlp_layer_types`` layout). + """ + # HF bookkeeping keys that are intentionally irrelevant to the engine. + ignore = {"__consumed_keys__", "torch_dtype", "dtype", "transformers_version", + "use_cache", "initializer_range", "pretraining_tp", "moe_router_dtype", + "index_topk_pattern"} + dropped = sorted(set(hf) - consumed - ignore) + if dropped: + logger.warning( + "%s.from_hf: checkpoint config.json keys seen but not consumed by " + "the mapping (silently ignored): %s. Verify none are structurally " + "load-bearing for this checkpoint.", + cls_name, dropped, + ) + + @staticmethod + def _assert_mlp_layer_types(cls_name: str, hf: Dict[str, Any], + first_k_dense_replace: int) -> None: + """Cross-check a per-layer dense/sparse pattern against the scalar + ``first_k_dense_replace`` the engine relies on. + + ``is_dense_layer`` treats layers ``< first_k_dense_replace`` as dense and + the rest as sparse (a contiguous dense prefix). If a checkpoint ships a + non-contiguous ``mlp_layer_types`` the scalar cannot represent, fail loud + rather than silently mis-modelling the MoE layout. + """ + layer_types = hf.get("mlp_layer_types") + if not isinstance(layer_types, list) or not layer_types: + return + expected = [ + "dense" if i < first_k_dense_replace else "sparse" + for i in range(len(layer_types)) + ] + if layer_types != expected: + raise ValueError( + f"{cls_name}.from_hf: mlp_layer_types is not a contiguous dense " + f"prefix of length first_k_dense_replace={first_k_dense_replace}; " + f"the engine's scalar first_k_dense_replace cannot represent this " + f"layout. Got {layer_types!r}." + ) + + @classmethod + def from_hf(cls, hf_dict: Dict[str, Any]) -> "GLM5Config": + """Build a rich GLM-5 config from a HuggingFace ``config.json`` dict. + + Maps HF keys onto the config fields explicitly (no silent + name-intersection). Fails loud when a required HF key is absent (rather + than masking it with a family default) and WARNs on checkpoint keys the + mapping never consumed. + """ + cls._assert_required_hf_keys(cls.__name__, hf_dict) + kwargs = cls._common_hf_kwargs(hf_dict) + consumed = hf_dict.pop("__consumed_keys__", set()) + cls._warn_dropped_hf_keys(cls.__name__, hf_dict, consumed) + cls._assert_mlp_layer_types( + cls.__name__, hf_dict, + kwargs.get("first_k_dense_replace", cls.first_k_dense_replace), + ) + # Guard against passing keys the (possibly-subclassed) dataclass does + # not declare. + known = {f.name for f in fields(cls)} + unknown = set(kwargs) - known + if unknown: # pragma: no cover - defensive + logger.debug("Dropping HF keys not declared on %s: %s", cls.__name__, unknown) + kwargs = {k: v for k, v in kwargs.items() if k in known} + return cls(**kwargs) + + def validate(self) -> None: + """Fail loud on missing required fields or self-inconsistency. + + This is the last line of defence before the rich config is projected + into the engine's minimal ModelConfig. + """ + required = { + "model_type": self.model_type, + "num_hidden_layers": self.num_hidden_layers, + "num_attention_heads": self.num_attention_heads, + "num_key_value_heads": self.num_key_value_heads, + "num_local_experts": self.num_local_experts, + "qk_head_dim": self.qk_head_dim, + "kv_lora_rank": self.kv_lora_rank, + "qk_rope_head_dim": self.qk_rope_head_dim, + "compressed_kv_dim": self.compressed_kv_dim, + "first_k_dense_replace": self.first_k_dense_replace, + } + missing = [name for name, value in required.items() if value is None] + if missing: + raise ValueError( + f"{type(self).__name__}.validate: missing required fields " + f"{missing} (model_type={self.model_type!r}). Config resolution " + f"produced an incomplete config; refusing to build the engine." + ) + + if self.num_key_value_heads > self.num_attention_heads: + raise ValueError( + f"{type(self).__name__}.validate: num_key_value_heads=" + f"{self.num_key_value_heads} > num_attention_heads=" + f"{self.num_attention_heads}." + ) + + expected_compressed = self.kv_lora_rank + self.qk_rope_head_dim + if self.compressed_kv_dim != expected_compressed: + raise ValueError( + f"{type(self).__name__}.validate: compressed_kv_dim=" + f"{self.compressed_kv_dim} != kv_lora_rank + qk_rope_head_dim=" + f"{self.kv_lora_rank} + {self.qk_rope_head_dim} = " + f"{expected_compressed}." + ) + + if self.first_k_dense_replace >= self.num_hidden_layers: + raise ValueError( + f"{type(self).__name__}.validate: first_k_dense_replace=" + f"{self.first_k_dense_replace} must be < num_hidden_layers=" + f"{self.num_hidden_layers}." + ) + + for name in ("num_hidden_layers", "num_attention_heads", + "num_local_experts", "qk_head_dim"): + value = getattr(self, name) + if value <= 0: + raise ValueError( + f"{type(self).__name__}.validate: {name}={value} must be > 0." + ) + + +@register_config("glm_moe_dsa_5_2") +@dataclass +class GLM52Config(GLM5Config): + """GLM-5.2 (GlmMoeDsaForCausalLM) configuration. + + GLM-5.2 shares the glm_moe_dsa MODEL graph with GLM-5 but is given its own + config identity (``model_type="glm_moe_dsa_5_2"``) so it can carry distinct + DSA-indexer scheduling knobs and a longer native context window without + perturbing GLM-5. Notable checkpoint differences vs GLM-5: + + - max_position_embeddings = 1,048,576 (vs 202,752) + - rope_theta = 8,000,000, delivered nested under ``rope_parameters`` + - DSA indexer gains ``index_topk_freq`` / ``index_skip_topk_offset`` / + ``index_share_for_mtp_iteration`` and a per-layer ``indexer_types`` list; + GLM-5 lacks these. + """ + + # ==================== Identity ==================== + model_type: str = "glm_moe_dsa_5_2" + + # ==================== Position Encoding ==================== + max_position_embeddings: int = 1048576 + rope_theta: float = 8000000.0 + + # ==================== DSA (GLM-5.2-only indexer scheduling) ==================== + index_topk_freq: int = 4 + index_skip_topk_offset: int = 3 + index_share_for_mtp_iteration: bool = True + indexer_types: Optional[List[str]] = None + + @classmethod + def from_hf(cls, hf_dict: Dict[str, Any]) -> "GLM52Config": + """Build a rich GLM-5.2 config from a HuggingFace ``config.json`` dict. + + Extends the shared GLM mapping with the GLM-5.2-only indexer fields, and + applies the same fail-loud / warn-on-drop diagnostics as GLM-5. + """ + cls._assert_required_hf_keys(cls.__name__, hf_dict) + kwargs = cls._common_hf_kwargs(hf_dict) + consumed = hf_dict.get("__consumed_keys__", set()) + + def put(field_name: str, hf_key: str) -> None: + consumed.add(hf_key) + if hf_key in hf_dict and hf_dict[hf_key] is not None: + kwargs[field_name] = hf_dict[hf_key] + + put("index_topk_freq", "index_topk_freq") + put("index_skip_topk_offset", "index_skip_topk_offset") + put("index_share_for_mtp_iteration", "index_share_for_mtp_iteration") + put("indexer_types", "indexer_types") + + consumed = hf_dict.pop("__consumed_keys__", set()) + cls._warn_dropped_hf_keys(cls.__name__, hf_dict, consumed) + cls._assert_mlp_layer_types( + cls.__name__, hf_dict, + kwargs.get("first_k_dense_replace", cls.first_k_dense_replace), + ) + + known = {f.name for f in fields(cls)} + kwargs = {k: v for k, v in kwargs.items() if k in known} + return cls(**kwargs) + + diff --git a/batchgen/models/glm/glm5/glm5_initializer.py b/batchgen/models/glm/glm5/glm5_initializer.py index 799cc1c86..cb9d0ddc6 100644 --- a/batchgen/models/glm/glm5/glm5_initializer.py +++ b/batchgen/models/glm/glm5/glm5_initializer.py @@ -22,6 +22,7 @@ import torch from batchgen.config.config import EngineConfig, ModelConfig +from batchgen.config.batchgen_model_config import BatchGenModelConfig from .configuration_glm5 import Glm5Config from .set_basic_config import set_basic_config from .planner import GLM5Planner @@ -40,6 +41,11 @@ def __init__(self, input_arguments): self.loaded_model_config._name_or_path = input_arguments.huggingface_ckpt_name self.loaded_model_config.architectures = ["GlmMoeDsaForCausalLM"] + self.model_name = input_arguments.huggingface_ckpt_name + # Local checkpoint dir holding config.json (pre-downloaded model files). + # None for a bare HF id -> resolver falls back to rich defaults. + self.checkpoint_path = input_arguments.get("cache_dir", None) + self.host_kv_cache_size = input_arguments.host_kv_cache_size self.host_kv_cache_byte_size = input_arguments.host_kv_cache_size * (1024**3) self.global_kv_cache_size_gb = input_arguments.global_host_kv_cache_size_gb @@ -156,15 +162,35 @@ def _default_engine_config(self): } def _parse_model_config(self): + """Resolve the rich, checkpoint-backed config via the single decoupled + resolver, then PROJECT it into the minimal engine ModelConfig. + + The rich config (GLM5Config / GLM52Config) is the single source of + truth; the engine's minimal ModelConfig is duck-typed on the C++ side + (core/utils.cpp parse_model_config reads model_type, num_hidden_layers, + num_local_experts, num_attention_heads, num_key_value_heads, head_dim). + + Traps honored here: + * head_dim <- rich.qk_head_dim (256), NOT rich.head_dim (64). The C++ + attention geometry expects the MLA qk head dim. + * compressed_kv_dim and first_k_dense_replace are NOT declared fields of + the minimal ModelConfig but ARE read downstream (this file @ + _default_engine_config, glm5_parameter_server). Set them as attributes + so Init does not crash with AttributeError. + """ + rich = BatchGenModelConfig.resolve(self.model_name, self.checkpoint_path) + model_config = ModelConfig() - model_config.model_type = "glm_moe_dsa" - model_config.num_hidden_layers = 78 - model_config.num_local_experts = 256 - model_config.num_attention_heads = 64 - model_config.num_key_value_heads = 64 - model_config.head_dim = 256 # qk_nope + qk_rope = 192 + 64 - model_config.compressed_kv_dim = 576 # kv_lora_rank + qk_rope_head_dim - model_config.first_k_dense_replace = 3 + model_config.model_type = rich.model_type + model_config.num_hidden_layers = rich.num_hidden_layers + model_config.num_local_experts = rich.num_local_experts + model_config.num_attention_heads = rich.num_attention_heads + model_config.num_key_value_heads = rich.num_key_value_heads + # TRAP: engine head_dim is the MLA qk head dim (256), not rich.head_dim (64). + model_config.head_dim = rich.qk_head_dim + # Not declared on ModelConfig but read downstream — set as attributes. + model_config.compressed_kv_dim = rich.compressed_kv_dim # kv_lora_rank + qk_rope_head_dim + model_config.first_k_dense_replace = rich.first_k_dense_replace return model_config def Init(self, weights_storage): diff --git a/tests/test_batchgen_model_config.py b/tests/test_batchgen_model_config.py new file mode 100644 index 000000000..506b9db58 --- /dev/null +++ b/tests/test_batchgen_model_config.py @@ -0,0 +1,257 @@ +# ---------------------------------------------------------------------------- # +# BatchGen # +# copyright (c) EfficientMoE team 2025 # +# ---------------------------------------------------------------------------- # + +"""Headless regression tests for the decoupled BatchGenModelConfig resolver. + +These tests must run WITHOUT torch / the C++ engine. The production +``batchgen.config`` package ``__init__`` eagerly imports the tokenizer stack +(torch) and the model registry (whose auto-import triggers an engine JIT +build), so — mirroring tests/test_glm5_planner.py — we pre-register stub +package modules and load the specific source files directly via importlib. + +Coverage: + * Regression lock: resolve() for GLM-5-FP8 reproduces the exact values the + old glm5_initializer._parse_model_config hardcoded, AND the engine + projection maps head_dim <- qk_head_dim (256, the head_dim TRAP). + * resolve() for GLM-5.2-FP8 reads the real checkpoint config.json and yields + the correct dims + GLM-5.2-only indexer fields + nested rope_theta, with a + distinct model_type. + * Pattern ordering: GLM-5.2 resolves to GLM52Config, not GLM5Config. +""" + +import importlib.util +import os +import sys +import types +from pathlib import Path + +import pytest + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_CONFIG_DIR = _REPO_ROOT / "batchgen" / "config" +_GLM5_CONFIG = _REPO_ROOT / "batchgen" / "models" / "glm" / "glm5" / "config.py" + +# Point GLM52_CKPT_DIR at a local GLM-5.2-FP8 checkout to exercise the +# real-config resolution tests; unset, they skip (as they do in CI). +_GLM52_CKPT = os.environ.get("GLM52_CKPT_DIR", "") + +# Old _parse_model_config hardcoded values — the regression baseline. +_LEGACY_GLM5 = { + "model_type": "glm_moe_dsa", + "num_hidden_layers": 78, + "num_local_experts": 256, + "num_attention_heads": 64, + "num_key_value_heads": 64, + "head_dim": 256, # == qk_head_dim (the projection uses qk_head_dim) + "compressed_kv_dim": 576, + "first_k_dense_replace": 3, +} + + +def _load_module(fqname: str, path: Path): + """Exec a source file as ``fqname`` without running package __init__.""" + if fqname in sys.modules: + return sys.modules[fqname] + spec = importlib.util.spec_from_file_location(fqname, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[fqname] = module + spec.loader.exec_module(module) + return module + + +def _stub_pkg(fqname: str, dir_path: Path): + if fqname not in sys.modules: + pkg = types.ModuleType(fqname) + pkg.__path__ = [str(dir_path)] + sys.modules[fqname] = pkg + + +def _bootstrap(): + """Wire up a torch-free import graph and return (BatchGenModelConfig, mod).""" + # Namespace stubs so no real package __init__ (torch/engine) runs. + _stub_pkg("batchgen", _REPO_ROOT / "batchgen") + _stub_pkg("batchgen.config", _CONFIG_DIR) + _stub_pkg("batchgen.models", _REPO_ROOT / "batchgen" / "models") + _stub_pkg("batchgen.models.glm", _REPO_ROOT / "batchgen" / "models" / "glm") + _stub_pkg("batchgen.models.glm.glm5", _REPO_ROOT / "batchgen" / "models" / "glm" / "glm5") + + # Minimal, torch-free stand-in for model_registry (the real one auto-imports + # every family config, dragging the engine). Provides just what glm5/config + # needs: a working register_config decorator + a registry dict. + if "batchgen.config.model_registry" not in sys.modules: + reg = types.ModuleType("batchgen.config.model_registry") + reg.CONFIG_REGISTRY = {} + + def register_config(model_type): + def deco(cls): + reg.CONFIG_REGISTRY[model_type] = cls + return cls + return deco + + reg.register_config = register_config + sys.modules["batchgen.config.model_registry"] = reg + + _load_module("batchgen.config.model_config", _CONFIG_DIR / "model_config.py") + # Pre-load glm5 config so the resolver's import_module hits the cache. + _load_module("batchgen.models.glm.glm5.config", _GLM5_CONFIG) + bmc = _load_module( + "batchgen.config.batchgen_model_config", + _CONFIG_DIR / "batchgen_model_config.py", + ) + return bmc.BatchGenModelConfig, bmc + + +BatchGenModelConfig, _bmc_mod = _bootstrap() + + +def _project(rich): + """Replicate glm5_initializer's engine projection (the load-bearing part). + + Mirrors the minimal ModelConfig the C++ engine reads. Crucially maps + head_dim <- qk_head_dim, not rich.head_dim. + """ + return { + "model_type": rich.model_type, + "num_hidden_layers": rich.num_hidden_layers, + "num_local_experts": rich.num_local_experts, + "num_attention_heads": rich.num_attention_heads, + "num_key_value_heads": rich.num_key_value_heads, + "head_dim": rich.qk_head_dim, # TRAP: qk_head_dim, not head_dim + "compressed_kv_dim": rich.compressed_kv_dim, + "first_k_dense_replace": rich.first_k_dense_replace, + } + + +def test_resolver_imports_without_torch(): + assert "torch" not in sys.modules, ( + "batchgen_model_config must import cleanly headless; something dragged torch." + ) + + +def test_glm5_fp8_regression_matches_legacy_hardcoded_values(): + # No GLM-5-FP8 checkpoint on disk -> resolver uses GLM5Config defaults, + # which must reproduce the old _parse_model_config hardcoded values. + rich = BatchGenModelConfig.resolve("zai-org/GLM-5-FP8", checkpoint_path=None) + assert type(rich).__name__ == "GLM5Config" + projected = _project(rich) + assert projected == _LEGACY_GLM5 + + +def test_glm5_head_dim_trap_projection_uses_qk_head_dim(): + rich = BatchGenModelConfig.resolve("GLM-5-FP8", checkpoint_path=None) + # rich.head_dim is 64; the engine must receive qk_head_dim (256). + assert rich.head_dim != 256 + assert rich.qk_head_dim == 256 + assert _project(rich)["head_dim"] == 256 + + +def test_glm52_pattern_resolves_to_glm52_config_not_glm5(): + target = BatchGenModelConfig._match_variant("zai-org/GLM-5.2-FP8") + assert target == ("batchgen.models.glm.glm5.config", "GLM52Config") + # And GLM-5.1 / GLM-5 do not get swallowed by GLM-5.2. + assert BatchGenModelConfig._match_variant("GLM-5.1-FP8")[1] == "GLM5Config" + assert BatchGenModelConfig._match_variant("GLM-5-FP8")[1] == "GLM5Config" + + +@pytest.mark.skipif( + not (Path(_GLM52_CKPT) / "config.json").exists(), + reason="GLM-5.2-FP8 checkpoint config.json not available", +) +def test_glm52_fp8_reads_checkpoint_dims(): + rich = BatchGenModelConfig.resolve("zai-org/GLM-5.2-FP8", checkpoint_path=_GLM52_CKPT) + assert type(rich).__name__ == "GLM52Config" + + # Distinct config identity. + assert rich.model_type == "glm_moe_dsa_5_2" + + # Core dims read from the real config.json. + assert rich.num_hidden_layers == 78 + assert rich.num_local_experts == 256 # from n_routed_experts + assert rich.n_routed_experts == 256 + assert rich.num_attention_heads == 64 + assert rich.num_key_value_heads == 64 + assert rich.qk_head_dim == 256 + assert rich.head_dim == 192 # HF head_dim (distinct from qk_head_dim) + assert rich.kv_lora_rank == 512 + assert rich.qk_rope_head_dim == 64 + assert rich.compressed_kv_dim == 576 # kv_lora_rank + qk_rope_head_dim + assert rich.first_k_dense_replace == 3 + + # GLM-5.2 specifics. + assert rich.max_position_embeddings == 1048576 + assert rich.rope_theta == 8000000 # nested under rope_parameters + assert rich.index_topk_freq == 4 + assert rich.index_skip_topk_offset == 3 + assert rich.indexer_types is not None and len(rich.indexer_types) == 78 + + # Engine projection still uses qk_head_dim. + assert _project(rich)["head_dim"] == 256 + + +def test_unlisted_variant_warns_and_falls_back(caplog): + import logging as _logging + with caplog.at_level(_logging.WARNING): + rich = BatchGenModelConfig.resolve("Some-Unknown-Model", checkpoint_path=None) + assert type(rich).__name__ == "GLM5Config" + assert any("matched no supported variant" in r.message for r in caplog.records) + + +def test_unlisted_glm5_minor_warns_before_falling_back_to_base(caplog): + """An unlisted GLM-5.x minor must warn loudly, not silently bind to base.""" + import logging as _logging + with caplog.at_level(_logging.WARNING): + rich = BatchGenModelConfig.resolve("zai/GLM-5.3", checkpoint_path=None) + assert type(rich).__name__ == "GLM5Config" # base fallback + assert any("unlisted GLM-5.3 variant" in r.message for r in caplog.records) + + +def test_glm5_superstring_warns(caplog): + """GLM-50 is a superstring of GLM-5 and must warn, not bind silently.""" + import logging as _logging + with caplog.at_level(_logging.WARNING): + BatchGenModelConfig.resolve("GLM-50-foo", checkpoint_path=None) + assert any("superstring" in r.message for r in caplog.records) + + +def _write_config(tmp_path, data): + import json + (tmp_path / "config.json").write_text(json.dumps(data)) + return str(tmp_path) + + +def test_empty_config_fails_loud(tmp_path): + """A truncated config.json ({}) must FAIL, not silently backfill defaults.""" + ckpt = _write_config(tmp_path, {}) + with pytest.raises(ValueError, match="missing required fields"): + BatchGenModelConfig.resolve("GLM-5.2-FP8", checkpoint_path=ckpt) + + +def test_null_required_field_fails_loud(tmp_path): + """A required field explicitly null in config.json must FAIL loud.""" + import json + if not (Path(_GLM52_CKPT) / "config.json").exists(): + pytest.skip("GLM-5.2-FP8 checkpoint config.json not available") + data = json.loads((Path(_GLM52_CKPT) / "config.json").read_text()) + data["num_hidden_layers"] = None + ckpt = _write_config(tmp_path, data) + with pytest.raises(ValueError, match="missing required fields"): + BatchGenModelConfig.resolve("GLM-5.2-FP8", checkpoint_path=ckpt) + + +def test_noncontiguous_mlp_layer_types_fails_loud(tmp_path): + """A non-contiguous dense/sparse layout the scalar first_k_dense_replace + cannot represent must FAIL loud rather than be silently mis-modelled.""" + import json + if not (Path(_GLM52_CKPT) / "config.json").exists(): + pytest.skip("GLM-5.2-FP8 checkpoint config.json not available") + data = json.loads((Path(_GLM52_CKPT) / "config.json").read_text()) + layer_types = list(data["mlp_layer_types"]) + layer_types[0] = "sparse" # break the contiguous dense prefix + layer_types[5] = "dense" + data["mlp_layer_types"] = layer_types + ckpt = _write_config(tmp_path, data) + with pytest.raises(ValueError, match="mlp_layer_types"): + BatchGenModelConfig.resolve("GLM-5.2-FP8", checkpoint_path=ckpt) From 27f151adfce8ec14fef4ee8db42a1ea6552a88a8 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Thu, 13 Aug 2026 00:48:12 +0800 Subject: [PATCH 2/4] feat(glm5): add dsa_layer_skips_topk schedule helper Add module-level dsa_layer_skips_topk(config, layer_id) and assert_indexer_schedule_consistent to the GLM-5 config. The helper is value-based: it keys off index_topk_freq presence, so it works on any config object and returns False when the frequency is absent/None/1, leaving GLM-5's uniform per-layer recompute unchanged. The frequency/offset formula is authoritative; the checkpoint indexer_types list is a startup cross-check oracle that fails loud on disagreement. --- batchgen/models/glm/glm5/config.py | 91 ++++++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) diff --git a/batchgen/models/glm/glm5/config.py b/batchgen/models/glm/glm5/config.py index 53b4b51b2..5bb3c86ff 100644 --- a/batchgen/models/glm/glm5/config.py +++ b/batchgen/models/glm/glm5/config.py @@ -489,3 +489,94 @@ def put(field_name: str, hf_key: str) -> None: return cls(**kwargs) +# ---------------------------------------------------------------------------- # +# DSA indexer top-k reuse schedule # +# ---------------------------------------------------------------------------- # +# GLM-5.2 does not run the DSA indexer on every layer. Only "full" layers carry +# indexer weights and recompute the sparse top-k token selection; "shared" +# layers reuse the most recent preceding full layer's top-k indices. GLM-5 (base +# GLM5Config) has no such schedule — every layer is a full layer (uniform +# recompute), so the helpers below are no-ops for it and GLM-5 behavior is +# unchanged. +# +# Ported from SGLang's ``dsa_layer_skips_topk`` (sglang/srt/configs/model_config.py): +# the freq/offset formula is authoritative; the per-layer ``indexer_types`` list +# shipped in the GLM-5.2 checkpoint is used only as a startup cross-check oracle. + + +def dsa_layer_skips_topk(config, layer_id: int) -> bool: + """Whether ``layer_id`` reuses the previous full layer's DSA top-k indices. + + Value-based (works on ANY config object — the rich :class:`GLM52Config` or + the ``configuration_glm5.Glm5Config`` PretrainedConfig used to build the + model graph): the decision keys off the presence of a positive + ``index_topk_freq``. GLM-5 / GLM-5.1 configs have no such field (or it is + ``None`` / 1), so every layer is a full layer (uniform recompute) and this + returns ``False`` — GLM-5 behavior unchanged. + + Authoritative formula (matches SGLang ``dsa_layer_skips_topk``): + - if ``index_topk_pattern`` is set: ``pattern[layer_id] == "S"`` + - else with ``freq = index_topk_freq`` (>=1) and optional + ``offset = index_skip_topk_offset``: + ``max(layer_id - offset + 1, 0) % freq != 0`` (offset present) + ``max(layer_id - 1, 0) % freq != 0`` (offset absent) + """ + pattern = getattr(config, "index_topk_pattern", None) + if pattern is not None: + return layer_id < len(pattern) and pattern[layer_id] == "S" + + freq = getattr(config, "index_topk_freq", None) + if freq is None: + freq = 1 + if freq <= 0: + raise ValueError(f"index_topk_freq must be positive, got {freq}") + if freq == 1: + # Uniform recompute (GLM-5): no layer is ever skipped. + return False + + offset = getattr(config, "index_skip_topk_offset", None) + if offset is not None: + if offset <= 0: + raise ValueError( + "index_skip_topk_offset must be positive; offset <= 0 marks " + "layer 0 as skip_topk with no prior topk to reuse" + ) + return max(layer_id - offset + 1, 0) % freq != 0 + + return max(layer_id - 1, 0) % freq != 0 + + +def assert_indexer_schedule_consistent(config) -> None: + """Fail loud if the freq/offset schedule disagrees with ``indexer_types``. + + GLM-5.2 checkpoints ship a per-layer ``indexer_types`` list ("full" / + "shared"). It must agree layer-for-layer with :func:`dsa_layer_skips_topk` + (shared == skips top-k). This is a startup guard so a future checkpoint whose + list diverges from the formula fails immediately instead of silently + mis-scheduling. No-op when ``indexer_types`` is absent/``None`` (e.g. GLM-5). + """ + indexer_types = getattr(config, "indexer_types", None) + if indexer_types is None: + return + num_layers = config.num_hidden_layers + if len(indexer_types) != num_layers: + raise ValueError( + f"indexer_types has {len(indexer_types)} entries but " + f"num_hidden_layers={num_layers} " + f"(model_type={getattr(config, 'model_type', '?')!r})" + ) + mismatched = [ + layer_id + for layer_id in range(num_layers) + if dsa_layer_skips_topk(config, layer_id) != (indexer_types[layer_id] == "shared") + ] + if mismatched: + raise ValueError( + "DSA indexer schedule mismatch: dsa_layer_skips_topk disagrees with " + f"indexer_types at layers {mismatched} " + f"(index_topk_freq={getattr(config, 'index_topk_freq', None)}, " + f"index_skip_topk_offset={getattr(config, 'index_skip_topk_offset', None)})" + ) + + + From 5b73333eea7df484b3f33c0130e21cff68fb58cd Mon Sep 17 00:00:00 2001 From: TairanXU Date: Thu, 13 Aug 2026 00:48:12 +0800 Subject: [PATCH 3/4] refactor(glm5): drop HF PretrainedConfig, use internal config end-to-end GLM-5 was the lone model building its graph from an HF transformers.PretrainedConfig; every other model already uses its internal BaseModelConfig dataclass. Migrate GLM-5 to match so GLM-5.2's DSA schedule fields (index_topk_freq/indexer_types) reach the graph without HF plumbing. - config.py: add rope_parameters + rope_type (with from_hf mapping), the only two graph-read attrs the internal config lacked versus the HF class. - glm5_initializer.py: build from the resolved internal config; reuse it in _parse_model_config (no double-resolve); assert the indexer schedule at build (a no-op for GLM-5). - glm5_parameter_server.py: build the graph from the resolved config. - model.py: import the internal GLM5Config for the type hints and the dsa_layer_skips_topk helper for the upcoming indexer-reuse step. - delete configuration_glm5.py. --- batchgen/models/glm/glm5/config.py | 21 ++- .../models/glm/glm5/configuration_glm5.py | 143 ------------------ batchgen/models/glm/glm5/glm5_initializer.py | 21 ++- .../models/glm/glm5/glm5_parameter_server.py | 11 +- batchgen/models/glm/glm5/model.py | 6 +- 5 files changed, 42 insertions(+), 160 deletions(-) delete mode 100644 batchgen/models/glm/glm5/configuration_glm5.py diff --git a/batchgen/models/glm/glm5/config.py b/batchgen/models/glm/glm5/config.py index 5bb3c86ff..c43dbef12 100644 --- a/batchgen/models/glm/glm5/config.py +++ b/batchgen/models/glm/glm5/config.py @@ -117,6 +117,11 @@ class GLM5Config(BaseModelConfig): max_position_embeddings: int = 202752 rope_theta: float = 1000000.0 rope_scaling: Optional[Dict[str, Any]] = None # No YaRN + # GLM checkpoints nest rope params under `rope_parameters`; the model graph + # reads both the raw block and `rope_type` (surfaced for parity with the + # former HF PretrainedConfig; defaults keep GLM-5 behavior unchanged). + rope_parameters: Optional[Dict[str, Any]] = None + rope_type: str = "default" # ==================== Normalization & Activation ==================== rms_norm_eps: float = 1e-5 @@ -246,6 +251,13 @@ def put(field_name: str, hf_key: str) -> None: rope_theta = GLM5Config._rope_theta_from_hf(hf) if rope_theta is not None: kwargs["rope_theta"] = rope_theta + # Carry the raw rope_parameters block and rope_type for the model graph + # (parity with the former HF PretrainedConfig, which surfaced both). + rope_params = hf.get("rope_parameters") + if isinstance(rope_params, dict): + kwargs["rope_parameters"] = rope_params + if rope_params.get("rope_type") is not None: + kwargs["rope_type"] = rope_params["rope_type"] # Normalization & activation put("rms_norm_eps", "rms_norm_eps") @@ -508,11 +520,10 @@ def dsa_layer_skips_topk(config, layer_id: int) -> bool: """Whether ``layer_id`` reuses the previous full layer's DSA top-k indices. Value-based (works on ANY config object — the rich :class:`GLM52Config` or - the ``configuration_glm5.Glm5Config`` PretrainedConfig used to build the - model graph): the decision keys off the presence of a positive - ``index_topk_freq``. GLM-5 / GLM-5.1 configs have no such field (or it is - ``None`` / 1), so every layer is a full layer (uniform recompute) and this - returns ``False`` — GLM-5 behavior unchanged. + the base :class:`GLM5Config` used to build the model graph): the decision + keys off the presence of a positive ``index_topk_freq``. GLM-5 / GLM-5.1 + configs have no such field (or it is ``None`` / 1), so every layer is a full + layer (uniform recompute) and this returns ``False`` — GLM-5 unchanged. Authoritative formula (matches SGLang ``dsa_layer_skips_topk``): - if ``index_topk_pattern`` is set: ``pattern[layer_id] == "S"`` diff --git a/batchgen/models/glm/glm5/configuration_glm5.py b/batchgen/models/glm/glm5/configuration_glm5.py deleted file mode 100644 index 3b524c836..000000000 --- a/batchgen/models/glm/glm5/configuration_glm5.py +++ /dev/null @@ -1,143 +0,0 @@ -# ---------------------------------------------------------------------------- # -# BatchGen # -# copyright (c) EfficientMoE team 2025 # -# # -# licensed under the apache license, version 2.0 (the "license"); # -# you may not use this file except in compliance with the license. # -# # -# you may obtain a copy of the license at # -# # -# http://www.apache.org/licenses/license-2.0 # -# # -# unless required by applicable law or agreed to in writing, software # -# distributed under the license is distributed on an "as is" basis, # -# without warranties or conditions of any kind, either express or implied. # -# see the license for the specific language governing permissions and # -# limitations under the license. # -# ---------------------------------------------------------------------------- # - -"""GLM-5 HuggingFace-style PretrainedConfig for checkpoint loading. - -This mirrors the config.json structure from zai-org/GLM-5-FP8. -Used by model.py to instantiate the model with correct dimensions. -""" - -from transformers.configuration_utils import PretrainedConfig - - -class Glm5Config(PretrainedConfig): - model_type = "glm_moe_dsa" - - def __init__( - self, - vocab_size=154880, - hidden_size=6144, - intermediate_size=12288, - moe_intermediate_size=2048, - num_hidden_layers=78, - num_attention_heads=64, - num_key_value_heads=64, - head_dim=64, - qk_head_dim=256, - qk_nope_head_dim=192, - qk_rope_head_dim=64, - v_head_dim=256, - q_lora_rank=2048, - kv_lora_rank=512, - rope_theta=1000000.0, - rope_interleave=True, - indexer_rope_interleave=True, - rope_scaling=None, - attention_bias=False, - attention_dropout=0.0, - n_routed_experts=256, - n_shared_experts=1, - num_experts_per_tok=8, - first_k_dense_replace=3, - moe_layer_freq=1, - n_group=1, - topk_group=1, - topk_method="noaux_tc", - norm_topk_prob=True, - routed_scaling_factor=2.5, - scoring_func="sigmoid", - index_n_heads=32, - index_head_dim=128, - index_topk=2048, - use_dense_mla=False, - hidden_act="silu", - max_position_embeddings=202752, - rms_norm_eps=1e-5, - initializer_range=0.02, - tie_word_embeddings=False, - num_nextn_predict_layers=1, - ep_size=1, - pad_token_id=154820, - bos_token_id=None, - eos_token_id=None, - **kwargs, - ): - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.moe_intermediate_size = moe_intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim - self.qk_head_dim = qk_head_dim - self.qk_nope_head_dim = qk_nope_head_dim - self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim - self.q_lora_rank = q_lora_rank - self.kv_lora_rank = kv_lora_rank - # GLM-5 checkpoints nest rope_theta under `rope_parameters` — unwrap it - # here so a future checkpoint that overrides the base does not silently - # fall back to the kwarg default. Top-level `rope_theta` (if ever - # passed) takes precedence. Also surface `rope_type`. - _rope_params = kwargs.get("rope_parameters") - if isinstance(_rope_params, dict): - rope_theta = _rope_params.get("rope_theta", rope_theta) - self.rope_type = _rope_params.get("rope_type", "default") - else: - self.rope_type = "default" - self.rope_parameters = _rope_params - self.rope_theta = rope_theta - self.rope_interleave = rope_interleave - self.indexer_rope_interleave = indexer_rope_interleave - self.rope_scaling = rope_scaling - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.n_routed_experts = n_routed_experts - self.n_shared_experts = n_shared_experts - self.num_experts_per_tok = num_experts_per_tok - self.first_k_dense_replace = first_k_dense_replace - self.moe_layer_freq = moe_layer_freq - self.n_group = n_group - self.topk_group = topk_group - self.topk_method = topk_method - self.norm_topk_prob = norm_topk_prob - self.routed_scaling_factor = routed_scaling_factor - self.scoring_func = scoring_func - self.index_n_heads = index_n_heads - self.index_head_dim = index_head_dim - self.index_topk = index_topk - self.use_dense_mla = bool(use_dense_mla) - self.hidden_act = hidden_act - self.max_position_embeddings = max_position_embeddings - self.rms_norm_eps = rms_norm_eps - self.initializer_range = initializer_range - self.num_nextn_predict_layers = num_nextn_predict_layers - self.ep_size = ep_size - - # Derived - self.num_local_experts = n_routed_experts - self.compressed_kv_dim = kv_lora_rank + qk_rope_head_dim # 576 - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) diff --git a/batchgen/models/glm/glm5/glm5_initializer.py b/batchgen/models/glm/glm5/glm5_initializer.py index cb9d0ddc6..2c338ee75 100644 --- a/batchgen/models/glm/glm5/glm5_initializer.py +++ b/batchgen/models/glm/glm5/glm5_initializer.py @@ -23,7 +23,7 @@ from batchgen.config.config import EngineConfig, ModelConfig from batchgen.config.batchgen_model_config import BatchGenModelConfig -from .configuration_glm5 import Glm5Config +from .config import assert_indexer_schedule_consistent from .set_basic_config import set_basic_config from .planner import GLM5Planner from batchgen.kv_cache.host_kv_mananger_config import build_host_kv_config @@ -37,15 +37,24 @@ class GLM5Initializer: def __init__(self, input_arguments): - self.loaded_model_config = Glm5Config() - self.loaded_model_config._name_or_path = input_arguments.huggingface_ckpt_name - self.loaded_model_config.architectures = ["GlmMoeDsaForCausalLM"] - self.model_name = input_arguments.huggingface_ckpt_name # Local checkpoint dir holding config.json (pre-downloaded model files). # None for a bare HF id -> resolver falls back to rich defaults. self.checkpoint_path = input_arguments.get("cache_dir", None) + # Single source of truth: the resolved internal config (GLM5Config / + # GLM52Config, checkpoint-backed). This IS the config used to build the + # model graph (loaded_model_config) — GLM-5 no longer uses an HF + # transformers.PretrainedConfig, matching kimi and every other model. + self.loaded_model_config = BatchGenModelConfig.resolve( + self.model_name, self.checkpoint_path + ) + self.loaded_model_config._name_or_path = self.model_name + self.loaded_model_config.architectures = ["GlmMoeDsaForCausalLM"] + # Fail loud at build if the DSA indexer schedule (freq/offset) disagrees + # with the checkpoint's per-layer indexer_types (no-op for GLM-5). + assert_indexer_schedule_consistent(self.loaded_model_config) + self.host_kv_cache_size = input_arguments.host_kv_cache_size self.host_kv_cache_byte_size = input_arguments.host_kv_cache_size * (1024**3) self.global_kv_cache_size_gb = input_arguments.global_host_kv_cache_size_gb @@ -178,7 +187,7 @@ def _parse_model_config(self): _default_engine_config, glm5_parameter_server). Set them as attributes so Init does not crash with AttributeError. """ - rich = BatchGenModelConfig.resolve(self.model_name, self.checkpoint_path) + rich = self.loaded_model_config model_config = ModelConfig() model_config.model_type = rich.model_type diff --git a/batchgen/models/glm/glm5/glm5_parameter_server.py b/batchgen/models/glm/glm5/glm5_parameter_server.py index 2a2aac26d..96b73581c 100644 --- a/batchgen/models/glm/glm5/glm5_parameter_server.py +++ b/batchgen/models/glm/glm5/glm5_parameter_server.py @@ -38,11 +38,9 @@ def _diag(msg): from tqdm import tqdm, trange _diag("tqdm done") -from .configuration_glm5 import Glm5Config -_diag("configuration_glm5 done") from .model import Glm5ForCausalLM _diag("model (Glm5ForCausalLM) done") -from batchgen.config.model_registry import load_config +from batchgen.config.batchgen_model_config import BatchGenModelConfig _diag("model_registry done") try: @@ -64,8 +62,11 @@ def __init__(self, huggingface_ckpt_name, cache_dir, converted_ckpt_dir, enable_ self.state_dict_name_map = {} self.enable_hugetlbfs = enable_hugetlbfs self.enable_memfd = enable_memfd - self.model_config = load_config(huggingface_ckpt_name) - self.hf_config = Glm5Config() + # Single resolved internal config (checkpoint-backed GLM5Config/GLM52Config). + # Used both for metadata reads and to build the model graph — GLM-5 no + # longer uses an HF transformers.PretrainedConfig. + self.model_config = BatchGenModelConfig.resolve(huggingface_ckpt_name, cache_dir) + self.hf_config = self.model_config self.hf_config._name_or_path = huggingface_ckpt_name free_memory, total_memory = torch.cuda.mem_get_info() diff --git a/batchgen/models/glm/glm5/model.py b/batchgen/models/glm/glm5/model.py index 979e6ac8b..bfd099e93 100644 --- a/batchgen/models/glm/glm5/model.py +++ b/batchgen/models/glm/glm5/model.py @@ -29,7 +29,11 @@ import torch.nn.functional as F from .decode_utils import clamp_token_indices_to_seqlens -from .configuration_glm5 import Glm5Config +# GLM-5 uses BatchGen's internal config (a plain BaseModelConfig dataclass), not +# an HF transformers.PretrainedConfig — matching every other model (kimi, etc.). +# Imported under the historical name `Glm5Config` so the __init__ type hints below +# need no churn; only attribute reads are used, so any config object works. +from .config import GLM5Config as Glm5Config, dsa_layer_skips_topk # ============================================================================ From ed8b545a270331d1714277edcb1a17cf71f1faf3 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Thu, 13 Aug 2026 00:48:12 +0800 Subject: [PATCH 4/4] feat(glm5): make eager decode None-safe for shared DSA layers GLM-5.2 runs its DSA indexer on the full layers only; shared layers carry no indexer weights and reuse the previous full layer's top-k. Prepare the model side of that reuse for the eager (non-CUDA-graph) decode path. - model.py Glm5MLA.__init__: set skip_topk/next_skip_topk; build Glm5Indexer only on full DSA layers, else self.indexer = None (use_dense_mla routing preserved). Glm5Model.__init__ guards the indexer.rotary_emb assignment for None (shared) layers. - wrappers.py: add ClassVar _dsa_prev_topk_indices (the carried top-k) and guard _maybe_init_fused_kernels and the prefill indexer-KV compute on indexer is None. - Parallel_Strategy_Manager.py: guard the indexer FP8 scale attach on None. GLM-5 has no such schedule, so every new branch is dead for it and its behavior is unchanged. --- .../glm/glm5/Parallel_Strategy_Manager.py | 5 ++-- batchgen/models/glm/glm5/model.py | 18 +++++++++---- batchgen/models/glm/glm5/wrappers.py | 27 ++++++++++++------- 3 files changed, 34 insertions(+), 16 deletions(-) diff --git a/batchgen/models/glm/glm5/Parallel_Strategy_Manager.py b/batchgen/models/glm/glm5/Parallel_Strategy_Manager.py index 260f11459..cab1b820b 100644 --- a/batchgen/models/glm/glm5/Parallel_Strategy_Manager.py +++ b/batchgen/models/glm/glm5/Parallel_Strategy_Manager.py @@ -688,8 +688,9 @@ def _setup_fp8_scales(self): # After wrapping, self_attn is GLM5AttnWrapper; original Glm5MLA is at .module inner = attn.module if hasattr(attn, 'module') else attn # When use_dense_mla is set, Glm5MLA skips indexer construction; - # skip the scale attach too (no destination). - if hasattr(inner, "indexer"): + # skip the scale attach too (no destination). GLM-5.2 "shared" layers + # carry no indexer weights either (indexer is None) — skip them. + if getattr(inner, "indexer", None) is not None: indexer = inner.indexer for proj, attr in [("wk", "wk_scale"), ("wq_b", "wq_b_scale")]: key = f"model.layers.{layer_idx}.self_attn.indexer.{proj}.weight_scale_inv" diff --git a/batchgen/models/glm/glm5/model.py b/batchgen/models/glm/glm5/model.py index bfd099e93..fc337116f 100644 --- a/batchgen/models/glm/glm5/model.py +++ b/batchgen/models/glm/glm5/model.py @@ -758,11 +758,17 @@ def __init__(self, config: Glm5Config, layer_idx: int): # Softmax scale self.softmax_scale = self.q_head_dim ** -0.5 - # DSA indexer — structurally absent when config.use_dense_mla is True. - # Downstream decode dispatch uses hasattr(self.module, 'indexer') as - # the signal to route to the dense-MLA (DeepSeek-V3 / Kimi style) path. + # DSA indexer — structurally absent when config.use_dense_mla is True + # (the `indexer` attribute is then NOT set, so hasattr(...) routes to the + # dense-MLA path — unchanged for DeepSeek-V3 / Kimi style configs). + # For DSA configs the attribute always exists, but GLM-5.2 "shared" layers + # carry no indexer weights: they reuse the previous full layer's top-k + # indices, so self.indexer is None on those layers (GLM-5: never a shared + # layer, so indexer is built on every layer — bit-identical). + self.skip_topk = dsa_layer_skips_topk(config, layer_idx) + self.next_skip_topk = dsa_layer_skips_topk(config, layer_idx + 1) if not getattr(config, "use_dense_mla", False): - self.indexer = Glm5Indexer(config, layer_idx) + self.indexer = None if self.skip_topk else Glm5Indexer(config, layer_idx) # Absorbed projections for decode (set by initialize()) self.q_absorb = None @@ -2302,7 +2308,9 @@ def __init__(self, config: Glm5Config): # Assign shared RoPE to attention and indexer for layer in self.layers: layer.self_attn.rotary_emb = self._shared_rotary_emb - if hasattr(layer.self_attn, 'indexer'): + # DSA layers have an `indexer` attribute; GLM-5.2 shared layers set it + # to None (no indexer weights), so guard the deref. + if getattr(layer.self_attn, 'indexer', None) is not None: layer.self_attn.indexer.rotary_emb = self._shared_rotary_emb self.norm = Glm5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/batchgen/models/glm/glm5/wrappers.py b/batchgen/models/glm/glm5/wrappers.py index a9efafc51..d51c3422d 100644 --- a/batchgen/models/glm/glm5/wrappers.py +++ b/batchgen/models/glm/glm5/wrappers.py @@ -413,6 +413,10 @@ class GLM5AttnWrapper(AttnWrapperBase): # Set once per decode step by the worker so per-layer _forward_decode_dsa # branches on it without doing a D2H .sum().item() 78 times per step. _dsa_short_count: ClassVar[Optional[int]] = None + # DSA top-k index REUSE (GLM-5.2): carried top-k indices from the most recent + # FULL layer, reused by subsequent shared layers. Reset per decode step by the + # worker. For GLM-5 (all layers full) this is never read by a shared branch. + _dsa_prev_topk_indices: ClassVar[Optional[torch.Tensor]] = None # Whole-model CUDA graph can pad local rows to a global NCCL bucket. These # graph-owned overrides let GLM-5 DSA use explicit slot sentinels for padded # rows instead of deriving slot count from cur_batch. @@ -547,7 +551,7 @@ def initialize_fused_kernels(self): AND after _setup_fp8_scales has attached indexer.wk_scale / wq_b_scale. """ attn = self.module - if not hasattr(attn, "indexer"): + if not hasattr(attn, "indexer") or attn.indexer is None: return indexer = attn.indexer @@ -700,15 +704,20 @@ def _forward_prefill(self, hidden_states: torch.Tensor, **kwargs) -> Tuple: raise RuntimeError( "GLM-5 DSA prefill requires indexer KV; refusing primary-only host offload" ) - indexer_kv = self.module.indexer.compute_indexer_kv( - hidden_states_2d.unsqueeze(0), - positions=self.position_ids.to(hidden_states_2d.device), - ) - if indexer_kv is None: - raise RuntimeError( - "GLM-5 DSA prefill indexer returned no KV; refusing primary-only host offload" + # Shared layers (GLM-5.2) have indexer is None: they reuse a full + # layer's top-k at decode and never read an aux indexer cache, so + # skip the indexer-K compute + offload entirely. GLM-5 layers always + # have a real indexer so this block always runs there. + if self.module.indexer is not None: + indexer_kv = self.module.indexer.compute_indexer_kv( + hidden_states_2d.unsqueeze(0), + positions=self.position_ids.to(hidden_states_2d.device), ) - self._offload_prepacked_indexer_kv(indexer_kv.squeeze(0)) + if indexer_kv is None: + raise RuntimeError( + "GLM-5 DSA prefill indexer returned no KV; refusing primary-only host offload" + ) + self._offload_prepacked_indexer_kv(indexer_kv.squeeze(0)) self._offload_prepacked_kv(offload_kv) attn_output = attn_output.unsqueeze(0)