[Bugfix][Model] Fix FP8 PLE loading in mixed ModelOpt checkpoints - #54882
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
This pull request has merge conflicts that must be resolved before it can be |
|
✅ @sychen52, CI is now available for this PR.
|
Select the FP8 PLE embedding method from per-layer mixed-precision metadata so NVFP4 checkpoints carrying FP8 PLE weights register and load the global weight scale. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
ae01e57 to
bac9c8e
Compare
|
/ci run |
|
✅ Triggered Buildkite CI #86902 for commit |
sychen52
changed the title
|
/ci retry |
|
✅ Queued 2 failed job(s) for retry in Buildkite CI #86902. |
|
/ci retry |
|
✅ Triggered Buildkite CI #86968 for commit |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📝 WalkthroughWalkthroughThe PLE embedding quant-method selector now supports ModelOpt mixed-precision configurations. Tests verify FP8 selection for the PLE prefix and no method selection for an NVFP4 prefix. ChangesQwen4Exp PLE mixed-precision support
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to The change enables FP8 PLE weights to load correctly from mixed-precision Qwen checkpoints while preserving existing behavior for other layer assignments. No actionable merge-blocking risk remains beyond normal checks and review. Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
/ci run |
|
✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/86968 |
|
/ci retry |
|
✅ Queued 3 failed job(s) for retry in Buildkite CI #86968. |
|
/ci retry |
|
✅ Triggered Buildkite CI #87046 for commit |
|
/ci run |
|
✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/87046 |
…ors target group Mixed checkpoints (INT4 experts via compressed-tensors) can carry the PLE n-gram table as FP8-e4m3 with one global scale, the same layout the official FP8 checkpoint ships and that PLE CPU offload keeps resident at 1 byte/param (~51 GB instead of ~102 GB host RAM for Qwen3.8-Flash-Next). Mirrors vllm-project#54882 for ModelOptMixedPrecisionConfig: a config group with an FP8 per-tensor weight scheme whose target matches the embedding prefix selects Qwen4ExpPLEFp8EmbeddingMethod; anything else keeps the BF16 embedding.
|
This fixed the After this merge, if isinstance(quant_config, ModelOptMixedPrecisionConfig):
if quant_config._resolve_quant_algo(prefix) == "FP8":
return Qwen4ExpPLEFp8EmbeddingMethod()
return NoneIf a mixed-precision checkpoint keeps its FP8 PLE table in the exclude / The if isinstance(quant_config, ModelOptMixedPrecisionConfig):
if quant_config._resolve_quant_algo(prefix) == "FP8":
return Qwen4ExpPLEFp8EmbeddingMethod()
# PLE kept in the exclude list rather than enumerated: trust the
# checkpoint's declared ple_embedding_dtype.
if quant_config.is_layer_excluded(prefix) and _ple_dtype_is_fp8(ple_embedding_dtype):
return Qwen4ExpPLEFp8EmbeddingMethod()
return NoneHappy to send this as a small follow-up PR with a test for the exclude-list convention if it's wanted. Found with AI assistance (Claude Code) while extending a Flash-Next checkpoint to mixed-precision; the fall-through was verified by reading the merged |
… experts in ModelOpt mixed checkpoints `nvidia/Qwen3.8-Flash-Next-NVFP4` is a three-precision MIXED_PRECISION checkpoint: NVFP4 routed experts, an FP8 PLE n-gram table (F8_E4M3 shards + a scalar weight_scale) and block-FP8 MTP experts. Two decisions were made before consulting the checkpoint's per-layer `quantized_layers` map and broke on it: 1. PLE table dtype. The table's storage dtype is fixed at construction and only switched to fp8 for a whole-checkpoint `fp8` quant name or an explicit `text_config.ple_embedding_dtype`. This checkpoint sets neither (RadixArk's NVFP4 repo sets the config key; nvidia's does not), so the table was built bf16, moved into pinned host memory by the default `--ple-offload-embedding`, and load_weights then raised "fp8 PLE auto-switch is unsupported with ple_offload_embedding". `_ple_table_is_fp8()` now also asks the mixed config's `_resolve_quant_algo()` for the embedding's own prefix, so the table is fp8 from the start (mirrors vllm-project/vllm#54882). 2. Draft quantization. `_mtp_quant_config` dropped the quant config for every modelopt_mixed checkpoint on the assumption that embedded MTP weights are bf16. Here `mtp.layers.0.mlp.experts` is block-FP8, so the draft's FusedMoE was built unquantized, the fp8 expert values were cast to bf16 without their scales and `weight_scale_inv` was silently skipped: the server ran and target accuracy was fine, but MTP accept length fell to ~1.6. The draft now keeps the quant config when the map lists any `mtp.` layer; every other draft module still resolves to bf16 through the same per-layer lookup. Verified on 4x B300 with the user launch command (TP4, NEXTN 3/1/4, no --quantization flag): loads, sgl-eval GSM8K (200, thinking) 0.97, accept length ~3.2-3.6, draft weights 1.53 GB (was 2.47 GB bf16). Depends on the block-FP8 FusedMoE dispatch in ModelOptMixedPrecisionConfig. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
3 participants
Select the FP8 PLE embedding method from per-layer mixed-precision metadata so NVFP4 checkpoints carrying FP8 PLE weights register and load the global weight scale.
Purpose
Enable Qwen3.8-Flash-Next-NVFP4 checkpoint
Test Plan
added unittest
tested by running checkpoint
Test Result
passed
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Summary by CodeRabbit
New Features
Tests