From b4586916367d926d6aa65a7309ed18feb7d4f77d Mon Sep 17 00:00:00 2001 From: TairanXU Date: Thu, 13 Aug 2026 00:41:47 +0800 Subject: [PATCH 1/2] feat(config): add BatchGenModelConfig resolver for GLM-5.2 Introduce a decoupled, pure-Python config-resolution module as the single producer of the engine's model config, replacing hand-hardcoded values in the GLM-5 initializer that silently drifted from the checkpoint. Enables GLM-5.2 config resolution (its own config identity, shared model graph). - batchgen/config/batchgen_model_config.py (new): resolve(model_name, ckpt) returns a rich BaseModelConfig subclass via a supported-variant registry plus per-model from_hf() mapping; validate() fails loud; unlisted variants warn loudly. Module body is torch/engine-free and imports model configs lazily. - batchgen/config/model_registry.py: add GLM-5.2 patterns before GLM-5 and a loud warning for unlisted GLM-5.x variants. - batchgen/config/tokenizer_registry.py: add the GLM-5.2 tokenizer pattern before GLM-5. GLM-only in this change; other model families keep the existing load_config path unchanged and will be migrated one per change. --- batchgen/config/batchgen_model_config.py | 225 +++++++++++++++++++++++ batchgen/config/model_registry.py | 60 +++++- batchgen/config/tokenizer_registry.py | 10 +- 3 files changed, 292 insertions(+), 3 deletions(-) create mode 100644 batchgen/config/batchgen_model_config.py diff --git a/batchgen/config/batchgen_model_config.py b/batchgen/config/batchgen_model_config.py new file mode 100644 index 000000000..0fa7fca80 --- /dev/null +++ b/batchgen/config/batchgen_model_config.py @@ -0,0 +1,225 @@ +# ---------------------------------------------------------------------------- # +# 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. # +# ---------------------------------------------------------------------------- # + +"""Single, decoupled model-config resolver for BatchGen. + +Historically GLM config was resolved in three inconsistent places and the +engine ran on values hardcoded in ``glm5_initializer._parse_model_config`` that +silently drifted from the checkpoint. This module is the ONE producer of a +rich, checkpoint-backed :class:`BaseModelConfig` subclass; the GLM initializer +then *projects* that rich config into the minimal engine ``ModelConfig``. + +Design constraints (do not break): + +* This module's OWN body is PURE PYTHON — no ``torch`` / engine / distributed + imports execute at module load. Family config classes are imported *lazily* + inside :meth:`BatchGenModelConfig.resolve` (they live under ``batchgen.models`` + and only pull heavier deps when actually needed). + + Caveat on import path: the enclosing ``batchgen.config`` package ``__init__`` + eagerly imports the tokenizer stack (torch) and the model registry (whose + auto-import triggers an engine JIT build), so a plain + ``import batchgen.config.batchgen_model_config`` is NOT torch-free — that cost + comes from the package ``__init__``, not this module. In the real engine the + GLM initializer already lives behind torch, so this is a non-issue there. + A genuinely headless caller (e.g. the unit tests) must exec this file directly + via ``importlib`` — see ``tests/test_batchgen_model_config.py`` — rather than + going through the package import. +* Supported variants are matched by an ordered name->class registry. Insertion + order == iteration order, so more-specific patterns (``GLM-5.2``) MUST precede + broader ones (``GLM-5``) — otherwise ``GLM-5.2-FP8`` would substring-match + ``GLM-5`` first. +* Only GLM routes through this resolver today. Other families keep their + existing ``load_config`` behaviour untouched. +""" + +from __future__ import annotations + +import json +import logging +import re +from importlib import import_module +from pathlib import Path +from typing import Dict, List, Optional, Tuple, TYPE_CHECKING + +if TYPE_CHECKING: # pragma: no cover - typing only, never imported at runtime + from batchgen.config.model_config import BaseModelConfig + +logger = logging.getLogger(__name__) + + +# Ordered registry: substring pattern -> (module path, class name) of the rich +# family config. ORDER MATTERS — dict insertion order is iteration order, and a +# more-specific pattern must come before a broader one it would substring-match. +# +# GLM-5 / GLM-5.1 share the architecturally-identical glm_moe_dsa graph and use +# GLM5Config. GLM-5.2 gets its OWN config identity (GLM52Config, model_type +# "glm_moe_dsa_5_2"); the model *code* stays shared, only the config differs. +_SUPPORTED_VARIANTS: Dict[str, Tuple[str, str]] = { + "GLM-5.2-FP8": ("batchgen.models.glm.glm5.config", "GLM52Config"), + "GLM-5.2": ("batchgen.models.glm.glm5.config", "GLM52Config"), + "GLM-5.1-FP8": ("batchgen.models.glm.glm5.config", "GLM5Config"), + "GLM-5.1": ("batchgen.models.glm.glm5.config", "GLM5Config"), + "GLM-5-FP8": ("batchgen.models.glm.glm5.config", "GLM5Config"), + "GLM-5": ("batchgen.models.glm.glm5.config", "GLM5Config"), +} + +# Best-effort fallback for an unlisted GLM variant (see resolve()). +_FALLBACK_VARIANT: Tuple[str, str] = ("batchgen.models.glm.glm5.config", "GLM5Config") + + +class BatchGenModelConfig: + """Namespace for the config-resolution entry point. + + This is intentionally a thin static facade rather than a dataclass: it holds + no state and exists only so callers get a single, discoverable API + (``BatchGenModelConfig.resolve(...)``) that returns the rich family config. + """ + + @staticmethod + def _match_variant(model_name: str) -> Optional[Tuple[str, str]]: + """Return the (module, class) for the first matching name pattern. + + When ``model_name`` matches only the broad ``GLM-5`` / ``GLM-5-FP8`` + pattern but carries an unlisted minor/patch token (e.g. ``GLM-5.3``, + ``GLM-50``), emit a loud warning before binding it to GLM5Config — a + silently mis-resolved future variant would build the engine with wrong + dims (mirrors model_registry._warn_if_unlisted_glm5_variant). + """ + for pattern, target in _SUPPORTED_VARIANTS.items(): + if pattern in model_name: + if pattern in ("GLM-5", "GLM-5-FP8"): + BatchGenModelConfig._warn_unlisted_glm5(model_name) + return target + return None + + @staticmethod + def _warn_unlisted_glm5(model_name: str) -> None: + """Loudly warn for a GLM-5.x identifier caught by the broad GLM-5 rule.""" + listed_minors = set() + for pat in _SUPPORTED_VARIANTS: + m = re.search(r"GLM-5(?:\.(\d+))?", pat) + if m and pat.startswith("GLM-5"): + listed_minors.add(m.group(1)) # None for bare GLM-5 + m = re.search(r"GLM-5(\d*)(?:\.(\d+))?", model_name) + if m is None: + return + glued, minor = m.group(1), m.group(2) + if glued: + logger.warning( + "BatchGenModelConfig.resolve: model_name=%r matched the broad " + "'GLM-5' pattern as a superstring (GLM-5%s...). This is almost " + "certainly NOT GLM-5; resolving to GLM5Config anyway. Add an " + "explicit variant if this is real.", + model_name, glued, + ) + elif minor not in listed_minors: + logger.warning( + "BatchGenModelConfig.resolve: model_name=%r is an unlisted " + "GLM-5.%s variant; it matched the broad 'GLM-5' pattern and is " + "resolving to GLM5Config (GLM-5 base). If its config.json " + "diverges the engine will be built with wrong dims. Add an " + "explicit 'GLM-5.%s' entry to _SUPPORTED_VARIANTS.", + model_name, minor, minor, + ) + + @staticmethod + def _read_hf_config(checkpoint_path: Optional[str]) -> Optional[Dict]: + """Load ``config.json`` from a checkpoint dir (or a direct json path). + + Returns None when no checkpoint path is given or no config.json exists + (e.g. a bare HuggingFace model id) — the caller then falls back to the + family config's built-in defaults. + """ + if not checkpoint_path: + return None + p = Path(checkpoint_path) + if p.is_file() and p.suffix == ".json": + config_file = p + else: + config_file = p / "config.json" + if not config_file.exists(): + return None + with open(config_file, "r") as f: + return json.load(f) + + @staticmethod + def resolve( + model_name: str, + checkpoint_path: Optional[str] = None, + ) -> "BaseModelConfig": + """Resolve a rich model config for ``model_name``. + + Steps: + 1. Match ``model_name`` against the supported-variant registry to pick + the family config class. An unlisted variant proceeds best-effort + (GLM5Config) but logs a loud warning. + 2. Read ``/config.json`` when available. + 3. Build the rich config via the class's ``from_hf`` classmethod when a + checkpoint config is present; otherwise use the class defaults. + 4. ``validate()`` — FAIL LOUD on missing required fields or + self-inconsistency. + 5. Return the rich subclass instance. + + Args: + model_name: Model identifier / checkpoint name used for pattern + matching (e.g. "zai-org/GLM-5.2-FP8", "GLM-5-FP8"). + checkpoint_path: Local path to the checkpoint dir (or its + config.json). None / non-local ids fall back to defaults. + + Returns: + A rich :class:`BaseModelConfig` subclass instance. + """ + target = BatchGenModelConfig._match_variant(model_name) + if target is None: + logger.warning( + "BatchGenModelConfig.resolve: model_name=%r matched no supported " + "variant %s. Proceeding best-effort with %s defaults; verify the " + "resolved config is correct for this checkpoint.", + model_name, + list(_SUPPORTED_VARIANTS.keys()), + _FALLBACK_VARIANT[1], + ) + target = _FALLBACK_VARIANT + + module_path, class_name = target + # Lazy import: keeps this module free of torch / engine deps at load. + config_module = import_module(module_path) + config_cls = getattr(config_module, class_name) + + hf_dict = BatchGenModelConfig._read_hf_config(checkpoint_path) + if hf_dict is not None: + logger.info( + "Resolving %s from checkpoint config for model_name=%r", + class_name, + model_name, + ) + config = config_cls.from_hf(hf_dict) + else: + logger.warning( + "BatchGenModelConfig.resolve: no config.json found for " + "model_name=%r (checkpoint_path=%r). Using %s built-in defaults.", + model_name, + checkpoint_path, + class_name, + ) + config = config_cls() + + config._name_or_path = model_name + config.validate() + return config diff --git a/batchgen/config/model_registry.py b/batchgen/config/model_registry.py index 0554fe00d..f8c817f21 100644 --- a/batchgen/config/model_registry.py +++ b/batchgen/config/model_registry.py @@ -43,6 +43,7 @@ import json import logging import os +import re from .model_name_utils import KIMI_K25_BACKEND_MODEL_IDS @@ -83,7 +84,12 @@ "gpt-oss": "gpt_oss", # GLM-5 / GLM-5.1 share an architecturally-identical glm_moe_dsa graph # (754B MoE + DSA, 78 layers, identical config.json apart from transformers_version). - # More-specific patterns first so `GLM-5.1-FP8` doesn't get swallowed by `GLM-5`. + # GLM-5.2 shares the model graph but gets its OWN config identity + # (glm_moe_dsa_5_2) — longer context, nested rope, extra DSA indexer knobs. + # More-specific patterns first so `GLM-5.2-FP8` doesn't get swallowed by + # `GLM-5` (and `GLM-5.1-FP8` not by `GLM-5`). + "GLM-5.2-FP8": "glm_moe_dsa_5_2", + "GLM-5.2": "glm_moe_dsa_5_2", "GLM-5.1-FP8": "glm_moe_dsa", "GLM-5.1": "glm_moe_dsa", "GLM-5-FP8": "glm_moe_dsa", @@ -116,6 +122,55 @@ def decorator(cls: Type["BaseModelConfig"]) -> Type["BaseModelConfig"]: return decorator +def _warn_if_unlisted_glm5_variant(model_identifier: str, matched_pattern: str) -> None: + """Loudly warn when a GLM-5.x identifier resolved via the broad ``GLM-5`` + catch-all rather than an explicitly-listed variant. + + ``"GLM-5"`` is a substring of every future minor (``GLM-5.3``, ``GLM-5.9``) + and of superstrings like ``GLM-50``. Without this guard such an identifier + silently binds to the GLM-5 base config (glm_moe_dsa) — and if its real + config.json diverges the engine would be built with wrong dims and no + diagnostic. We only warn when the matched pattern is the broad ``GLM-5`` / + ``GLM-5-FP8`` AND the identifier carries a minor/patch token that is not one + of the explicitly-listed GLM-5.x patterns. + """ + if matched_pattern not in ("GLM-5", "GLM-5-FP8"): + return + + # Version tokens that are explicitly listed (and therefore trusted). + listed_versions = set() + for pat in MODEL_NAME_PATTERNS: + m = re.search(r"GLM-5(?:\.(\d+))?", pat) + if m and pat.startswith("GLM-5"): + listed_versions.add(m.group(1)) # None for bare "GLM-5", "1", "2", ... + + # Extract the GLM-5 version token from the identifier: an optional trailing + # integer glued to the 5 (GLM-50) or a dotted minor (GLM-5.3). + m = re.search(r"GLM-5(\d*)(?:\.(\d+))?", model_identifier) + if m is None: + return + glued, minor = m.group(1), m.group(2) + # GLM-50 / GLM-51 ... — a superstring of "GLM-5", never a real GLM-5 minor. + if glued: + logger.warning( + "Model identifier %r matched the broad 'GLM-5' pattern as a " + "superstring (GLM-5%s...). This is almost certainly NOT GLM-5; it " + "is being resolved to the GLM-5 base config (glm_moe_dsa). Add an " + "explicit pattern if this is a real variant.", + model_identifier, glued, + ) + return + if minor not in listed_versions: + logger.warning( + "Model identifier %r is an unlisted GLM-5.%s variant; it matched the " + "broad 'GLM-5' pattern and is being resolved to the GLM-5 base config " + "(glm_moe_dsa). If its config.json diverges from GLM-5 the engine will " + "be built with wrong dims. Add an explicit 'GLM-5.%s' pattern to " + "MODEL_NAME_PATTERNS.", + model_identifier, minor, minor, + ) + + def _detect_model_type_from_identifier(model_identifier: str) -> Optional[str]: """Detect model type from HuggingFace model identifier. @@ -130,6 +185,9 @@ def _detect_model_type_from_identifier(model_identifier: str) -> Optional[str]: for pattern, model_type in MODEL_NAME_PATTERNS.items(): if pattern in model_identifier: logger.debug(f"Detected model_type={model_type} from identifier={model_identifier}") + # Loud guard: an unlisted GLM-5.x that fell through to the broad + # 'GLM-5' catch-all must not silently bind to the base config. + _warn_if_unlisted_glm5_variant(model_identifier, pattern) return model_type return None diff --git a/batchgen/config/tokenizer_registry.py b/batchgen/config/tokenizer_registry.py index 719d7f8c4..afc52e243 100644 --- a/batchgen/config/tokenizer_registry.py +++ b/batchgen/config/tokenizer_registry.py @@ -71,8 +71,14 @@ # EOS/pad), but GLM-5.1 ships a richer chat template (tool_to_json macro, # thinking_indices tracking, tool_reference responses). We route them to # separate tokenizer types so each loads its own Jinja template. - # More-specific patterns first so `GLM-5.1-FP8` doesn't get swallowed by - # `GLM-5`. + # GLM-5.2 reserves its own tokenizer identity (glm_moe_dsa_5_2) here; it + # shares the GLM-5 vocab, so until a dedicated GLM52Tokenizer class is + # registered, load_tokenizer() gracefully falls through this (unregistered) + # type to the shared "glm_moe_dsa" tokenizer via the `GLM-5` pattern below. + # More-specific patterns first so `GLM-5.2-FP8` / `GLM-5.1-FP8` don't get + # swallowed by `GLM-5`. + "GLM-5.2-FP8": "glm_moe_dsa_5_2", + "GLM-5.2": "glm_moe_dsa_5_2", "GLM-5.1-FP8": "glm_moe_dsa_5_1", "GLM-5.1": "glm_moe_dsa_5_1", "GLM-5-FP8": "glm_moe_dsa", From f1f7efb65bdafa0e1934dadb98f23d421a2babdb Mon Sep 17 00:00:00 2001 From: TairanXU Date: Thu, 13 Aug 2026 00:41:47 +0800 Subject: [PATCH 2/2] feat(glm5): reset DSA indexer top-k carry each decode step Clear GLM5AttnWrapper._dsa_prev_topk_indices once per decode step (before layer 0) in both decode dispatch paths, so GLM-5.2 shared DSA layers never reuse a stale top-k carried over from the previous step. GLM-5 never publishes a carry, so this is a no-op for it. --- batchgen/batchgen_worker.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 35427a22b..92cf18bb5 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -8163,6 +8163,9 @@ def _bind_decode_attention_metadata_for_graph_config( else: GLM5AttnWrapper._dsa_short_count = 0 if cache_view.numel() == 0 else None + # GLM-5.2 DSA indexer reuse: clear prev top-k once per decode step (before layer 0) + # so shared layers never reuse a stale value carried over from the previous step. + GLM5AttnWrapper._dsa_prev_topk_indices = None gpu_manager = self._get_cuda_graph_gpu_manager() if gpu_manager is not None and cur_batch: manager = getattr(gpu_manager, "primary", gpu_manager) @@ -9811,6 +9814,12 @@ def decoding_continuous( else: GLM5AttnWrapper._dsa_short_count = None + # GLM-5.2 DSA indexer reuse: clear prev top-k once per decode + # step (before layer 0) so shared layers never reuse a stale + # value from the previous step. (Second decode path; the + # graph-config path resets it separately.) + GLM5AttnWrapper._dsa_prev_topk_indices = None + if new_tokens.shape[0] != len(batch): new_tokens = self._rebuild_input_tokens(batch) else: