[Bugfix][Model] Fix FP8 PLE loading in mixed ModelOpt checkpoints (#54882)
Signed-off-by: Shiyang Chen <shiychen@nvidia.com> Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
sychen52 committed
Sep 3, 2026 at 10:09 UTC
d4d703caf908786416585ceb1f369e2e0363358b
2 files changed
+42
tests/models/qwen4_exp/test_ple.py
+34
@@ -14,6 +14,9 @@ from torch.nn import functional as F
14
import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module
15
import vllm.model_executor.parameter as parameter_module
16
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
17
+from vllm.model_executor.layers.quantization.modelopt import (
18
+ ModelOptMixedPrecisionConfig,
19
+)
20
from vllm.models.qwen4_exp.common.ple import (
21
PLEShardOverlap,
22
compute_ple_shard_overlap,
@@ -288,6 +291,37 @@ def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None:
291
assert _get_ple_embedding_quant_method(quant_config, prefix) is None
292
293
294
+def test_ple_fp8_embedding_supports_mixed_precision_config() -> None:
295
+ prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"
296
+ quant_config = ModelOptMixedPrecisionConfig.from_config(
297
+ {
298
+ "quantization": {
299
+ "quant_algo": "MIXED_PRECISION",
300
+ "exclude_modules": [],
301
+ "group_size": 16,
302
+ "quantized_layers": {
303
+ prefix: {"quant_algo": "FP8"},
304
+ "model.language_model.layers.2.moe.gate_proj": {
305
+ "quant_algo": "NVFP4"
306
+ },
307
+ },
308
+ }
309
+ }
310
+ )
311
+
312
+ assert isinstance(
313
+ _get_ple_embedding_quant_method(quant_config, prefix),
314
+ Qwen4ExpPLEFp8EmbeddingMethod,
315
+ )
316
+ assert (
317
+ _get_ple_embedding_quant_method(
318
+ quant_config,
319
+ "model.language_model.layers.2.moe.gate_proj",
320
+ )
321
+ is None
322
+ )
323
+
324
+
325
def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
326
conv_state_len = 6
327
dilation = 2
vllm/models/qwen4_exp/nvidia/ple_layer.py
+8
@@ -22,6 +22,9 @@ from vllm.model_executor.layers.quantization.base_config import (
22
QuantizeMethodBase,
23
)
24
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
25
+from vllm.model_executor.layers.quantization.modelopt import (
26
+ ModelOptMixedPrecisionConfig,
27
+)
28
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
29
create_fp8_scale_parameter,
30
create_fp8_weight_parameter,
@@ -133,6 +136,11 @@ def _get_ple_embedding_quant_method(
136
) -> QuantizeMethodBase | None:
137
"""Select global-scale FP8 only for quantized PLE checkpoint shards."""
138
139
+ if isinstance(quant_config, ModelOptMixedPrecisionConfig):
140
+ if quant_config._resolve_quant_algo(prefix) == "FP8":
141
+ return Qwen4ExpPLEFp8EmbeddingMethod()
142
+ return None
143
+
144
if not isinstance(quant_config, Fp8Config):
145
return None
146
if not quant_config.is_checkpoint_fp8_serialized: