A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 19:30 UTC
Add LM head quantization support for ModelOpt (#42124)
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Wei-Ming Chen committed
May 26, 2026 at 09:21 UTC
6f5b533241929d26f13948e6a48290f6a4f4eefe
8 files changed
+220
-5
tests/model_executor/test_nemotron_h_quantization.py
new
+34
@@ -0,0 +1,34 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from unittest.mock import Mock, patch
5
+
6
+
7
+def test_nemotron_h_lm_head_receives_quant_config():
8
+ from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
9
+
10
+ mock_quant_config = Mock()
11
+
12
+ mock_hf_config = Mock()
13
+ mock_hf_config.vocab_size = 128
14
+ mock_hf_config.hidden_size = 64
15
+
16
+ mock_vllm_config = Mock()
17
+ mock_vllm_config.model_config.hf_config = mock_hf_config
18
+ mock_vllm_config.model_config.dtype = None
19
+ mock_vllm_config.scheduler_config = Mock()
20
+ mock_vllm_config.quant_config = mock_quant_config
21
+
22
+ with (
23
+ patch("vllm.model_executor.models.nemotron_h.NemotronHModel") as MockModel,
24
+ patch("vllm.model_executor.models.nemotron_h.ParallelLMHead") as MockLMHead,
25
+ patch("vllm.model_executor.models.nemotron_h.LogitsProcessor"),
26
+ ):
27
+ MockModel.return_value.make_empty_intermediate_tensors = Mock()
28
+ MockModel.return_value.has_moe = False
29
+
30
+ NemotronHForCausalLM(vllm_config=mock_vllm_config)
31
+
32
+ MockLMHead.assert_called_once()
33
+ call_kwargs = MockLMHead.call_args.kwargs
34
+ assert call_kwargs["quant_config"] is mock_quant_config
tests/model_executor/test_qwen3_5_quantization.py
new
+78
@@ -0,0 +1,78 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from unittest.mock import Mock, patch
5
+
6
+
7
+def test_qwen3_5_lm_head_receives_quant_config():
8
+ from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLMBase
9
+
10
+ mock_quant_config = Mock()
11
+
12
+ mock_hf_config = Mock()
13
+ mock_hf_config.tie_word_embeddings = False
14
+ mock_hf_config.vocab_size = 128
15
+ mock_hf_config.hidden_size = 64
16
+
17
+ mock_vllm_config = Mock()
18
+ mock_vllm_config.model_config.hf_text_config = mock_hf_config
19
+ mock_vllm_config.cache_config.mamba_cache_mode = "align"
20
+ mock_vllm_config.scheduler_config = Mock()
21
+ mock_vllm_config.quant_config = mock_quant_config
22
+ mock_vllm_config.lora_config = None
23
+
24
+ mock_pp_group = Mock()
25
+ mock_pp_group.is_last_rank = True
26
+
27
+ with (
28
+ patch("vllm.model_executor.models.qwen3_5.Qwen3_5Model") as MockModel,
29
+ patch("vllm.model_executor.models.qwen3_5.ParallelLMHead") as MockLMHead,
30
+ patch("vllm.model_executor.models.qwen3_5.LogitsProcessor"),
31
+ patch(
32
+ "vllm.model_executor.models.qwen3_5.get_pp_group",
33
+ return_value=mock_pp_group,
34
+ ),
35
+ ):
36
+ MockModel.return_value.make_empty_intermediate_tensors = Mock()
37
+
38
+ Qwen3_5ForCausalLMBase(vllm_config=mock_vllm_config)
39
+
40
+ MockLMHead.assert_called_once()
41
+ call_kwargs = MockLMHead.call_args.kwargs
42
+ assert call_kwargs["quant_config"] is mock_quant_config
43
+
44
+
45
+def test_qwen3_5_mtp_lm_head_receives_quant_config():
46
+ from vllm.config import CompilationMode
47
+ from vllm.model_executor.models.qwen3_5_mtp import Qwen3_5MTP
48
+
49
+ mock_quant_config = Mock()
50
+
51
+ mock_hf_config = Mock()
52
+ mock_hf_config.tie_word_embeddings = False
53
+ mock_hf_config.vocab_size = 128
54
+ mock_hf_config.hidden_size = 64
55
+
56
+ mock_vllm_config = Mock()
57
+ mock_vllm_config.model_config.hf_text_config = mock_hf_config
58
+ mock_vllm_config.cache_config.mamba_cache_mode = "align"
59
+ mock_vllm_config.compilation_config.mode = CompilationMode.NONE
60
+ mock_vllm_config.quant_config = mock_quant_config
61
+
62
+ mock_pp_group = Mock()
63
+ mock_pp_group.is_last_rank = True
64
+
65
+ with (
66
+ patch("vllm.model_executor.models.qwen3_5_mtp.Qwen3_5MultiTokenPredictor"),
67
+ patch("vllm.model_executor.models.qwen3_5_mtp.ParallelLMHead") as MockLMHead,
68
+ patch("vllm.model_executor.models.qwen3_5_mtp.LogitsProcessor"),
69
+ patch(
70
+ "vllm.model_executor.models.qwen3_5_mtp.get_pp_group",
71
+ return_value=mock_pp_group,
72
+ ),
73
+ ):
74
+ Qwen3_5MTP(vllm_config=mock_vllm_config)
75
+
76
+ MockLMHead.assert_called_once()
77
+ call_kwargs = MockLMHead.call_args.kwargs
78
+ assert call_kwargs["quant_config"] is mock_quant_config
tests/quantization/test_modelopt.py
+93
-1
@@ -7,13 +7,24 @@ Run `pytest tests/quantization/test_modelopt.py`.
7
8
import os
9
from typing import Any, NoReturn
10
-from unittest.mock import MagicMock, patch
10
+from unittest.mock import MagicMock, Mock, patch
11
12
import pytest
13
import torch
14
15
from tests.quantization.utils import is_quant_method_supported
16
from vllm.config.model import ModelConfig
17
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
18
+from vllm.model_executor.layers.quantization.modelopt import (
19
+ ModelOptFp8Config,
20
+ ModelOptMixedPrecisionConfig,
21
+ ModelOptNvFp4Config,
22
+ ModelOptNvFp4LinearMethod,
23
+)
24
+from vllm.model_executor.layers.vocab_parallel_embedding import (
25
+ ParallelLMHead,
26
+ VocabParallelEmbedding,
27
+)
28
29
30
@pytest.fixture(scope="function", autouse=True)
@@ -44,6 +55,87 @@ def _snapshot_download_or_skip(model_id: str) -> str:
55
_skip(f"Failed to download {model_id} from the HF Hub: {e}")
56
57
58
+def _mock_lm_head() -> Mock:
59
+ lm_head = Mock(spec=ParallelLMHead)
60
+ lm_head.__class__ = ParallelLMHead
61
+ return lm_head
62
+
63
+
64
+def _mixed_precision_config(quantized_layers: dict) -> ModelOptMixedPrecisionConfig:
65
+ return ModelOptMixedPrecisionConfig(
66
+ kv_cache_quant_method=None,
67
+ exclude_modules=[],
68
+ quantized_layers=quantized_layers,
69
+ fp8_config=ModelOptFp8Config(
70
+ quant_method="FP8",
71
+ is_checkpoint_fp8_serialized=True,
72
+ kv_cache_quant_method=None,
73
+ exclude_modules=[],
74
+ ),
75
+ nvfp4_config=ModelOptNvFp4Config(
76
+ is_checkpoint_nvfp4_serialized=True,
77
+ kv_cache_quant_algo=None,
78
+ exclude_modules=[],
79
+ ),
80
+ w4a16_nvfp4_config=ModelOptNvFp4Config(
81
+ quant_method="W4A16_NVFP4",
82
+ is_checkpoint_nvfp4_serialized=True,
83
+ kv_cache_quant_algo=None,
84
+ exclude_modules=[],
85
+ ),
86
+ )
87
+
88
+
89
+def test_modelopt_nvfp4_quantizes_parallel_lm_head():
90
+ config = ModelOptNvFp4Config(
91
+ is_checkpoint_nvfp4_serialized=True,
92
+ kv_cache_quant_algo=None,
93
+ exclude_modules=[],
94
+ )
95
+
96
+ with patch(
97
+ "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel"
98
+ ):
99
+ method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
100
+
101
+ assert isinstance(method, ModelOptNvFp4LinearMethod)
102
+
103
+
104
+def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized():
105
+ config = ModelOptNvFp4Config(
106
+ is_checkpoint_nvfp4_serialized=True,
107
+ kv_cache_quant_algo=None,
108
+ exclude_modules=["lm_head"],
109
+ )
110
+
111
+ method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
112
+
113
+ assert isinstance(method, UnquantizedLinearMethod)
114
+
115
+
116
+def test_modelopt_mixed_precision_quantizes_parallel_lm_head():
117
+ config = _mixed_precision_config(
118
+ {"lm_head": {"quant_algo": "NVFP4", "group_size": 16}}
119
+ )
120
+
121
+ with patch(
122
+ "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel"
123
+ ):
124
+ method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
125
+
126
+ assert isinstance(method, ModelOptNvFp4LinearMethod)
127
+
128
+
129
+def test_vocab_parallel_embedding_weight_loader_accepts_scalar_scale():
130
+ holder = Mock()
131
+ scale = torch.nn.Parameter(torch.empty(1))
132
+ loaded_scale = torch.tensor(2.0)
133
+
134
+ VocabParallelEmbedding.weight_loader(holder, scale, loaded_scale)
135
+
136
+ assert torch.equal(scale, loaded_scale.reshape(1))
137
+
138
+
139
@pytest.mark.skipif(
140
not is_quant_method_supported("modelopt"),
141
reason="ModelOpt FP8 is not supported on this GPU type.",
vllm/model_executor/layers/quantization/modelopt.py
+5
-4
@@ -81,6 +81,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
81
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
82
requantize_with_max_scale,
83
)
84
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
85
from vllm.model_executor.parameter import (
86
BlockQuantScaleParameter,
87
ChannelQuantScaleParameter,
@@ -182,7 +183,7 @@ class ModelOptQuantConfigBase(QuantizationConfig):
183
184
# handle exclusion
185
if self.is_layer_excluded(prefix):
185
- if isinstance(layer, LinearBase):
186
+ if isinstance(layer, (LinearBase, ParallelLMHead)):
187
return UnquantizedLinearMethod()
188
return None
189
@@ -195,7 +196,7 @@ class ModelOptQuantConfigBase(QuantizationConfig):
196
return UnquantizedLinearMethod()
197
198
# now, the layer is quantized, handle it here
198
- if isinstance(layer, LinearBase):
199
+ if isinstance(layer, (LinearBase, ParallelLMHead)):
200
quant_method = self.LinearMethodCls(self)
201
if getattr(quant_method, "backend", "") == "marlin":
202
quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
@@ -2371,13 +2372,13 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2372
2373
# Excluded layers
2374
if self.is_layer_excluded(prefix):
2374
- if isinstance(layer, LinearBase):
2375
+ if isinstance(layer, (LinearBase, ParallelLMHead)):
2376
return UnquantizedLinearMethod()
2377
return None
2378
2379
quant_algo = self._resolve_quant_algo(prefix)
2380
2380
- if isinstance(layer, LinearBase):
2381
+ if isinstance(layer, (LinearBase, ParallelLMHead)):
2382
if quant_algo == "FP8":
2383
return ModelOptFp8LinearMethod(self.fp8_config)
2384
if quant_algo == "NVFP4":
vllm/model_executor/layers/vocab_parallel_embedding.py
+7
@@ -290,6 +290,7 @@ class VocabParallelEmbedding(PluggableLayer):
290
291
if params_dtype is None:
292
params_dtype = torch.get_default_dtype()
293
+ self.params_dtype = params_dtype
294
# Divide the weight matrix along the vocabulary dimension.
295
self.num_added_embeddings = self.num_embeddings - self.org_vocab_size
296
self.num_embeddings_per_partition = divide(
@@ -438,6 +439,12 @@ class VocabParallelEmbedding(PluggableLayer):
439
# If parameter does not have output dim, then it should
440
# be copied onto all gpus (e.g. g_idx for act_order gptq).
441
if output_dim is None:
442
+ if (
443
+ loaded_weight.ndim == 0
444
+ and param.data.ndim == 1
445
+ and param.data.numel() == 1
446
+ ):
447
+ loaded_weight = loaded_weight.reshape(1)
448
assert param.data.shape == loaded_weight.shape
449
param.data.copy_(loaded_weight)
450
return
vllm/model_executor/models/nemotron_h.py
+1
@@ -875,6 +875,7 @@ class NemotronHForCausalLM(
875
self.lm_head = ParallelLMHead(
876
config.vocab_size,
877
config.hidden_size,
878
+ quant_config=self.quant_config,
879
prefix=maybe_prefix(prefix, "lm_head"),
880
)
881
vllm/model_executor/models/qwen3_5.py
+1
@@ -477,6 +477,7 @@ class Qwen3_5ForCausalLMBase(
477
self.lm_head = ParallelLMHead(
478
config.vocab_size,
479
config.hidden_size,
480
+ quant_config=self.quant_config,
481
prefix=maybe_prefix(prefix, "lm_head"),
482
)
483
else:
vllm/model_executor/models/qwen3_5_mtp.py
+1
@@ -381,6 +381,7 @@ class Qwen3_5MTP(nn.Module, SupportsMultiModal):
381
self.lm_head = ParallelLMHead(
382
config.vocab_size,
383
config.hidden_size,
384
+ quant_config=self.quant_config,
385
prefix=maybe_prefix(prefix, "lm_head"),
386
)
387
else: