From 2279b79f35b9b938317f02ec344c2f3b4fb10382 Mon Sep 17 00:00:00 2001 From: Xingyu Liu <38244988+charlotte12l@users.noreply.github.com> Date: Thu, 14 May 2026 01:54:06 -0700 Subject: [PATCH] Add --model-config-parser registry for pluggable config formats (#25050) Signed-off-by: Xingyu Liu --- python/sglang/srt/configs/model_config.py | 4 + .../configs/model_config_parser_registry.py | 62 +++++ python/sglang/srt/server_args.py | 71 +++-- .../srt/utils/hf_transformers/config.py | 248 +++++++++++------- .../test_model_config_parser_registry.py | 62 +++++ 5 files changed, 326 insertions(+), 121 deletions(-) create mode 100644 python/sglang/srt/configs/model_config_parser_registry.py create mode 100644 test/registered/unit/configs/test_model_config_parser_registry.py diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 7ba9421da..5a9a93101 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -174,6 +174,7 @@ class ModelConfig: encoder_only: bool = False, language_only: bool = False, disable_hybrid_swa_memory: bool = False, + model_config_parser: str = "auto", ) -> None: # Parse args self.model_path = model_path @@ -185,6 +186,7 @@ class ModelConfig: self.quantize_and_serve = quantize_and_serve self.is_multi_layer_eagle = is_multi_layer_eagle self.disable_hybrid_swa_memory = disable_hybrid_swa_memory + self.model_config_parser = model_config_parser # Validate quantize_and_serve configuration self._validate_quantize_and_serve_config() @@ -201,6 +203,7 @@ class ModelConfig: trust_remote_code=trust_remote_code, revision=revision, model_override_args=self.model_override_args, + model_config_parser=model_config_parser, **kwargs, ) self.hf_text_config = get_hf_text_config(self.hf_config) @@ -403,6 +406,7 @@ class ModelConfig: encoder_only=server_args.encoder_only, is_draft_model=is_draft_model, disable_hybrid_swa_memory=server_args.disable_hybrid_swa_memory, + model_config_parser=server_args.model_config_parser, **kwargs, ) diff --git a/python/sglang/srt/configs/model_config_parser_registry.py b/python/sglang/srt/configs/model_config_parser_registry.py new file mode 100644 index 000000000..27a15f714 --- /dev/null +++ b/python/sglang/srt/configs/model_config_parser_registry.py @@ -0,0 +1,62 @@ +"""Named registry for model-config parsers. + +Mirrors the ``LoadFormat.PRIVATE`` escape hatch in +:mod:`sglang.srt.configs.load_config` but registry-shaped, so multiple +plugins can coexist without colliding on a single private import path. +""" + +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Optional + +from transformers import PretrainedConfig + +logger = logging.getLogger(__name__) + + +class ModelConfigParserBase(ABC): + @abstractmethod + def parse( + self, + model: str | Path, + trust_remote_code: bool, + revision: Optional[str] = None, + **kwargs, + ) -> PretrainedConfig: + raise NotImplementedError + + +_MODEL_CONFIG_PARSER_REGISTRY: dict[str, type[ModelConfigParserBase]] = {} + + +def register_model_config_parser(name: str): + """Returned instances are freshly constructed on each call -- parsers + should be stateless or carry only per-instance state.""" + + def _wrapper(cls): + if not issubclass(cls, ModelConfigParserBase): + raise ValueError("Model-config parser must subclass ModelConfigParserBase.") + if name in _MODEL_CONFIG_PARSER_REGISTRY: + logger.warning( + "Model-config parser %r already registered; overwriting with %s", + name, + cls, + ) + _MODEL_CONFIG_PARSER_REGISTRY[name] = cls + logger.debug("Registered model-config parser %r -> %s", name, cls.__name__) + return cls + + return _wrapper + + +def get_model_config_parser(name: str) -> ModelConfigParserBase: + """``"auto"`` is not handled here -- the caller must resolve it first.""" + if name not in _MODEL_CONFIG_PARSER_REGISTRY: + raise ValueError( + f"Unknown model-config parser {name!r}. " + f"Registered: {sorted(_MODEL_CONFIG_PARSER_REGISTRY)}" + ) + return _MODEL_CONFIG_PARSER_REGISTRY[name]() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d0e1b1b91..66aff8edf 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -17,6 +17,7 @@ from __future__ import annotations import argparse import dataclasses +import glob import importlib import importlib.util import json @@ -378,6 +379,7 @@ class ServerArgs: enable_multimodal: Optional[bool] = None revision: Optional[str] = None model_impl: str = "auto" + model_config_parser: str = "auto" # HTTP server host: str = "127.0.0.1" @@ -3953,44 +3955,60 @@ class ServerArgs: ) def _is_mistral_native_format(self) -> bool: - """Detect if the model uses Mistral native format (params.json + consolidated weights). + """True iff the checkpoint requires load_format=mistral. - When both params.json and config.json exist, default to HF format to - avoid weight-name mismatches (e.g. Mistral-7B-Instruct-v0.3). + Looks for ``consolidated*.safetensors`` with no competing + ``model-*.safetensors``; when both weight formats ship in the + same checkpoint (e.g. Mistral-7B-Instruct-v0.3) the HF path is + preferred to avoid loading Mistral-named weights into an + HF-named architecture. - Exception: models routed through ``_load_mistral_large_3_for_causal_LM`` - (mistral-large-3, mistral-small-4, leanstral) build their config from - params.json and expect native weight names, so native format is required - even when config.json is also present. + Name override: ``mistral-large-3`` / ``mistral-small-4`` / + ``leanstral`` always treat as Mistral-native when ``params.json`` + is present -- those families need Mistral weight loading + regardless of which weight files happen to be present. """ - # Keep in sync with the name checks in - # hf_transformers_utils.py::get_config / get_tokenizer. - _MISTRAL_NATIVE_CONFIG_PATTERNS = ( + _MISTRAL_NATIVE_PATTERNS = ( "mistral-large-3", "mistral-small-4", "leanstral", ) + name_matches = any( + p in str(self.model_path).lower() for p in _MISTRAL_NATIVE_PATTERNS + ) - def _check_format(has_params: bool, has_hf_config: bool) -> bool: - if has_params and not has_hf_config: + def _check_format(has_params, has_consolidated, has_hf_weights) -> bool: + if has_params and name_matches: return True - if has_params and has_hf_config: - model_lower = str(self.model_path).lower() - if any(name in model_lower for name in _MISTRAL_NATIVE_CONFIG_PATTERNS): - return True - return False + return has_consolidated and not has_hf_weights if os.path.isdir(self.model_path): - has_params = os.path.exists(os.path.join(self.model_path, "params.json")) - has_hf_config = os.path.exists(os.path.join(self.model_path, "config.json")) - return _check_format(has_params, has_hf_config) + return _check_format( + has_params=os.path.exists(os.path.join(self.model_path, "params.json")), + has_consolidated=bool( + glob.glob( + os.path.join(self.model_path, "consolidated*.safetensors") + ) + ), + has_hf_weights=bool( + glob.glob(os.path.join(self.model_path, "model-*.safetensors")) + ), + ) - # For hub models, check remote files try: from huggingface_hub import HfApi files = {s.rfilename for s in HfApi().model_info(self.model_path).siblings} - return _check_format("params.json" in files, "config.json" in files) + return _check_format( + has_params="params.json" in files, + has_consolidated=any( + f.startswith("consolidated") and f.endswith(".safetensors") + for f in files + ), + has_hf_weights=any( + f.startswith("model-") and f.endswith(".safetensors") for f in files + ), + ) except Exception: return False @@ -4585,6 +4603,15 @@ class ServerArgs: '* "mindspore" will use the MindSpore model ' "implementation.\n", ) + parser.add_argument( + "--model-config-parser", + type=str, + default=ServerArgs.model_config_parser, + help='Which model-config parser to use. "auto" picks "mistral" ' + 'via the is_mistral_model name heuristic, else "hf" ' + "(AutoConfig over config.json). Plugins can register additional " + "parsers via @register_model_config_parser.", + ) # HTTP server parser.add_argument( diff --git a/python/sglang/srt/utils/hf_transformers/config.py b/python/sglang/srt/utils/hf_transformers/config.py index 6e1743e1c..a46c4a45d 100644 --- a/python/sglang/srt/utils/hf_transformers/config.py +++ b/python/sglang/srt/utils/hf_transformers/config.py @@ -18,6 +18,11 @@ from typing import Optional from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES +from sglang.srt.configs.model_config_parser_registry import ( + ModelConfigParserBase, + get_model_config_parser, + register_model_config_parser, +) from sglang.srt.connector import create_remote_connector from sglang.srt.utils import is_remote_url, lru_cache_frozenset @@ -46,19 +51,156 @@ def _apply_deepseek_ocr_overrides(config, model): config._name_or_path = model +@register_model_config_parser("hf") +class HfModelConfigParser(ModelConfigParserBase): + def parse( + self, + model, + trust_remote_code: bool, + revision: Optional[str] = None, + **kwargs, + ): + config = AutoConfig.from_pretrained( + model, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + + if ( + config.architectures is not None + and config.architectures[0] == "Phi4MMForCausalLM" + ): + from transformers import SiglipVisionConfig + + config.vision_config = SiglipVisionConfig( + hidden_size=1152, + image_size=448, + intermediate_size=4304, + model_type="siglip_vision_model", + num_attention_heads=16, + num_hidden_layers=26, + patch_size=14, + ) + + if config.architectures in [ + ["LongcatCausalLM"], + ["LongcatFlashForCausalLM"], + ["LongcatFlashNgramForCausalLM"], + ]: + config.model_type = "longcat_flash" + + text_config = get_hf_text_config(config=config) + + if isinstance(model, str) and text_config is not None: + items = ( + text_config.items() + if hasattr(text_config, "items") + else vars(text_config).items() + ) + for key, val in items: + if not hasattr(config, key) and val is not None: + setattr(config, key, val) + + is_ocr = _is_deepseek_ocr_model(config) + is_ocr2 = _is_deepseek_ocr2_model(config) + + if is_ocr2: + _override_v_head_dim_if_zero(config) + config.model_type = "deepseek-ocr" + _set_architectures(config, "DeepseekOCRForCausalLM") + config = DeepseekVLV2Config.from_pretrained(model, revision=revision) + _apply_deepseek_ocr_overrides(config, model) + elif config.model_type in _CONFIG_REGISTRY: + model_type = config.model_type + if model_type == "deepseek_vl_v2" and is_ocr: + model_type = "deepseek-ocr" + config = _CONFIG_REGISTRY[model_type].from_pretrained( + model, revision=revision + ) + + # Re-check after reloading config from registry + if _is_deepseek_ocr_model(config) or _is_deepseek_ocr2_model(config): + _apply_deepseek_ocr_overrides(config, model) + else: + config._name_or_path = model + + if isinstance(model, str) and config.model_type == "internvl_chat": + for key, val in config.llm_config.__dict__.items(): + if not hasattr(config, key): + setattr(config, key, val) + + if config.model_type == "multi_modality": + _set_architectures(config, "MultiModalityCausalLM") + + if config.model_type in ("gemma4", "gemma4_assistant"): + # Gemma4 configs use base attributes for SWA layers and `global_*` + # variants for full-attention layers. SGLang expects the opposite: + # base = full-attention, `swa_*` = sliding-window overrides. + text_config = config.text_config + global_head_dim = getattr(text_config, "global_head_dim", None) + global_kv_heads = getattr(text_config, "num_global_key_value_heads", None) + + swa_head_dim = text_config.head_dim + swa_kv_heads = text_config.num_key_value_heads + + text_config.swa_head_dim = swa_head_dim + text_config.swa_v_head_dim = swa_head_dim + text_config.swa_num_key_value_heads = swa_kv_heads + + if global_head_dim is not None: + text_config.head_dim = global_head_dim + if global_kv_heads is not None: + text_config.num_key_value_heads = global_kv_heads + + if not hasattr(text_config, "v_head_dim"): + text_config.v_head_dim = text_config.head_dim + if not hasattr(text_config, "swa_v_head_dim"): + text_config.swa_v_head_dim = text_config.swa_head_dim + + if config.model_type == "longcat_flash": + _set_architectures(config, "LongcatFlashForCausalLM") + + return config + + +@register_model_config_parser("mistral") +class MistralModelConfigParser(ModelConfigParserBase): + def parse( + self, + model, + trust_remote_code: bool, + revision: Optional[str] = None, + **kwargs, + ): + del kwargs + return load_mistral_config( + model, trust_remote_code=trust_remote_code, revision=revision + ) + + @lru_cache_frozenset(maxsize=32) def get_config( model: str, trust_remote_code: bool, revision: Optional[str] = None, model_override_args: Optional[dict] = None, + model_config_parser: str = "auto", **kwargs, ): is_gguf = check_gguf_file(model) if is_gguf: + if model_config_parser not in ("auto", "hf"): + raise ValueError( + f"model_config_parser={model_config_parser!r} is incompatible " + "with GGUF inputs; only 'hf' (or 'auto') is supported." + ) _ensure_gguf_version() kwargs["gguf_file"] = model model = Path(model).parent + # Skip auto-resolution for GGUF: the name-based Mistral heuristic + # would misfire on the rewritten parent dir. + model_config_parser = "hf" model = resolve_runai_obj_uri(model) @@ -67,106 +209,14 @@ def get_config( client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"]) model = client.get_local_dir() - if is_mistral_model(model): - config = load_mistral_config( - model, trust_remote_code=trust_remote_code, revision=revision - ) - else: - config = AutoConfig.from_pretrained( - model, trust_remote_code=trust_remote_code, revision=revision, **kwargs - ) + if model_config_parser == "auto": + # `model` is post-rewrite (gguf parent / runai uri / remote pull). + model_config_parser = "mistral" if is_mistral_model(model) else "hf" - if ( - config.architectures is not None - and config.architectures[0] == "Phi4MMForCausalLM" - ): - from transformers import SiglipVisionConfig - - config.vision_config = SiglipVisionConfig( - hidden_size=1152, - image_size=448, - intermediate_size=4304, - model_type="siglip_vision_model", - num_attention_heads=16, - num_hidden_layers=26, - patch_size=14, - ) - - if config.architectures in [ - ["LongcatCausalLM"], - ["LongcatFlashForCausalLM"], - ["LongcatFlashNgramForCausalLM"], - ]: - config.model_type = "longcat_flash" - - text_config = get_hf_text_config(config=config) - - if isinstance(model, str) and text_config is not None: - items = ( - text_config.items() - if hasattr(text_config, "items") - else vars(text_config).items() - ) - for key, val in items: - if not hasattr(config, key) and val is not None: - setattr(config, key, val) - - is_ocr = _is_deepseek_ocr_model(config) - is_ocr2 = _is_deepseek_ocr2_model(config) - - if is_ocr2: - _override_v_head_dim_if_zero(config) - config.model_type = "deepseek-ocr" - _set_architectures(config, "DeepseekOCRForCausalLM") - config = DeepseekVLV2Config.from_pretrained(model, revision=revision) - _apply_deepseek_ocr_overrides(config, model) - elif config.model_type in _CONFIG_REGISTRY: - model_type = config.model_type - if model_type == "deepseek_vl_v2" and is_ocr: - model_type = "deepseek-ocr" - config = _CONFIG_REGISTRY[model_type].from_pretrained(model, revision=revision) - - # Re-check after reloading config from registry - if _is_deepseek_ocr_model(config) or _is_deepseek_ocr2_model(config): - _apply_deepseek_ocr_overrides(config, model) - else: - config._name_or_path = model - - if isinstance(model, str) and config.model_type == "internvl_chat": - for key, val in config.llm_config.__dict__.items(): - if not hasattr(config, key): - setattr(config, key, val) - - if config.model_type == "multi_modality": - _set_architectures(config, "MultiModalityCausalLM") - - if config.model_type in ("gemma4", "gemma4_assistant"): - # Gemma4 configs use base attributes for SWA layers and `global_*` - # variants for full-attention layers. SGLang expects the opposite: - # base = full-attention, `swa_*` = sliding-window overrides. - text_config = config.text_config - global_head_dim = getattr(text_config, "global_head_dim", None) - global_kv_heads = getattr(text_config, "num_global_key_value_heads", None) - - swa_head_dim = text_config.head_dim - swa_kv_heads = text_config.num_key_value_heads - - text_config.swa_head_dim = swa_head_dim - text_config.swa_v_head_dim = swa_head_dim - text_config.swa_num_key_value_heads = swa_kv_heads - - if global_head_dim is not None: - text_config.head_dim = global_head_dim - if global_kv_heads is not None: - text_config.num_key_value_heads = global_kv_heads - - if not hasattr(text_config, "v_head_dim"): - text_config.v_head_dim = text_config.head_dim - if not hasattr(text_config, "swa_v_head_dim"): - text_config.swa_v_head_dim = text_config.swa_head_dim - - if config.model_type == "longcat_flash": - _set_architectures(config, "LongcatFlashForCausalLM") + parser = get_model_config_parser(model_config_parser) + config = parser.parse( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) if model_override_args: config.update(model_override_args) diff --git a/test/registered/unit/configs/test_model_config_parser_registry.py b/test/registered/unit/configs/test_model_config_parser_registry.py new file mode 100644 index 000000000..3c1329858 --- /dev/null +++ b/test/registered/unit/configs/test_model_config_parser_registry.py @@ -0,0 +1,62 @@ +"""Unit tests for srt/configs/model_config_parser_registry.py""" + +import unittest + +from transformers import PretrainedConfig + +from sglang.srt.configs.model_config_parser_registry import ( + _MODEL_CONFIG_PARSER_REGISTRY, + ModelConfigParserBase, + get_model_config_parser, + register_model_config_parser, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="stage-a-test-cpu") + + +class _FakeParser(ModelConfigParserBase): + def parse(self, model, trust_remote_code, revision=None, **kwargs): + return PretrainedConfig() + + +class _AnotherFakeParser(ModelConfigParserBase): + def parse(self, model, trust_remote_code, revision=None, **kwargs): + return PretrainedConfig() + + +class TestModelConfigParserRegistry(CustomTestCase): + def setUp(self): + self._saved_registry = dict(_MODEL_CONFIG_PARSER_REGISTRY) + _MODEL_CONFIG_PARSER_REGISTRY.clear() + + def tearDown(self): + _MODEL_CONFIG_PARSER_REGISTRY.clear() + _MODEL_CONFIG_PARSER_REGISTRY.update(self._saved_registry) + + def test_register_then_get_roundtrip(self): + register_model_config_parser("fake")(_FakeParser) + self.assertIsInstance(get_model_config_parser("fake"), _FakeParser) + + def test_register_rejects_non_subclass(self): + class NotAParser: + pass + + with self.assertRaises(ValueError) as ctx: + register_model_config_parser("bad")(NotAParser) + self.assertIn("ModelConfigParserBase", str(ctx.exception)) + + def test_unknown_name_raises_with_registered_list(self): + register_model_config_parser("fake")(_FakeParser) + register_model_config_parser("another")(_AnotherFakeParser) + with self.assertRaises(ValueError) as ctx: + get_model_config_parser("does-not-exist") + msg = str(ctx.exception) + self.assertIn("does-not-exist", msg) + self.assertIn("another", msg) + self.assertIn("fake", msg) + + +if __name__ == "__main__": + unittest.main()