Add --model-config-parser registry for pluggable config formats (#25050)

Signed-off-by: Xingyu Liu <charlotteliu12x@gmail.com>
This commit is contained in:
Xingyu Liu
2026-05-14 01:54:06 -07:00
committed by GitHub
parent 4be25f2428
commit 2279b79f35
5 changed files with 326 additions and 121 deletions
@@ -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,
)
@@ -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]()
+49 -22
View File
@@ -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(
@@ -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,34 +51,20 @@ def _apply_deepseek_ocr_overrides(config, model):
config._name_or_path = model
@lru_cache_frozenset(maxsize=32)
def get_config(
model: str,
@register_model_config_parser("hf")
class HfModelConfigParser(ModelConfigParserBase):
def parse(
self,
model,
trust_remote_code: bool,
revision: Optional[str] = None,
model_override_args: Optional[dict] = None,
**kwargs,
):
is_gguf = check_gguf_file(model)
if is_gguf:
_ensure_gguf_version()
kwargs["gguf_file"] = model
model = Path(model).parent
model = resolve_runai_obj_uri(model)
if is_remote_url(model):
client = create_remote_connector(model)
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
model,
trust_remote_code=trust_remote_code,
revision=revision,
**kwargs,
)
if (
@@ -124,7 +115,9 @@ def get_config(
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)
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):
@@ -168,6 +161,63 @@ def get_config(
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)
if is_remote_url(model):
client = create_remote_connector(model)
client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"])
model = client.get_local_dir()
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"
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)
@@ -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()