Fix/nemotron mtp quantaized (#19433)

2026-03-03 01:07:46 -08:00
parent af0d35b224
commit 4c95953b77
5 changed files with 73 additions and 3 deletions
@@ -561,5 +561,64 @@ class TestModelOptLoaderIntegration(CustomTestCase):
        self.assertEqual(server_args.model_path, "TinyLlama/TinyLlama-1.1B-Chat-v1.0")


+class TestParseQuantHfConfig(CustomTestCase):
+    """Tests for _parse_quant_hf_config and _parse_modelopt_quant_config.
+
+    Regression tests for the fix where quant_method='modelopt' ignoring quant_algo.
+    """
+
+    # (quant_config_input, expected_quant_method)
+    _MODELOPT_CASES = [
+        ({"quant_method": "modelopt", "quant_algo": "FP8"}, "modelopt_fp8"),
+        ({"quant_method": "modelopt", "quant_algo": "FP4"}, "modelopt_fp4"),
+        ({"quant_method": "modelopt", "quant_algo": "NVFP4"}, "modelopt_fp4"),
+        ({"quant_method": "modelopt", "quant_algo": "MIXED_PRECISION"}, "w4afp8"),
+        ({"quant_algo": "FP8"}, "modelopt_fp8"),
+        ({"quant_algo": "FP4"}, "modelopt_fp4"),
+        ({"quant_algo": "MIXED_PRECISION"}, "w4afp8"),
+        ({"quant_method": "modelopt"}, "modelopt"),
+    ]
+
+    def setUp(self):
+        """Set up a real ModelConfig using TinyLlama (already used elsewhere)."""
+        self.mock_tp_rank = patch(
+            "sglang.srt.distributed.parallel_state.get_tensor_model_parallel_rank",
+            return_value=0,
+        )
+        self.mock_tp_rank.start()
+
+        self.mock_mp_is_initialized = patch(
+            "sglang.srt.distributed.parallel_state.model_parallel_is_initialized",
+            return_value=True,
+        )
+        self.mock_mp_is_initialized.start()
+
+        self.model_config = ModelConfig(
+            model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+        )
+
+    def tearDown(self):
+        self.mock_tp_rank.stop()
+        self.mock_mp_is_initialized.stop()
+
+    def test_modelopt_quant_parsing(self):
+        """Modelopt quant configs must resolve to the correct quant_method."""
+        for quant_cfg_input, expected in self._MODELOPT_CASES:
+            with self.subTest(quant_cfg=quant_cfg_input):
+                self.model_config.hf_config.quantization_config = dict(quant_cfg_input)
+                result = self.model_config._parse_quant_hf_config()
+                self.assertEqual(result["quant_method"], expected)
+
+    def test_non_modelopt_quant_method_unchanged(self):
+        """Non-modelopt quant_method (e.g. 'gptq') must NOT enter the modelopt path."""
+        self.model_config.hf_config.quantization_config = {
+            "quant_method": "gptq",
+            "bits": 4,
+        }
+        result = self.model_config._parse_quant_hf_config()
+        self.assertEqual(result["quant_method"], "gptq")
+        self.assertNotIn("quant_algo", result)
+
+
 if __name__ == "__main__":
    unittest.main()