Fix/nemotron mtp quantaized (#19433)

This commit is contained in:
Shaun Kotek
2026-03-03 01:07:46 -08:00
committed by GitHub
parent af0d35b224
commit 4c95953b77
5 changed files with 73 additions and 3 deletions
+4 -1
View File
@@ -665,7 +665,10 @@ class ModelConfig:
quant_cfg = quant_cfg.to_dict()
if quant_cfg is not None:
# Identify modelopt quantization
if "quant_method" not in quant_cfg:
if (
"quant_method" not in quant_cfg
or quant_cfg["quant_method"] == "modelopt"
):
parsed_cfg = self._parse_modelopt_quant_config(
{"quantization": quant_cfg}
)
@@ -418,7 +418,9 @@ class FusedMoE(torch.nn.Module):
# w3, up_proj: Load into second logical weight of w13.
# trtllm cutlass kernel assumes differently
switch_w13 = getattr(self.quant_method, "load_up_proj_weight_first", False)
if (switch_w13 and shard_id == "w1") or (not switch_w13 and shard_id == "w3"):
if (
(switch_w13 and shard_id == "w1") or (not switch_w13 and shard_id == "w3")
) and self.moe_runner_config.is_gated:
start = shard_size
else:
start = 0
@@ -58,6 +58,7 @@ if _is_npu:
try:
from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe
from flashinfer.fused_moe.core import ActivationType
except ImportError:
flashinfer_cutlass_fused_moe = None
@@ -384,6 +385,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
tp_size=layer.moe_tp_size,
tp_rank=layer.moe_tp_rank,
tune_max_num_tokens=next_power_of_2(x.shape[0]),
activation_type=(
ActivationType.Relu2
if moe_runner_config.activation == "relu2"
else ActivationType.Swiglu
),
)[0]
return StandardCombineInput(hidden_states=output)
elif self.use_flashinfer_trtllm_moe:
+1 -1
View File
@@ -297,7 +297,7 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM):
self.model = NemotronHMultiTokenPredictor(
config=config,
quant_config=quant_config,
prefix=add_prefix("model", prefix),
prefix=add_prefix("mtp", prefix),
)
self.lm_head = ParallelLMHead(
@@ -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()