Add --model-config-parser registry for pluggable config formats (#25050)
Signed-off-by: Xingyu Liu <charlotteliu12x@gmail.com>
This commit is contained in:
@@ -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]()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user