[Bugfix][Humming] Preserve ModelOpt FP8 weight dimensions (#51093)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
Netanel Haber committed
Aug 6, 2026 at 00:08 UTC
482c6134236d96ebab3c749cd02d245516adad1d
2 files changed
+25
tests/quantization/test_modelopt.py
+23
@@ -22,6 +22,7 @@ from vllm.model_executor.kernels.linear import (
22
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
23
from vllm.model_executor.layers.quantization.modelopt import (
24
ModelOptFp8Config,
25
+ ModelOptFp8LinearMethod,
26
ModelOptMixedPrecisionConfig,
27
ModelOptMxFp8Config,
28
ModelOptNvFp4Config,
@@ -114,6 +115,28 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head():
115
assert isinstance(method, ModelOptNvFp4LinearMethod)
116
117
118
+def test_modelopt_fp8_updates_weight_dims_after_transpose():
119
+ layer = torch.nn.Module()
120
+ layer.register_parameter(
121
+ "weight", torch.nn.Parameter(torch.empty(3, 2), requires_grad=False)
122
+ )
123
+ layer.register_parameter(
124
+ "weight_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False)
125
+ )
126
+ layer.register_parameter(
127
+ "input_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False)
128
+ )
129
+
130
+ method = ModelOptFp8LinearMethod.__new__(ModelOptFp8LinearMethod)
131
+ method.fp8_linear = Mock()
132
+ method.process_weights_after_loading(layer)
133
+
134
+ assert layer.weight.shape == (2, 3)
135
+ assert layer.weight.input_dim == 0
136
+ assert layer.weight.output_dim == 1
137
+ method.fp8_linear.process_weights_after_loading.assert_called_once_with(layer)
138
+
139
+
140
def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized():
141
config = ModelOptNvFp4Config(
142
is_checkpoint_nvfp4_serialized=True,
vllm/model_executor/layers/quantization/modelopt.py
+2
@@ -519,6 +519,8 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
519
layer.weight, layer.weight_scale, layer.logical_widths
520
)
521
layer.weight = Parameter(weight.t(), requires_grad=False)
522
+ layer.weight.input_dim = 0
523
+ layer.weight.output_dim = 1
524
layer.weight_scale = Parameter(max_w_scale, requires_grad=False)
525
layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
526
self.fp8_linear.process_weights_after_loading(layer)