diff --git a/python/sglang/srt/model_loader/checkpoint_quantization.py b/python/sglang/srt/model_loader/checkpoint_quantization.py new file mode 100644 index 000000000..3578e4f1b --- /dev/null +++ b/python/sglang/srt/model_loader/checkpoint_quantization.py @@ -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, + ) diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 2d25f2add..965123fcc 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -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 diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index b6fc958ef..eddc599f2 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -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( { diff --git a/test/registered/unit/test_checkpoint_quantization.py b/test/registered/unit/test_checkpoint_quantization.py new file mode 100644 index 000000000..aae75d58d --- /dev/null +++ b/test/registered/unit/test_checkpoint_quantization.py @@ -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()