[NVBug 6108145] Fix PTQ calibration and export for fused-experts MoE (Qwen3.5-MoE VLM)#1340
Merged
[NVBug 6108145] Fix PTQ calibration and export for fused-experts MoE (Qwen3.5-MoE VLM)#1340
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds end-to-end support for HuggingFace fused-experts in PTQ: normalizes per-expert quantizer names for wildcard matching, forces eager experts implementation at runtime, discovers representative weight quantizers, moves fused-expert export earlier, updates a recipe, and adds tests covering calibration and export. Changes
Sequence Diagram(s)sequenceDiagram
participant Model as Model<br/><span style="background-color:rgba(52,152,219,0.5)">Model</span>
participant Plugin as Plugin<br/><span style="background-color:rgba(46,204,113,0.5)">force_eager_experts_impl_on_the_fly</span>
participant Calibrator as Calibrator<br/><span style="background-color:rgba(155,89,182,0.5)">PTQ Calibrator</span>
participant Exporter as Exporter<br/><span style="background-color:rgba(241,196,15,0.5)">Unified Exporter</span>
Model->>Plugin: detect fused-experts
Plugin-->>Model: set _experts_implementation = "eager" (recursively)
Model->>Calibrator: run calibration
Calibrator->>Model: invoke _QuantFusedExperts.forward (eager path)
Calibrator->>Calibrator: collect per-expert quantizer amax via hooks
Model->>Exporter: request export
Exporter->>Exporter: resolve representative_weight_quantizer for weight names
Exporter->>Exporter: detect fused-experts early
Exporter->>Exporter: call _export_fused_experts under fsdp2_aware_weight_update
Exporter-->>Caller: produce export metadata (includes fused-experts quant spec)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Contributor
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 649-657: The elif branch that checks for the same attribute
gate_up_proj_weight_quantizers is dead code because the preceding if block
handles that case and continues; remove the unreachable elif block (the second
check for gate_up_proj_weight_quantizers and its body) so only the initial
handling using fsdp2_aware_weight_update and _export_fused_experts(sub_module,
dtype) remains, leaving no duplicate checks for gate_up_proj_weight_quantizers
in the loop.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6c028304-1560-45eb-bcba-de04e7c03a20
📒 Files selected for processing (4)
examples/llm_ptq/hf_ptq.pymodelopt/torch/export/unified_export_hf.pymodelopt/torch/quantization/conversion.pymodelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1340 +/- ##
==========================================
+ Coverage 76.93% 76.96% +0.03%
==========================================
Files 471 471
Lines 50404 50482 +78
==========================================
+ Hits 38776 38855 +79
+ Misses 11628 11627 -1
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
meenchen
added
bug
cjluo-nv
reviewed
Apr 24, 2026
Collaborator
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
This is a well-structured bug fix with good test coverage. The 4-part fix (eager impl forcing, quantizer name normalization, export path hoisting, YAML layerwise change) is logically coherent and well-tested. However, there's a duplicate code path in unified_export_hf.py that should be cleaned up.
Design: Despite the complexity gate firing, this is a bug fix within existing systems, not an architectural change. No new abstractions are introduced.
Tests: Comprehensive — covers force_eager_experts_impl_on_the_fly edge cases, and an end-to-end calibration test that guards the full pipeline (name normalization → wildcard matching → amax collection). Good.
Issue: The new early-exit block in _process_quantized_modules makes the existing elif hasattr(sub_module, "gate_up_proj_weight_quantizers") block (deeper in the same function, inside the get_quantization_format() != QUANTIZATION_NONE guard) dead code. One of these should be removed.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
modelopt/torch/quantization/plugins/huggingface.py (1)
1441-1449: Harden recursive config traversal with a cycle guard.The recursive
_force()walk can loop forever on cyclic config graphs. Add a visited-set guard to make this robust.Proposed patch
nested_cfg_attrs = ("text_config", "vision_config", "audio_config", "speech_config") + visited_cfg_ids = set() def _force(cfg): if cfg is None: return + cfg_id = id(cfg) + if cfg_id in visited_cfg_ids: + return + visited_cfg_ids.add(cfg_id) if hasattr(cfg, "_experts_implementation"): cfg._experts_implementation = "eager" for sub in nested_cfg_attrs: if hasattr(cfg, sub): _force(getattr(cfg, sub))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelopt/torch/quantization/plugins/huggingface.py` around lines 1441 - 1449, The recursive helper _force can loop on cyclic config graphs; modify _force to accept and maintain a visited set (e.g., of object ids) and skip recursing into cfg instances already seen, so each cfg is processed only once. Specifically, update the _force signature to take an optional visited set, add the current cfg to visited (use id(cfg) or cfg itself), return early if already visited, and keep the existing behavior of setting cfg._experts_implementation = "eager" and iterating nested_cfg_attrs; ensure recursive calls pass the same visited set. This change hardens _force and prevents infinite recursion on cycles while still touching the same symbols (_force, nested_cfg_attrs, cfg._experts_implementation).tests/unit/torch/quantization/plugins/test_fused_experts.py (1)
385-441: Make registry cleanup exception-safe in the calibration test.If an assertion fails before the last line, the temporary registry entry may leak into subsequent tests.
Proposed patch
def test_calibration_populates_all_expert_quantizers(self): """After PTQ, every input/weight quantizer on the fused-experts module has amax set.""" import modelopt.torch.quantization as mtq model = _TinyMoEModel() expert_type = type(model.moe.experts) self._cleanup_registry(expert_type) - - quant_cfg = { + try: + quant_cfg = { "quant_cfg": [ {"quantizer_name": "*", "enable": False}, { "quantizer_name": "*gate_up_proj_input_quantizer", "cfg": {"num_bits": 8, "axis": None}, @@ for idx in range(NUM_EXPERTS): assert experts.gate_up_proj_weight_quantizers[idx].amax is not None, ( f"gate_up_proj_weight_quantizers[{idx}].amax is None — " "plural ModuleList name normalization in _match_quantizer likely broken." ) assert experts.down_proj_weight_quantizers[idx].amax is not None, ( f"down_proj_weight_quantizers[{idx}].amax is None." ) - - self._cleanup_registry(expert_type) + finally: + self._cleanup_registry(expert_type)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/torch/quantization/plugins/test_fused_experts.py` around lines 385 - 441, The test creates a temporary registry entry via expert_type and calls self._cleanup_registry(expert_type) at the end but can leak if an assertion fails; wrap the main test actions (quant_cfg setup, forward_loop, mtq.quantize, and all asserts) in a try/finally and move the final self._cleanup_registry(expert_type) into the finally block so cleanup always runs; keep expert_type assigned before the try and leave the initial cleanup call (before quantization) as-is.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@modelopt/torch/quantization/plugins/huggingface.py`:
- Around line 1441-1449: The recursive helper _force can loop on cyclic config
graphs; modify _force to accept and maintain a visited set (e.g., of object ids)
and skip recursing into cfg instances already seen, so each cfg is processed
only once. Specifically, update the _force signature to take an optional visited
set, add the current cfg to visited (use id(cfg) or cfg itself), return early if
already visited, and keep the existing behavior of setting
cfg._experts_implementation = "eager" and iterating nested_cfg_attrs; ensure
recursive calls pass the same visited set. This change hardens _force and
prevents infinite recursion on cycles while still touching the same symbols
(_force, nested_cfg_attrs, cfg._experts_implementation).
In `@tests/unit/torch/quantization/plugins/test_fused_experts.py`:
- Around line 385-441: The test creates a temporary registry entry via
expert_type and calls self._cleanup_registry(expert_type) at the end but can
leak if an assertion fails; wrap the main test actions (quant_cfg setup,
forward_loop, mtq.quantize, and all asserts) in a try/finally and move the final
self._cleanup_registry(expert_type) into the finally block so cleanup always
runs; keep expert_type assigned before the try and leave the initial cleanup
call (before quantization) as-is.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 310a1d2b-745c-4df5-8317-038bf82c199f
📒 Files selected for processing (2)
modelopt/torch/quantization/plugins/huggingface.pytests/unit/torch/quantization/plugins/test_fused_experts.py
cjluo-nv
reviewed
Apr 24, 2026
cjluo-nv
reviewed
Apr 24, 2026
cjluo-nv
reviewed
Apr 24, 2026
cjluo-nv
reviewed
Apr 24, 2026
Collaborator
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
All critical previous review comments have been addressed: dead code removed, import moved to top, control flow restructured, and singular/plural quantizer naming covered. The PR is a well-structured 4-part bug fix (eager impl forcing, quantizer name normalization, export path hoisting, YAML layerwise change) with comprehensive tests. No new abstractions — purely surgical fixes within existing systems. Code is correct and clean.
Complex PR: spans 6 directories (≥ 5). Looping in a human for approval.
cjluo-nv
approved these changes
Apr 27, 2026
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@modelopt/torch/quantization/utils/core_utils.py`:
- Around line 232-257: weight_attr_names now yields plural fused-experts names
but consumers still assume a singular "<name>_weight_quantizer" attr; update
those consumers (e.g. the getattr call in quant_module.py and the logic in
model_calib.py) to use representative_weight_quantizer(module, name) to fetch
the correct quantizer (handles singular, plural ModuleList, or None) or
explicitly check for "<name>_weight_quantizers" and handle ModuleList entries,
ensuring you avoid bare getattr(...) without a default and preserve existing
behavior for singular quantizers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 376b95ff-4015-46e6-81e9-e9ae47a24a05
📒 Files selected for processing (2)
modelopt/torch/quantization/plugins/huggingface.pymodelopt/torch/quantization/utils/core_utils.py
sychen52
approved these changes
Apr 27, 2026
meenchen
force-pushed
the
weimingc/fix_qwen36_moe_ptq
branch
from
c1242ce to
3eb0f21
Compare
Contributor
Author
|
/ok to test 3eb0f21 |
meenchen
force-pushed
the
weimingc/fix_qwen36_moe_ptq
branch
from
e65e48e to
ce3e0e0
Compare
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
…standard export branches Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
The base QuantModule.fold_weight only walks singular *_weight_quantizer attributes, so per-expert quantizers in _QuantFusedExperts' gate_up_proj_weight_quantizers / down_proj_weight_quantizers ModuleLists are never folded, leaving _amax behind. The vLLM fake-quant export test test_hf_vllm_export_tiny_qwen3_moe[FP8] surfaces this. Override fold_weight on _QuantFusedExperts to walk the per-expert ModuleList, apply each quantizer to its 3-D slice, disable, and drop _amax / _pre_quant_scale. Also add the blank line between docstring intro and bullet list in representative_weight_quantizer / weight_attr_names so sphinx stops emitting "Block quote ends without a blank line" warnings during build-docs. Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
meenchen
force-pushed
the
weimingc/fix_qwen36_moe_ptq
branch
2 times, most recently
from
58b897a to
29fdcec
Compare
The vllm-fakequant export's _fakequant_module_weights iterated
module.named_children() looking for attr_name.endswith("weight_quantizer")
plus isinstance(quantizer, TensorQuantizer). Both filters skip the plural
*_weight_quantizers ModuleList that _QuantFusedExperts uses for per-expert
quantizers. As a result, fused MoE experts had their fused 3-D
gate_up_proj / down_proj saved un-fake-quantized while the rest of the
pipeline (fold_weight, modelopt state) treated them as fake-quanted ->
weight mismatch on round-trip and the gpu test
test_hf_vllm_export_tiny_qwen3_moe[FP8_DEFAULT_CFG] failed by 0.00244
(the FP8 quant noise that should have been zero).
Fix:
- Add _fakequant_fused_experts_weights helper that walks the per-expert
ModuleList and applies each TensorQuantizer to its slice of the 3-D
fused weight, writing into state_dict (non-mutating export) or
in-place on the parameter, mirroring _QuantFusedExperts.fold_weight.
- Extend _WEIGHT_QUANTIZER_STATE_KEY regex to match the plural form
weight_quantizers? so per-expert quantizer state (_amax etc.) is
recognized as weight-quantizer state and stripped before save -
otherwise the saved modelopt_state_weights would carry per-expert
amax tensors and trip the test's "weight quantizer ... should have
empty state after fold" invariant.
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
meenchen
force-pushed
the
weimingc/fix_qwen36_moe_ptq
branch
from
29fdcec to
fda1e20
Compare
Contributor
Author
|
/ok to test fda1e20 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: Bug fix
Fixes a 4-bug cascade that caused silent PTQ failure on Qwen3.5-MoE VLMs (Qwen3.6-35B-A3B): calibration
appeared to succeed but produced token-salad at inference. Root cause: HF's @use_experts_implementation
dispatches expert forward to torch._grouped_mm / torch.bmm, bypassing the F.linear hook that captures
activations — so gate_up_proj_input_quantizer / down_proj_input_quantizer never calibrated and no input_scale
tensors were emitted.
Changes:
vision_config / …) so per-expert F.linear calls are visible to the calibration hook.
→ weight_quantizer) before fnmatch, so wildcards like mlp.expertsweight_quantizer match fused-expert
quantizers.
get_quantization_format() gate so _export_fused_experts() runs even when the top-level format query returns
QUANTIZATION_NONE (happens for experts-only recipes).
breaks the layerwise walker).
Usage
Testing
Testing
End-to-end PTQ → vLLM deploy → NEL eval on Qwen3.6-35B-A3B (256 experts × 40 layers, 35B params):
Hook-call diagnostic: 0 → 6720 per-expert F.linear calls during calibration after the fix; 0 → 30720
input_scale tensors emitted in the exported checkpoint.
FP8 fused-MoE path still produces gibberish — separate follow-up (vLLM per-expert weight_scale handling).
vLLM full-FP8: the FlashInfer TRTLLM Fp8MoE loader doesn't stack the 256 per-expert scalar weight_scale tensors
into a [num_experts] per-expert vector — it ends up applying one expert's scale across all 256, so every
routed expert dequants with the wrong amplitude → coherent token stream collapses into multilingual gibberish.
SGLang full-FP8: qwen3_5.py::_make_packed_weight_loader rejects with AssertionError: Unexpected scalar for
tuple shard load: loaded_shard_id=(0,1,2), split_sizes=[1,1,1] — its packed-loader has no path for "N
independent per-tensor source scalars combining into one fused-shard parameter," so the fused QKV (or
in_proj_qkvz) load is structurally refused and the model never finishes loading.
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅ / ❌ / N/AAdditional Information
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests