[Qwen4] validate FP8 PLE weight scale after loading (#54722)
Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
Harry Huang committed
Sep 2, 2026 at 12:30 UTC
1e300895acecd29fa037616b59d793aae164c4ea
2 files changed
+26
-5
tests/models/qwen4_exp/test_ple.py
+19
-4
@@ -72,7 +72,7 @@ def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding:
72
)
73
embedding.register_parameter(
74
"weight_scale",
75
- nn.Parameter(torch.zeros(1, dtype=torch.bfloat16), requires_grad=False),
75
+ nn.Parameter(torch.zeros(1, dtype=torch.float32), requires_grad=False),
76
)
77
_set_test_embedding_weight_loader(embedding)
78
module.ngram_embedding = embedding
@@ -156,11 +156,17 @@ def test_ngram_embedding_loads_fp8_shards_and_global_scale() -> None:
156
module.ngram_embedding.weight.float(),
157
torch.cat((shard_0[2:4], shard_1[0:2])).float(),
158
)
159
- assert torch.equal(module.ngram_embedding.weight_scale, weight_scale)
159
+ assert module.ngram_embedding.weight_scale.dtype == torch.float32
160
+ torch.testing.assert_close(
161
+ module.ngram_embedding.weight_scale,
162
+ weight_scale.float(),
163
+ )
164
165
166
def _make_fp8_embedding_layer(
167
monkeypatch: pytest.MonkeyPatch,
168
+ *,
169
+ load_scale: bool = True,
170
) -> embedding_module.VocabParallelEmbedding:
171
monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
172
monkeypatch.setattr(
@@ -180,12 +186,14 @@ def _make_fp8_embedding_layer(
186
)
187
weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]])
188
layer.weight.data.copy_(weight.to(torch.float8_e4m3fn))
183
- layer.weight_scale.data.copy_(torch.tensor([0.25], dtype=torch.bfloat16))
189
+ if load_scale:
190
+ layer.weight_scale.data.copy_(torch.tensor([0.25], dtype=torch.bfloat16))
191
return layer
192
193
194
def test_ple_fp8_embedding_dequantizes_in_ple_layer(monkeypatch) -> None:
195
layer = _make_fp8_embedding_layer(monkeypatch)
196
+ layer.quant_method.process_weights_after_loading(layer)
197
quantized_output = layer(torch.tensor([2, 0]))
198
ple_layer = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
199
nn.Module.__init__(ple_layer)
@@ -198,13 +206,20 @@ def test_ple_fp8_embedding_dequantizes_in_ple_layer(monkeypatch) -> None:
206
)
207
208
assert layer.weight.dtype == torch.float8_e4m3fn
201
- assert layer.weight_scale.dtype == torch.bfloat16
209
+ assert layer.weight_scale.dtype == torch.float32
210
assert quantized_output.dtype == torch.float8_e4m3fn
211
assert output.dtype == torch.bfloat16
212
weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]])
213
torch.testing.assert_close(output, (weight[[2, 0]] * 0.25).bfloat16())
214
215
216
+def test_ple_fp8_embedding_rejects_missing_global_scale(monkeypatch) -> None:
217
+ layer = _make_fp8_embedding_layer(monkeypatch, load_scale=False)
218
+
219
+ with pytest.raises(ValueError, match="missing its global scale"):
220
+ layer.quant_method.process_weights_after_loading(layer)
221
+
222
+
223
def test_ple_fp8_embedding_uses_int8_for_tp_reduce(monkeypatch) -> None:
224
monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
225
monkeypatch.setattr(
vllm/models/qwen4_exp/nvidia/ple_layer.py
+7
-1
@@ -106,10 +106,16 @@ class Qwen4ExpPLEFp8EmbeddingMethod(QuantizeMethodBase):
106
input_size_per_partition,
107
None,
108
weight_loader,
109
- scale_dtype=torch.bfloat16,
109
+ scale_dtype=torch.float32,
110
)
111
layer.register_parameter("weight_scale", weight_scale)
112
113
+ def process_weights_after_loading(self, layer: nn.Module) -> None:
114
+ """Reject FP8 PLE checkpoints without a global scale."""
115
+ sentinel = torch.finfo(torch.float32).min
116
+ if torch.any(layer.weight_scale == sentinel):
117
+ raise ValueError("FP8 PLE checkpoint is missing its global scale")
118
+
119
def apply(
120
self,
121
layer: nn.Module,