[Bugfix][Quantization] Give LM heads standard linear metadata (#58444)
Signed-off-by: mgoin <mgoin64@gmail.com>
Misha Goin committed
Sep 24, 2026 at 19:39 UTC
4f72bd0152e74f3485216aebb3f5d276b6726a9a
8 files changed
+93
-70
tests/quantization/test_lm_head.py
+74
@@ -5,14 +5,26 @@
5
Run `pytest tests/quantization/test_quant_lm_head_true.py --forked`.
6
"""
7
8
+from types import SimpleNamespace
9
+
10
import pytest
11
import torch
12
+from compressed_tensors.quantization import preset_name_to_scheme
13
14
from tests.quantization.utils import load_model_without_vllm_runner
15
+from vllm.model_executor.layers.linear import ColumnParallelLinear
16
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
17
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod
18
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
19
+ CompressedTensorsConfig,
20
+)
21
+from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config
22
from vllm.model_executor.layers.vocab_parallel_embedding import (
23
+ ParallelLMHead,
24
UnquantizedEmbeddingMethod,
25
)
26
+from vllm.platforms import current_platform
27
+from vllm.utils.import_utils import has_humming
28
29
PROMPT = "On the surface of Mars, we found"
30
@@ -44,3 +56,65 @@ def test_lm_head(
56
assert isinstance(lm_head_layer.quant_method, AutoGPTQLinearMethod)
57
else:
58
assert isinstance(lm_head_layer.quant_method, UnquantizedEmbeddingMethod)
59
+
60
+
61
+@pytest.mark.skipif(
62
+ not current_platform.is_cuda() or not has_humming(), reason="requires Humming/CUDA"
63
+)
64
+@pytest.mark.parametrize(
65
+ "preset,quant_format",
66
+ [
67
+ ("FP8_DYNAMIC", "float-quantized"),
68
+ ("NVFP4A16", "nvfp4-pack-quantized"),
69
+ ("W4A16", "pack-quantized"),
70
+ ("W4A16_NVFP4", None),
71
+ ],
72
+)
73
+@pytest.mark.parametrize("bias", [False, True])
74
+@torch.inference_mode()
75
+def test_quantized_lm_head_matches_linear(
76
+ dist_init, default_vllm_config, preset, quant_format, bias
77
+):
78
+ """An explicitly quantized head must produce the same logits as a linear layer."""
79
+ default_vllm_config.model_config = SimpleNamespace(
80
+ dtype=torch.bfloat16, head_dtype=None
81
+ )
82
+ default_vllm_config.kernel_config.linear_backend = "humming"
83
+ if quant_format is None:
84
+ quant_config = ModelOptNvFp4Config(
85
+ quant_method=preset,
86
+ is_checkpoint_nvfp4_serialized=True,
87
+ kv_cache_quant_algo=None,
88
+ exclude_modules=[],
89
+ )
90
+ else:
91
+ scheme = preset_name_to_scheme(preset, targets=["Linear", "lm_head"])
92
+ quant_config = CompressedTensorsConfig.from_config(
93
+ {"config_groups": {"group_0": scheme.model_dump()}, "format": quant_format}
94
+ )
95
+ with torch.device("cuda"):
96
+ kwargs = dict(
97
+ bias=bias,
98
+ params_dtype=torch.bfloat16,
99
+ quant_config=quant_config,
100
+ disable_tp=True,
101
+ )
102
+ head = ParallelLMHead(500, 256, prefix="lm_head", **kwargs)
103
+ linear = ColumnParallelLinear(256, 512, prefix="proj", **kwargs)
104
+ for name, param in linear.named_parameters():
105
+ if "scale" in name:
106
+ param.fill_(1.0)
107
+ elif name == "weight_shape":
108
+ param.copy_(torch.tensor([512, 256]))
109
+ elif param.is_floating_point():
110
+ param.copy_(torch.randn(param.shape, dtype=torch.float32))
111
+ else:
112
+ param.random_(0, 127)
113
+ head.load_state_dict(linear.state_dict())
114
+ for layer in (head, linear):
115
+ layer.quant_method.process_weights_after_loading(layer)
116
+ x = torch.randn(8, 256, dtype=torch.bfloat16)
117
+ expected, _ = linear(x)
118
+ actual = LogitsProcessor(500)(head, x, head.bias)
119
+ assert torch.isfinite(actual).all()
120
+ torch.testing.assert_close(actual, expected[:, :500], rtol=0, atol=0)
tests/quantization/test_modelopt.py
-43
@@ -718,49 +718,6 @@ def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls):
718
assert isinstance(kernel, kernel_cls)
719
720
721
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
722
-def test_modelopt_linear_exposes_humming_layer_attrs(dist_init, monkeypatch):
723
- """``prepare_humming_linear_layer_config`` reads ``output_partition_sizes``
724
- and ``has_bias`` straight off the layer, so ``--linear-backend=humming``
725
- needs create_weights to leave both there. Nothing else sets
726
- ``output_partition_sizes``; ``LinearBase`` sets ``has_bias`` but
727
- ``ParallelLMHead`` does not.
728
- """
729
- from vllm.config.quantization import QuantSpec
730
- from vllm.model_executor.layers.quantization import modelopt as mo
731
-
732
- monkeypatch.setattr(
733
- mo, "select_linear_kernel", lambda spec, layer, rt, **kwargs: Mock()
734
- )
735
- monkeypatch.setattr(mo, "expose_input_quant_key", lambda layer, kernel: None)
736
-
737
- def build(layer):
738
- method = ModelOptLinearMethod.__new__(ModelOptLinearMethod)
739
- method.spec = QuantSpec(weight=kNvfp4Static, activation=None)
740
- method.ctx = mo.CkptCtx(group_size=16)
741
- method.fmt = mo.FormatScheme()
742
- method.wkey = mo.SCHEME_FOR[kNvfp4Static]
743
- method.akey = None
744
- method.input_dtype = method.out_dtype = torch.bfloat16
745
- method.marlin_input_dtype = None
746
- method.create_weights(
747
- layer, 64, [32, 32], 64, 64, torch.bfloat16, weight_loader=Mock()
748
- )
749
-
750
- # ParallelLMHead-style: a bias slot but no has_bias attribute.
751
- lm_head = torch.nn.Module()
752
- lm_head.register_parameter("bias", None)
753
- build(lm_head)
754
- assert lm_head.output_partition_sizes == [32, 32]
755
- assert lm_head.has_bias is False
756
-
757
- # LinearBase already decided has_bias; we must not overwrite it.
758
- linear = torch.nn.Module()
759
- linear.has_bias = True
760
- build(linear)
761
- assert linear.has_bias is True
762
-
763
-
721
@pytest.mark.parametrize(
722
"quant_method, expected_use_a16, act_key_is_none",
723
[
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py
-2
@@ -119,8 +119,6 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
119
layer.output_size_per_partition = output_size_per_partition
120
layer.output_partition_sizes = output_partition_sizes
121
layer.params_dtype = params_dtype
122
- if not hasattr(layer, "has_bias"):
123
- layer.has_bias = False
122
123
mp_linear_kernel_config = MPLinearLayerConfig(
124
full_weight_shape=(input_size, output_size),
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py
-2
@@ -147,8 +147,6 @@ class CompressedTensorsWNA4Int(CompressedTensorsScheme):
147
layer.output_size_per_partition = output_size_per_partition
148
layer.output_partition_sizes = output_partition_sizes
149
layer.params_dtype = params_dtype
150
- if not hasattr(layer, "has_bias"):
151
- layer.has_bias = False
150
151
mp_config = MPLinearLayerConfig(
152
full_weight_shape=(input_size, output_size),
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py
-2
@@ -147,8 +147,6 @@ class CompressedTensorsWNA8Int(CompressedTensorsScheme):
147
layer.output_size_per_partition = output_size_per_partition
148
layer.output_partition_sizes = output_partition_sizes
149
layer.params_dtype = params_dtype
150
- if not hasattr(layer, "has_bias"):
151
- layer.has_bias = False
150
151
mp_config = MPLinearLayerConfig(
152
full_weight_shape=(input_size, output_size),
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py
-4
@@ -97,12 +97,8 @@ class CompressedTensorsWNA8O8Int(CompressedTensorsScheme):
97
output_size_per_partition = sum(output_partition_sizes)
98
layer.input_size_per_partition = input_size_per_partition
99
layer.output_size_per_partition = output_size_per_partition
100
- # Set for kernels' weight prep; also covers ParallelLMHead, which does
101
- # not set these in __init__.
100
layer.output_partition_sizes = output_partition_sizes
101
layer.params_dtype = params_dtype
104
- if not hasattr(layer, "has_bias"):
105
- layer.has_bias = False
102
103
mp_config = MPLinearLayerConfig(
104
full_weight_shape=(input_size, output_size),
vllm/model_executor/layers/quantization/modelopt.py
-6
@@ -2508,12 +2508,6 @@ class ModelOptLinearMethod(LinearMethodBase):
2508
layer.logical_widths = output_partition_sizes
2509
layer.input_size_per_partition = input_size_per_partition
2510
layer.output_size_per_partition = sum(output_partition_sizes)
2511
- # Humming reads both off the layer in
2512
- # prepare_humming_linear_layer_config. LinearBase sets them itself;
2513
- # ParallelLMHead does not, so supply them here.
2514
- layer.output_partition_sizes = output_partition_sizes
2515
- if not hasattr(layer, "has_bias"):
2516
- layer.has_bias = getattr(layer, "bias", None) is not None
2511
shapes = Shapes(output_partition_sizes, input_size_per_partition, params_dtype)
2512
2513
self.wkey.create_weights(layer, WEIGHT, self.ctx, shapes, weight_loader)
vllm/model_executor/layers/vocab_parallel_embedding.py
+19
-11
@@ -278,12 +278,12 @@ class VocabParallelEmbedding(PluggableLayer):
278
self.num_embeddings = num_embeddings
279
self.padding_size = padding_size
280
self.org_vocab_size = org_num_embeddings or num_embeddings
281
- num_added_embeddings = num_embeddings - self.org_vocab_size
281
+ self.num_added_embeddings = num_embeddings - self.org_vocab_size
282
self.org_vocab_size_padded = pad_vocab_size(
283
self.org_vocab_size, self.padding_size
284
)
285
self.num_embeddings_padded = pad_vocab_size(
286
- self.org_vocab_size_padded + num_added_embeddings, self.padding_size
286
+ self.org_vocab_size_padded + self.num_added_embeddings, self.padding_size
287
)
288
assert self.org_vocab_size_padded <= self.num_embeddings_padded
289
@@ -296,6 +296,18 @@ class VocabParallelEmbedding(PluggableLayer):
296
self.tp_size,
297
)
298
self.embedding_dim = embedding_dim
299
+ # Divide the weight matrix along the vocabulary dimension.
300
+ self.num_embeddings_per_partition = divide(
301
+ self.num_embeddings_padded, self.tp_size
302
+ )
303
+
304
+ # Quantization methods share the same weight factory as linear layers,
305
+ # so setup standard linear metadata.
306
+ self.input_size = self.input_size_per_partition = embedding_dim
307
+ self.output_size = self.num_embeddings_padded
308
+ self.output_size_per_partition = self.num_embeddings_per_partition
309
+ self.output_partition_sizes = [self.output_size_per_partition]
310
+ self.prefix = prefix
311
312
# Avoid overriding a preselected model-specific method with generic
313
# config-based dispatch.
@@ -322,11 +334,6 @@ class VocabParallelEmbedding(PluggableLayer):
334
if params_dtype is None:
335
params_dtype = torch.get_default_dtype()
336
self.params_dtype = params_dtype
325
- # Divide the weight matrix along the vocabulary dimension.
326
- self.num_added_embeddings = self.num_embeddings - self.org_vocab_size
327
- self.num_embeddings_per_partition = divide(
328
- self.num_embeddings_padded, self.tp_size
329
- )
337
assert (
338
self.shard_indices.num_elements_padded == self.num_embeddings_per_partition
339
)
@@ -349,10 +356,10 @@ class VocabParallelEmbedding(PluggableLayer):
356
357
self.quant_method.create_weights(
358
self,
352
- self.embedding_dim,
353
- [self.num_embeddings_per_partition],
354
- self.embedding_dim,
355
- self.num_embeddings_padded,
359
+ self.input_size_per_partition,
360
+ self.output_partition_sizes,
361
+ self.input_size,
362
+ self.output_size,
363
params_dtype=params_dtype,
364
weight_loader=self.weight_loader,
365
)
@@ -600,6 +607,7 @@ class ParallelLMHead(VocabParallelEmbedding):
607
*,
608
disable_tp: bool = False,
609
):
610
+ self.has_bias = bias
611
super().__init__(
612
num_embeddings,
613
embedding_dim,