Skip weight quantizers with no stored weight in fold_weight

[babyplutokurt]

Skip weight quantizers with no stored weight in fold_weight

What does this PR do?

Type of change: Bug fix

Fixes #2131.

On Megatron-Core models with tied word embeddings
(share_embeddings_and_output_weights=True), the output_layer is built with
skip_weight_param_allocation: it stores weight = None and borrows the embedding
weight at forward time, while ModelOpt still attaches a weight_quantizer to it.
QuantModule.fold_weight matches the pair on the *_weight_quantizer attribute name
and fake_quant alone, then dereferences weight.data, so mtq.fold_weight crashes:

modelopt/torch/quantization/nn/modules/quant_module.py:145, in _fold_weight_quantizer
    weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype))
AttributeError: 'NoneType' object has no attribute 'data'

This PR skips pairs whose weight attribute is not a tensor and leaves their quantizer
untouched: there is nothing stored to fold, and the shared weight keeps being
quantized at forward time through the still-enabled quantizer (the embedding module
folds its own stored weight as before). HF models are unaffected — they express tying
as a shared tensor rather than None, which is why the crash only surfaced on
Megatron.

Usage

import modelopt.torch.quantization as mtq
from megatron.core.models.gpt import GPTModel

# Any GPTModel with tied embeddings, e.g. Qwen3-0.6B:
# output_layer.weight is None (borrowed from the embedding at forward time).
model = GPTModel(..., share_embeddings_and_output_weights=True)
model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop)

mtq.fold_weight(model)  # previously: AttributeError; now folds all stored weights

Testing

• New unit test test_fold_weight_skips_none_weight in
  tests/unit/torch/quantization/test_tensor_quant_cpu.py builds the tied-layer shape
  (a converted QuantLinear with weight = None) — it reproduces the exact
  AttributeError at quant_module.py:145 without the fix and passes with it,
  asserting the skipped quantizer stays enabled with its _amax intact.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ (previously-crashing calls now succeed; folding behavior for stored weights is unchanged)
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: ❌ (external contribution — cannot self-trigger /claude review)

Additional Information

Related issue: #2131. Downstream workaround that this fix would let us drop:
NVIDIA-NeMo/RL#3441 (nemo_rl/modelopt/models/policy/workers/weight_folding.py).

Summary by CodeRabbit

• Bug Fixes

  • Fixed weight folding crashes for quantized models with missing or non-tensor weights.
  • Such weights are now safely skipped without changing their associated quantizer settings.
  • Normal tensor weights continue to fold as expected.

• Tests

  • Added regression coverage to verify quantizers and calibration values remain unchanged when weights are unavailable.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4fa7c4a3-c74c-48b7-890d-c57854417350

📥 Commits

Reviewing files that changed from the base of the PR and between 44e0d34 and 8606f41.

📒 Files selected for processing (2)

• CHANGELOG.rst
• modelopt/torch/quantization/nn/modules/quant_module.py

🚧 Files skipped from review as they are similar to previous changes (2)

• modelopt/torch/quantization/nn/modules/quant_module.py
• CHANGELOG.rst

---

📝 Walkthrough

Walkthrough

QuantModule.fold_weight now skips quantized modules whose stored weight is not a tensor. The change prevents crashes for tied Megatron-Core embeddings and preserves the associated quantizer state. A regression test covers the None weight case.

Changes

Weight folding safety

Layer / File(s)	Summary
Skip non-tensor weights during folding modelopt/torch/quantization/nn/modules/quant_module.py, tests/unit/torch/quantization/test_tensor_quant_cpu.py, CHANGELOG.rst	fold_weight skips non-tensor weights without changing their quantizers. The regression test verifies that calibration state remains intact. The changelog documents the fix.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• NVIDIA/Model-Optimizer#2112: Both PRs address Megatron quantization edge cases involving tied or untied output embeddings, but they modify different code paths.

Suggested reviewers: sychen52, jenchen13, juhi10071998

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the main change: skipping weight quantizers when no stored weight exists during fold_weight.
Linked Issues check	✅ Passed	The changes fix issue #2131 by skipping non-tensor weights, preserving the quantizer, and preventing tied-embedding fold_weight crashes.
Out of Scope Changes check	✅ Passed	The code and regression test directly support issue #2131 and the stated fold_weight bug fix; no unrelated changes are evident.
Security Anti-Patterns	✅ Passed	The PR adds only a tensor guard, documentation, changelog, and a test; no prohibited loading, pickle, remote-code, eval/exec, nosec, or dependency changes were added.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Warning

CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.

Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.

👉 Steps to fix this

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unit/torch/quantization/test_tensor_quant_cpu.py`:
- Around line 344-350: Update the test around qlinear.fold_weight() to snapshot
qlinear.weight_quantizer.amax before folding, then assert afterward that the
calibration value is unchanged. Keep the existing enabled-state assertion and
use the captured value to verify the skipped quantizer retains its data.

🪄 Autofix

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: 88c010b9-7b55-4229-a416-950f55267b79

📥 Commits

Reviewing files that changed from the base of the PR and between 6b02f52 and df9220d.

📒 Files selected for processing (3)

• CHANGELOG.rst
• modelopt/torch/quantization/nn/modules/quant_module.py
• tests/unit/torch/quantization/test_tensor_quant_cpu.py

[cjluo-nv]

Bot review (gpt-5.6-sol) — DM the bot to share feedback.

Straightforward, well-scoped fix. fold_weight now skips missing/non-tensor stored weights before dereferencing them, preserving the forward-time quantizer needed by Megatron tied-output layers. The regression test reproduces the weight=None shape and verifies the quantizer and calibration state remain intact, and the changelog entry is under the current development release.

[kevalmorabia97]

/ok to test 44e0d34

[kevalmorabia97]

/claude review

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.14%. Comparing base (c6889c6) to head (44e0d34).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2132      +/-   ##
==========================================
- Coverage   78.73%   78.14%   -0.60%     
==========================================
  Files         522      522              
  Lines       60342    60344       +2     
==========================================
- Hits        47508    47153     -355     
- Misses      12834    13191     +357     

Flag	Coverage Δ
examples-diffusers	20.80% <0.00%> (-0.01%)	⬇️
examples-gpt-oss	13.28% <0.00%> (-0.01%)	⬇️
examples-hf_ptq	21.48% <0.00%> (-0.04%)	⬇️
examples-llm_distill	13.34% <0.00%> (-0.01%)	⬇️
examples-llm_eval	17.10% <0.00%> (-0.01%)	⬇️
examples-llm_qat	16.65% <0.00%> (-0.01%)	⬇️
examples-llm_sparsity	15.93% <0.00%> (-0.01%)	⬇️
examples-megatron_bridge	25.69% <0.00%> (-0.15%)	⬇️
examples-specdec_bench	13.02% <0.00%> (-0.01%)	⬇️
examples-speculative_decoding	17.53% <0.00%> (-0.08%)	⬇️
examples-torch_onnx	21.89% <0.00%> (-0.01%)	⬇️
examples-torch_trt	15.10% <0.00%> (-0.01%)	⬇️
gpu	58.60% <50.00%> (-0.71%)	⬇️
regression	14.91% <0.00%> (+0.06%)	⬆️
unit	55.30% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[claude]

[SUGGESTION] The skip is the right call, but it introduces the first documented exception to an invariant that two other places in the tree assert on: after fold_weight, no weight quantizer is enabled.

• examples/vllm_serve/fakequant_worker.py:146-150 raises RuntimeError("... is still enabled after fold_weight — double-quantization would corrupt activations.") for any still-enabled weight-quantizer state key.
• modelopt/torch/export/plugins/vllm_fakequant_hf.py:128-140 (_check_all_weight_quantizers_disabled) raises before writing metadata.

I traced this and believe it is currently unreachable: both are HF/vLLM-side paths, and vLLM's ParallelLMHead ties by sharing the embedding tensor (allocated weight), not by weight = None, so neither will see a skipped pair. Megatron's output_layer never flows through those checkers. So this is not a blocking issue.

Still, worth a line in the fold_weight docstring exception you added — something like "such quantizers stay enabled, so callers that assert all weight quantizers are disabled after folding must special-case them" — so the next person hitting one of those RuntimeErrors can find the reason here instead of re-deriving it. A pointer in the inline comment to _check_all_weight_quantizers_disabled would work equally well.

[claude]

Claude review summary

Scope: full review (trigger comment carried no scoping instructions). The PR touches 3 files — modelopt/torch/quantization/nn/modules/quant_module.py, tests/unit/torch/quantization/test_tensor_quant_cpu.py, CHANGELOG.rst — all reviewed. (Note: a two-dot diff against the base tip also surfaces unrelated attention_sparsity/plugins/vllm.py churn from base movement; that is not part of this PR and was excluded.)

Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 1

What I verified

• The guard is at the right place. isinstance(weight, torch.Tensor) sits between the hasattr assert and _fold_weight_quantizer, i.e. exactly where quant_module.py:145 dereferences weight.data. It admits every legitimate weight representation in the tree — nn.Parameter, DTensor (HF TP, transformers<5.0), and QTensorWrapper (real-quant packed weights, itself an nn.Parameter) all remain torch.Tensor instances, so no currently-folding module starts getting skipped. continue (rather than falling through) correctly leaves disable() / disable_rotate() / attr-drop unrun, which is what keeps the borrowed weight quantized at forward time.
• No double-quantization on the Megatron path. Megatron's VocabParallelEmbedding is not registered in the quantization QuantModuleRegistry (only in modelopt/torch/nas/plugins/megatron.py), so the shared embedding table is never folded on that side. The tied weight therefore gets quantized exactly once — at the still-enabled output_layer.weight_quantizer during forward. No leftover-state or double-fold hazard.
• Sibling fold paths are unaffected. _QuantHFParallelLinear.fold_weight (huggingface.py:514) enters enable_weight_access_and_writeback, whose isinstance(self.weight, DTensor) check is simply False for None and falls to the pass-through branch. _QuantFusedExperts.fold_weight (huggingface.py:1101) and _QuantFusedMoEBase.fold_weight (plugins/vllm.py:673) bypass the base loop entirely and already guard weight is None where relevant.
• Save/restore round-trip is self-consistent. After folding, the skipped output_layer quantizer stays enabled with _amax intact, so quant_module_get_extra_state (plugins/megatron.py:132) serializes its quantizer state instead of the {} its folded siblings emit — which matches what restore will re-apply. Prior to this PR the call crashed outright, so there is no checkpoint-compat regression.
• The test genuinely reproduces the bug. QuantModuleRegistry.convert(nn.Linear) yields a module using the base QuantModule.fold_weight (the only override in quant_linear.py lives on SVDQuantLinear), the default weight_quantizer is enabled with fake_quant=True, and register_parameter("weight", None) keeps hasattr(self, "weight") true — so the loop reaches the dereference and fails on main. Both post-conditions (still enabled, _amax preserved) are the right invariants to pin.

I did not execute the test suite in this environment (the sandbox declined the pytest invocation); correctness above is from reading the code paths, not from a local run.

Suggestion (non-blocking)

One inline comment on quant_module.py: the skip creates the first documented exception to the "after fold_weight, no weight quantizer is enabled" invariant that two other sites assert on — examples/vllm_serve/fakequant_worker.py:146 and _check_all_weight_quantizers_disabled in modelopt/torch/export/plugins/vllm_fakequant_hf.py:128. I traced both and believe they are unreachable for this case (vLLM's ParallelLMHead ties by sharing the allocated tensor, not by weight = None, and Megatron's output_layer never flows through either checker), so it is not a blocker — but a sentence in the docstring pointing at those assertions would save the next person hitting one of those RuntimeErrors from re-deriving why.

Separately, and outside this diff: SVDQuantLinear.fold_weight (quant_linear.py:168-181) guards with hasattr(self, "weight"), which is truthy for weight = None, then dereferences self.weight.data. It is unreachable today (SVDQuant calibration would have to produce svdquant_lora_a/b for a module with no stored weight, which would itself crash first), so no change is warranted here — just noting it as the same shape of latent issue in case that class ever meets a borrowed-weight module.

Risk assessment

Low. Nine added lines, purely additive, guarding a dereference that previously raised unconditionally. The behavior change is confined to modules whose stored weight is not a tensor — a set that was 100% crashing before — so every path that worked prior to this PR is bit-identical after it. Well-scoped, correctly narrow, and matched by a regression test that fails without the fix.

[claude]

Claude review passed — no blocking issues found. LGTM

[babyplutokurt]

Added the docstring suggestion from claude.

Also rebased on latest main already and re-push

@kevalmorabia97 ^

[babyplutokurt]

Added the docstring suggestion from claude.

Also rebased on latest main already and re-push

@kevalmorabia97 ^