minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

quant: extract shared checkpoint quant metadata resolver (#35172)

Browse Source
This commit is contained in:
Mick
2026-08-19 08:26:41 +08:00
committed by GitHub
parent e73201e462
commit ef490853bb
4 changed files with 332 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_loader/checkpoint_quantization.py
+100
View File
@@ -0,0 +1,100 @@
# SPDX-License-Identifier: Apache-2.0
"""Pure-data helpers for quantization metadata in Hugging Face configs."""
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass
from typing import Any, Literal, Mapping, TypeAlias
__all__ = [
"CheckpointQuantSpec",
"QuantMetadataSource",
"resolve_checkpoint_quant_spec",
]
QuantMetadataSource: TypeAlias = Literal[
"quantization_config",
"text_config.quantization_config",
"compression_config",
]
@dataclass(slots=True)
class CheckpointQuantSpec:
"""Quantization metadata declared by a checkpoint.
``declared_method`` preserves ``quant_method`` verbatim and is never inferred
from backend-specific fields. This intentionally contains no runtime
quantization classes, model construction, or layer hierarchy.
"""
declared_method: str | None
config: dict[str, Any]
source: QuantMetadataSource
def _get_field(config: object, name: str) -> Any:
if isinstance(config, Mapping):
return config.get(name)
return getattr(config, name, None)
def _to_metadata_dict(value: object, source: QuantMetadataSource) -> dict[str, Any]:
if isinstance(value, Mapping):
return deepcopy(dict(value))
to_dict = getattr(value, "to_dict", None)
if callable(to_dict):
metadata = to_dict()
if isinstance(metadata, Mapping):
return deepcopy(dict(metadata))
raise TypeError(
f"{source} must be a mapping or expose to_dict(), "
f"got {type(value).__name__}"
)
def _select_hf_quant_metadata(
hf_config: object,
) -> tuple[QuantMetadataSource, object] | None:
value = _get_field(hf_config, "quantization_config")
if value is not None:
return "quantization_config", value
text_config = _get_field(hf_config, "text_config")
value = _get_field(text_config, "quantization_config")
if value is not None:
return "text_config.quantization_config", value
value = _get_field(hf_config, "compression_config")
if value is not None:
return "compression_config", value
return None
def resolve_checkpoint_quant_spec(hf_config: object) -> CheckpointQuantSpec | None:
"""Resolve checkpoint quantization metadata from an HF config.
The lookup order matches SRT's checkpoint loader: top-level
``quantization_config``, the text sub-config used by some multimodal
checkpoints, then ``compression_config``. The returned metadata is deep-copied
so callers can attach runtime-only fields without mutating the HF config.
"""
selected = _select_hf_quant_metadata(hf_config)
if selected is None:
return None
source, value = selected
config = _to_metadata_dict(value, source)
declared_method = config.get("quant_method")
return CheckpointQuantSpec(
declared_method=(declared_method if isinstance(declared_method, str) else None),
config=config,
source=source,
)
python/sglang/srt/model_loader/weight_utils.py
+7 -15
View File
@@ -42,15 +42,16 @@ from tqdm.auto import tqdm
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.configs.model_config import REQUANTIZATION_METHODS, ModelConfig
from sglang.srt.distributed import (
get_world_group,
)
from sglang.srt.distributed import get_world_group
from sglang.srt.layers.quantization import QuantizationConfig, get_quantization_config
from sglang.srt.layers.quantization.fp8 import Fp8Config
from sglang.srt.layers.quantization.modelopt_quant import (
ModelOptFp4Config,
ModelOptFp8Config,
)
from sglang.srt.model_loader.checkpoint_quantization import (
resolve_checkpoint_quant_spec,
)
from sglang.srt.model_loader.ci_weight_validation import (
ci_download_with_validation_and_retry,
ci_validate_and_cleanup_local_snapshot,
@@ -271,18 +272,9 @@ def get_quant_config(
if model_config.quantization == "gguf":
return quant_cls.from_config({})
# Read the quantization config from the HF model config, if available.
hf_quant_config = getattr(model_config.hf_config, "quantization_config", None)
# some vision model may keep quantization_config in their text_config
hf_text_config = getattr(model_config.hf_config, "text_config", None)
if hf_quant_config is None and hf_text_config is not None:
hf_quant_config = getattr(hf_text_config, "quantization_config", None)
if hf_quant_config is None:
# compressed-tensors uses a compressions_config
hf_quant_config = getattr(model_config.hf_config, "compression_config", None)
if hf_quant_config is not None:
if not isinstance(hf_quant_config, dict):
hf_quant_config = hf_quant_config.to_dict()
checkpoint_quant_spec = resolve_checkpoint_quant_spec(model_config.hf_config)
if checkpoint_quant_spec is not None:
hf_quant_config = checkpoint_quant_spec.config
# For modelopt_mixed, config.json's quantization_config may not
# contain all runtime metadata. Fall through to the file-based
# hf_quant_config.json path when the per-layer map or KV-cache
test/registered/unit/model_loader/test_modelopt_loader.py
+95
View File
@@ -5,7 +5,10 @@ This test module verifies the functionality of ModelOptModelLoader, which
applies NVIDIA Model Optimizer quantization to models during loading.
"""
import json
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
@@ -693,6 +696,98 @@ class TestModelOptFp4LoaderSelection(CustomTestCase):
class TestModelOptMixedPrecisionConfig(CustomTestCase):
def test_incomplete_inline_config_falls_back_to_hf_quant_config_file(self):
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
}
file_quantized_layers = {
"model.layers.0.self_attn.q_proj": {"quant_algo": "FP8"}
}
file_config = {
"producer": {"name": "modelopt"},
"quantization": {
"quant_algo": "MIXED_PRECISION",
"kv_cache_quant_algo": "FP8",
"exclude_modules": [],
"quantized_layers": file_quantized_layers,
},
}
inline_configs = (
{
"quant_method": "modelopt_mixed",
"quant_algo": "MIXED_PRECISION",
"kv_cache_quant_algo": "NVFP4",
},
{
"quant_method": "modelopt_mixed",
"quant_algo": "MIXED_PRECISION",
"quantized_layers": {
"inline.layer": {"quant_algo": "NVFP4", "group_size": 16}
},
},
)
with tempfile.TemporaryDirectory() as model_path:
Path(model_path, "hf_quant_config.json").write_text(
json.dumps(file_config), encoding="utf-8"
)
for inline_config in inline_configs:
with self.subTest(inline_config=inline_config):
model_config = SimpleNamespace(
quantization="modelopt_mixed",
hf_config=SimpleNamespace(
quantization_config=inline_config,
),
model_path=model_path,
revision=None,
is_draft_model=False,
is_draft_quantization_explicit=False,
)
config = get_quant_config(
model_config, LoadConfig(), packed_modules_mapping
)
self.assertIsInstance(config, ModelOptMixedPrecisionConfig)
self.assertEqual(config.quantized_layers, file_quantized_layers)
self.assertEqual(config.kv_cache_quant_algo, "FP8")
self.assertEqual(
config.packed_modules_mapping, packed_modules_mapping
)
@patch("sglang.srt.model_loader.weight_utils.snapshot_download")
def test_complete_inline_config_does_not_download_metadata(self, mock_download):
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
}
inline_quantized_layers = {
"model.layers.0.self_attn.q_proj": {"quant_algo": "FP8"}
}
model_config = SimpleNamespace(
quantization="modelopt_mixed",
hf_config=SimpleNamespace(
quantization_config={
"quant_method": "modelopt_mixed",
"quant_algo": "MIXED_PRECISION",
"kv_cache_scheme": {"type": "float", "num_bits": 8},
"exclude_modules": [],
"quantized_layers": inline_quantized_layers,
}
),
model_path="remote/model",
revision=None,
is_draft_model=False,
is_draft_quantization_explicit=False,
)
config = get_quant_config(model_config, LoadConfig(), packed_modules_mapping)
self.assertIsInstance(config, ModelOptMixedPrecisionConfig)
self.assertEqual(config.quantized_layers, inline_quantized_layers)
self.assertEqual(config.kv_cache_quant_algo, "FP8")
self.assertEqual(config.packed_modules_mapping, packed_modules_mapping)
mock_download.assert_not_called()
def test_minimax_mixed_precision_resolves_runtime_names_and_mxfp8(self):
quant_config = ModelOptMixedPrecisionConfig.from_config(
{
test/registered/unit/test_checkpoint_quantization.py
+130
View File
@@ -0,0 +1,130 @@
# SPDX-License-Identifier: Apache-2.0
import unittest
from sglang.srt.model_loader.checkpoint_quantization import (
CheckpointQuantSpec,
resolve_checkpoint_quant_spec,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class _ConfigObject:
def __init__(self, **values):
self.__dict__.update(values)
class _QuantConfigObject:
def __init__(self, values):
self._values = values
def to_dict(self):
return self._values
class TestResolveCheckpointQuantSpec(CustomTestCase):
def test_top_level_quantization_config(self):
config = {
"quantization_config": {
"quant_method": "fp8",
"activation_scheme": "dynamic",
}
}
spec = resolve_checkpoint_quant_spec(config)
self.assertEqual(
spec,
CheckpointQuantSpec(
declared_method="fp8",
config={"quant_method": "fp8", "activation_scheme": "dynamic"},
source="quantization_config",
),
)
def test_text_config_fallback_supports_config_objects(self):
config = _ConfigObject(
text_config=_ConfigObject(
quantization_config={"quant_method": "gptq", "bits": 4}
),
compression_config={"quant_method": "compressed-tensors"},
)
spec = resolve_checkpoint_quant_spec(config)
self.assertIsNotNone(spec)
self.assertEqual(spec.declared_method, "gptq")
self.assertEqual(spec.source, "text_config.quantization_config")
def test_compression_config_fallback(self):
config = _ConfigObject(
compression_config={"quant_method": "compressed-tensors"}
)
spec = resolve_checkpoint_quant_spec(config)
self.assertIsNotNone(spec)
self.assertEqual(spec.declared_method, "compressed-tensors")
self.assertEqual(spec.source, "compression_config")
def test_modelopt_quant_algo_does_not_infer_declared_method(self):
config = {
"quantization_config": {
"quant_algo": "FP8",
"exclude_modules": ["lm_head"],
}
}
spec = resolve_checkpoint_quant_spec(config)
self.assertIsNotNone(spec)
self.assertIsNone(spec.declared_method)
self.assertEqual(spec.config["quant_algo"], "FP8")
def test_quant_config_object_is_converted(self):
config = _ConfigObject(
quantization_config=_QuantConfigObject(
{"quant_method": "bitsandbytes", "load_in_4bit": True}
)
)
spec = resolve_checkpoint_quant_spec(config)
self.assertIsNotNone(spec)
self.assertEqual(spec.config["load_in_4bit"], True)
def test_lookup_priority_matches_srt_loader(self):
config = {
"quantization_config": {},
"text_config": {"quantization_config": {"quant_method": "gptq"}},
"compression_config": {"quant_method": "compressed-tensors"},
}
spec = resolve_checkpoint_quant_spec(config)
self.assertIsNotNone(spec)
self.assertEqual(spec.config, {})
self.assertEqual(spec.source, "quantization_config")
def test_metadata_is_deep_copied(self):
metadata = {"quant_method": "fp8", "modules_to_not_convert": ["lm_head"]}
spec = resolve_checkpoint_quant_spec({"quantization_config": metadata})
self.assertIsNotNone(spec)
spec.config["modules_to_not_convert"].append("embed_tokens")
self.assertEqual(metadata["modules_to_not_convert"], ["lm_head"])
def test_missing_metadata_returns_none(self):
self.assertIsNone(resolve_checkpoint_quant_spec({"model_type": "qwen3_vl"}))
def test_invalid_metadata_type_has_clear_error(self):
with self.assertRaisesRegex(TypeError, "quantization_config must be a mapping"):
resolve_checkpoint_quant_spec({"quantization_config": "fp8"})
if __name__ == "__main__":
unittest.main()