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

2026-08-19 08:26:41 +08:00
parent e73201e462
commit ef490853bb
4 changed files with 332 additions and 15 deletions
@@ -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(
            {