Bug fix: 6542481 - #2064
Conversation
|
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:
📝 WalkthroughWalkthroughThe QLoRA export flow conditionally restores ModelOpt state, resolves reparented quantized modules, normalizes exported keys, preserves NVFP4 scales, and validates base-model and adapter layouts. ChangesQLoRA export flow
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ExportScript
participant ModeloptStateManager
participant QLoRAModel
participant ExportedCheckpoint
ExportScript->>ModeloptStateManager: Check converted model state
ModeloptStateManager->>QLoRAModel: Restore state when required
QLoRAModel->>ExportedCheckpoint: Export normalized base and adapter state
ExportedCheckpoint->>ExportedCheckpoint: Preserve NVFP4 scale tensors
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2064 +/- ##
===========================================
+ Coverage 62.19% 78.07% +15.87%
===========================================
Files 521 521
Lines 59857 59864 +7
===========================================
+ Hits 37230 46739 +9509
+ Misses 22627 13125 -9502
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
d53322b to
c610479
Compare
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Small, well-scoped fix for the QLoRA export path; the three sub-fixes each look correct (verified postprocess_state_dict's endswith matching is order-safe for weight_scale vs weight_scale_2, and that modelopt_state_train.pth is indeed written into the QLoRA output dir by QATTrainer.__init__, so the new is_converted() guard is the one doing the work). A few things worth addressing before merge:
postprocess_state_dictsilently drops anybase_layer.*key that isn't in the hand-maintained rename map — the exact failure mode being fixed here.base_layer.bias(e.g. Qwen2-style qkv biases) and the AWQbase_layer.pre_quant_scale(produced by the earlierinput_quantizer._pre_quant_scalerename, which leavesbase_layerin the key) are still mishandled. A generic strip of.base_layer— or at least a warning/assert on unmatchedbase_layerkeys — would stop this from recurring.- The library-side fix in
modelopt/torch/opt/plugins/transformers.py(.base_layerfallback in_restore_qtensor_wrappers) has no unit test; it's only covered by the heavy GPU example test. The added unit test covers the export rename, not this path. - Minor: in
examples/llm_qat/export.py,modelopt_weightsis popped outside the branch that uses it, and a LoRA checkpoint that is neither converted nor hasmodelopt_state_train.pthnow silently exports unquantized instead of failing loudly.
No licensing concerns and no prompt-injection content in the PR metadata.
| "base_layer.input_scale": "input_scale", | ||
| "base_layer.weight_scale": "weight_scale", | ||
| # NVFP4 global scale; the exported model cannot be dequantized without it. | ||
| "base_layer.weight_scale_2": "weight_scale_2", |
There was a problem hiding this comment.
Bot comment.
This fixes weight_scale_2, but the underlying mechanism is still fragile: skip_keys contains "base_layer", so any base_layer.* key that doesn't end with one of the four enumerated suffixes is dropped from the exported state dict with no warning. Two concrete cases remain broken:
...base_layer.bias— dropped entirely (matters for architectures with qkv/mlp biases, e.g. Qwen2; the tiny-Qwen3 e2e test won't catch it).- AWQ:
...base_layer.input_quantizer._pre_quant_scalematches the earlierinput_quantizer._pre_quant_scalereplacement (checked first in dict order), producing...base_layer.pre_quant_scale— i.e.base_layersurvives into the exported key.
Consider replacing the enumerated map with a generic strip (e.g. rewrite ".base_layer." → "." after the quantizer-key filtering) or, at minimum, warn/assert when a base_layer key reaches the end of the replacement loop unmatched so the next occurrence fails loudly instead of silently producing an undeployable checkpoint.
There was a problem hiding this comment.
Confirmed both cases and fixed generically in cc60139 — thanks, this one was worse than the bug the PR started with.
I ran postprocess_state_dict directly against both claims:
base_layer.bias→ dropped entirely.base_layer.input_quantizer._pre_quant_scale→ exported asl.base_layer.pre_quant_scale, i.e.base_layersurvives into the key, exactly as described.
The bias case matters for the model in the originating bug (Qwen2-7B-Instruct): Qwen2's q/k/v projections are bias=True unconditionally. My tiny-Qwen3 e2e model has no biases, which is why the new example test missed it. Re-ran the full quantize -> QLoRA train -> export flow on a tiny Qwen3 built with attention_bias=True: 0 bias tensors exported before the fix, 8 after, with no base_layer left in any key. So this would have shipped Qwen2 QLoRA checkpoints missing their attention biases.
Went with your first suggestion rather than the warn/assert: the enumerated map and the "base_layer" entry in skip_keys are both gone, replaced by a generic .base_layer. -> . strip applied to the output key in both branches of the loop. The weight_quantizer._scale / ._double_scale keys that the blanket skip used to catch are still removed by the existing RealQuantLinear.list_of_scale_tensors cleanup at the end, so nothing new leaks through.
test_postprocess_state_dict_qlora_strips_base_layer now covers bias, pre_quant_scale, and the quantizer-internal keys that must still be dropped.
| continue | ||
| # PEFT renames the quantized linear to `<name>.base_layer`, but `q_tensor_state` is | ||
| # keyed by the name it was saved with, so fall back to the stripped name. | ||
| key = name if name in q_tensor_state else name.removesuffix(".base_layer") |
There was a problem hiding this comment.
Bot comment.
This is the only library-code behavior change in the PR and it currently has no unit test — coverage comes solely from the GPU-heavy test_qwen3_qlora_nvfp4 example test. A small test that builds a compressed linear, renames it to <name>.base_layer (or constructs a q_tensor_state keyed without the suffix) and asserts the weight is re-wrapped as QTensorWrapper would pin this cheaply.
Also: the fallback only handles .base_layer as a suffix. If PEFT injection ever nests further (or the state was saved from a get_peft_model-wrapped model with a base_model.model. prefix), the lookup silently misses again and you get the same opaque shape error. A name.replace(".base_layer", "") normalization plus a debug log on miss would be more robust.
There was a problem hiding this comment.
Added the unit test in cc60139; pushing back on the .replace() suggestion.
Test: tests/unit/torch/opt/plugins/test_hf_patching.py::test_restore_qtensor_wrappers, parametrized on PEFT-wrapped vs not. It runs on CPU — NVFP4 quantize + compress work without a GPU — so it's a real compressed RealQuantLinear, not a mock. Verified it catches the regression: with the .base_layer fallback reverted, the wrap_in_lora=True case fails and the plain case still passes.
Worth noting for anyone touching this test: my first version passed even with the fix reverted, because RealQuantParameterDict.__setitem__ re-wraps a plain Parameter back into a QTensorWrapper when an existing wrapper of the same shape is present. The setup now deletes the entry first so the wrapper is genuinely gone, matching what the real from_pretrained path produces (I'd observed 0/15 modules wrapped there).
On name.replace(".base_layer", ""): keeping removesuffix. The lookup already tries the exact name first, so both key layouts are covered — state saved pre-adapters (keys without the suffix) and state saved from an adapter-wrapped model (keys with it). replace would also rewrite any module legitimately named base_layer mid-path. The base_model.model. prefix case needs get_peft_model, which QATTrainer has a TODO for but doesn't use yet (it calls add_adapter); happy to revisit as part of that migration, when there's a real checkpoint to test against.
| print_rank_0("Restored modelopt quantizer state dict") | ||
|
|
||
| if not ModeloptStateManager.is_converted(model): | ||
| restore_from_modelopt_state(model, modelopt_state) |
There was a problem hiding this comment.
Bot comment.
Two small things here:
modelopt_weightsis popped before theis_convertedcheck but only used inside it, so in the QLoRA (already-converted) path we load and discard it. Moving thepopinside the branch makes the intent clearer.- Adding
os.path.isfile(...)to the condition turns a previously loud failure into a silent no-op: a LoRA checkpoint that is not converted and has nomodelopt_state_train.pthwill now export as an unquantized model (quant_algo: null) instead of raising. Sinceis_converted()already handles the QLoRA case this PR is fixing, consider keeping the error (or emitting a warning) when the state file is missing and the model isn't converted.
There was a problem hiding this comment.
Both addressed in cc60139.
You're right that the os.path.isfile guard traded a loud failure for a silently unquantized export — that's a worse outcome than the FileNotFoundError it replaced. Removed it, so a LoRA checkpoint that isn't converted and has no modelopt_state_train.pth fails loudly again. is_converted() is doing the actual work for the QLoRA case this PR fixes, so the guard wasn't buying anything.
Also moved the modelopt_state_weights pop inside the branch that uses it. Net result is closer to the original code — the only change is not ModeloptStateManager.is_converted(model) in the condition — which makes the fix easier to review:
if hasattr(model, "peft_config") and not ModeloptStateManager.is_converted(model):
modelopt_state = mto.load_modelopt_state(f"{ckpt_path}/modelopt_state_train.pth")
restore_from_modelopt_state(model, modelopt_state)
...The import os added earlier is gone again.
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
All three previous comments are addressed, and addressed well:
postprocess_state_dictnow does a generic.base_layer.→.strip (enumerated map +"base_layer"skip-key removed). I verified the quantizer-internal keys that the blanket skip used to swallow (weight_quantizer._scale/._double_scale) are still dropped by the existingRealQuantLinear.list_of_scale_tensorscleanup, andtest_postprocess_state_dict_qlora_strips_base_layerpins the exact surviving/dropped sets includingbiasand AWQpre_quant_scale._restore_qtensor_wrappersnow has a real CPU unit test parametrized on PEFT-wrapped vs plain; deleting_parameters["weight"]before reassigning is the right setup (otherwiseRealQuantParameterDict.__setitem__re-wraps and the test passes with the fix reverted). Keepingremovesuffixoverreplaceis justified.examples/llm_qat/export.pydropped theos.path.isfileguard (loud failure restored) and moved themodelopt_state_weightspop inside the branch that uses it.
One thing I'd like the owner to confirm before merge (new, not previously raised): the not ModeloptStateManager.is_converted(model) guard applies to every peft checkpoint, not just QLoRA. QATTrainer.save_model writes modelopt_state.pth whenever is_converted, so a plain LoRA-QAT (fake-quant, uncompressed) output dir also restores state in from_pretrained and now skips set_quantizer_state_dict(model, modelopt_weights) from modelopt_state_train.pth. For QLoRA the scales live in the checkpoint buffers (_scale/_double_scale), so nothing is lost; for fake-quant LoRA the calibrated _amax is exactly what modelopt_state_train.pth carries (mto.modelopt_state stores buffer shapes, not values), so if adapter-only saves omit those buffers the export would proceed with uncalibrated quantizers instead of raising. This is not a regression (the old code hit AssertionError: Model already has modelopt state! on the same input), and the QLoRA e2e test now covers export — but test_qwen3_lora_qat_nvfp4 still stops after training, so the other checkpoint shape that this condition changed is untested. A quick confirmation that LoRA-QAT export still gets its amaxes (or an export step added to that test) would close this out.
The
|
|
/claude review |
There was a problem hiding this comment.
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.
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/examples/llm_qat/test_llm_qat.py`:
- Around line 193-202: Extend the scale validation in the exported base-model
test around the existing weight_scales loop to compare each weight_scale,
weight_scale_2, and input_scale tensor against a trusted calibrated reference,
such as the direct PTQ export or calibrated checkpoint state. Preserve the
current key-presence assertions and the leakage checks for base_layer and _amax,
and ensure the test verifies calibrated _amax-derived values survive restoration
rather than only verifying tensor names.
🪄 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: 18f9eb97-0d77-4653-bfc7-05e1b2884961
📒 Files selected for processing (1)
tests/examples/llm_qat/test_llm_qat.py
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| def _export_key(key: str) -> str: | ||
| # PEFT nests the quantized module under `base_layer`, which deployment does not expect. | ||
| # Strip it generically so new key types (bias, scales) do not need to be enumerated here. | ||
| return key.replace(".base_layer.", ".") if is_modelopt_qlora else key |
There was a problem hiding this comment.
[IMPORTANT Export] The generic .base_layer. strip is a real improvement over the enumerated map, but it can now silently overwrite an existing key instead of dropping one.
post_state_dict is keyed by the stripped name while iteration is over the raw state_dict. For a PEFT lora.Linear, named_parameters() yields both the nested …q_proj.base_layer.weight and, for any module that also exposes a same-named attribute at the wrapper level, …q_proj.weight. Whichever arrives second wins, with no warning:
post_state_dict[_export_key(key)] = value # last-write-wins on collisionThe old code had the same class of bug but failed loudly downstream (missing key → dequant error), whereas this fails silently with a plausible-looking checkpoint. Since the whole point of this PR is that a wrong base_layer mapping produces an undeployable checkpoint, it's worth making the collision detectable:
def _set(target: dict, key: str, value):
export_key = _export_key(key)
if export_key in target:
logger.warning(
f"Export key collision: '{key}' maps to '{export_key}', which is already "
f"populated. Keeping the first value."
)
return
target[export_key] = valueand call _set(post_state_dict, key, value) at both assignment sites. Even if no currently-supported architecture collides, an assert/warning here converts a future silent-corruption bug into an immediate, diagnosable failure — which is exactly the recurrence guard the comment on line 995-996 claims to provide.
There was a problem hiding this comment.
Agreed, and thanks for posting the correction below rather than leaving it standing — I reached the same conclusion independently while checking it.
Not taking the _set collision helper. Your own trace is the reason: _process_quantized_modules skips modules with a base_layer attribute (unified_export_hf.py:878), and peft exposes weight on the wrapper as a property while state_dict() emits only _parameters/_buffers, so no duplicate key reaches the strip. I would rather not add a warning path with no reachable trigger to guard it.
The zero-match warning from your other comment is different and I did take it — there the silent failure was demonstrably reachable, since it is exactly how this NVBug hid.
Your key-by-key trace of what the generic strip now newly admits matches mine, including that weight_quantizer._scale / ._double_scale survive the strip only to be dropped by the list_of_scale_tensors cleanup.
| # keyed by the name it was saved with, so fall back to the stripped name. | ||
| key = name if name in q_tensor_state else name.removesuffix(".base_layer") | ||
| if key not in q_tensor_state: | ||
| continue |
There was a problem hiding this comment.
[IMPORTANT ModeState] The fix is correct, but the loop still fails silently when nothing matches — which is the precise failure mode this PR exists to fix.
Before this PR, q_tensor_state was fully populated and zero modules matched, so every compressed weight stayed an unwrapped Parameter and the bug only surfaced much later as an opaque shape error deep in the NVFP4 dequant. Nothing in this function noticed. After the fix the .base_layer case is handled, but the next renaming layer (a different PEFT wrapper, ParamWrapper, a nested base_layer.base_layer, or the reverse direction where q_tensor_state keys carry .base_layer and the live module names do not) reproduces the identical silent miss.
Since you already know how many entries you expect to re-wrap, the check is nearly free:
state = load_modelopt_state(modelopt_state_path)
for _, mode_config in state["modelopt_state_dict"]:
q_tensor_state = mode_config.get("metadata", {}).get("q_tensor_state", {})
if not q_tensor_state:
continue
matched = 0
for name, module in model.named_modules():
if not isinstance(module, RealQuantLinear) or isinstance(module.weight, QTensorWrapper):
continue
key = name if name in q_tensor_state else name.removesuffix(".base_layer")
if key not in q_tensor_state:
continue
module._parameters["weight"] = QTensorWrapper(
qtensor=module.weight.data,
metadata=q_tensor_state[key]["metadata"],
)
matched += 1
if not matched:
warnings.warn(
f"Found {len(q_tensor_state)} saved compressed weight(s) in {modelopt_state_path} "
"but re-wrapped none — module names may have been remapped by a wrapper. "
"The model will likely fail when the packed weights are used."
)Note matched == 0 is the only safe condition to warn on: a partial match is legitimate here, because modules already holding a QTensorWrapper are skipped by the guard on line 114 and so never counted. This turns a silent, far-from-the-cause shape error into a message that names the file and the cause.
There was a problem hiding this comment.
Taken in 42c4671 — thanks for the follow-up narrowing it to the right condition.
I used pending and not matched rather than plain matched == 0: pending is the list of RealQuantLinear modules that still need wrapping, so a model whose weights are all already wrapped (skipped by the QTensorWrapper guard) does not warn spuriously. Your note that partial matches are legitimate is what made that distinction necessary.
if pending and not matched:
warnings.warn(
f"Found {len(q_tensor_state)} compressed weight(s) in {modelopt_state_path} but "
f"re-wrapped none of the {len(pending)} candidate module(s); their names may have "
"been remapped. The model will likely fail when the packed weights are used."
)Covered by test_restore_qtensor_wrappers_warns_when_nothing_matches, which double-nests the stand-in wrapper to produce a layout the lookup does not know about.
The reverse-direction case you raise in the same comment is also fixed — normalizing the .base_layer suffix on both the saved keys and the live names, per @Edwardf0t1's suggestion on the same line.
| def _export_key(key: str) -> str: | ||
| # PEFT nests the quantized module under `base_layer`, which deployment does not expect. | ||
| # Strip it generically so new key types (bias, scales) do not need to be enumerated here. | ||
| return key.replace(".base_layer.", ".") if is_modelopt_qlora else key |
There was a problem hiding this comment.
Correction to my comment above — downgrading this from IMPORTANT to SUGGESTION.
I asserted a concrete collision (…q_proj.base_layer.weight vs …q_proj.weight both landing on …q_proj.weight). I checked, and that path is not reachable on the current QLoRA export:
_process_quantized_modulesskips any module with abase_layerattribute (unified_export_hf.py:878), so the_QuantLoraLinearwrapper never registers its ownweight_scale/input_scalebuffers to clash with the ones onbase_layer.- PEFT exposes
weighton the wrapper as a property, andstate_dict()only emits_parameters/_buffers— so no duplicateweightentry exists.
So there is no live bug here, and the strip itself is correct. I also traced every key the generic strip now newly admits, and they all behave right:
base_layer.weight_quantizer._amax→ contains_amax(skip key), falls through to the replacements loop, matches no suffix, dropped ✅base_layer.weight_quantizer._scale/._double_scale→ stripped to…weight_quantizer._scale, then still removed by thelist_of_scale_tensorscleanup below, sinceendswith("weight_quantizer._scale")survives the strip ✅ (nicely pinned by the new test)base_layer.bias/base_layer.weight_scale_2→ newly retained — the fix ✅
The _set collision-warning helper I sketched is therefore optional hardening, not a blocker. The only argument for it: the old enumerated map failed loudly (an unmapped key vanished and blew up in dequant — how this NVBug surfaced), whereas a future silent overwrite would produce a plausible-looking checkpoint with one wrong tensor. Your call; it does not block approval.
There was a problem hiding this comment.
Acknowledged — answered on the parent thread. Short version: agreed the collision is unreachable, so I am not adding the _set helper; your key-by-key trace of what the generic strip newly admits matches mine. The zero-match warning from your transformers.py comment was taken, since that silent failure is reachable and is how this NVBug hid.
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
Reviewed the three fixes against the code on main. The diagnosis looks right in all three cases, the diff is minimal, and the plain-LoRA path (unquantized base -> not converted -> old behavior) is preserved. Checked style locally too: ruff format reports no changes and there are no new lint errors under the repo config (this test file is under # fmt: off, so the hand-aligned arg lists are intentional).
Five comments inline, none of them blocking correctness of the reported bug. Smaller nits not worth threads: export.py already imports pathlib.Path, so Path(ckpt_path) / "modelopt_state_train.pth" would avoid the new import os; assert processed_state_dict["layer1.weight_scale_2"] == torch.tensor([0.5]) asserts on a tensor and would read better as torch.equal(...); and load_file(base_model_dir / "model.safetensors") assumes an unsharded export (fine for tiny-qwen3). Also worth adding the **Bug Fixes** entry under 0.47 in CHANGELOG.rst -- and given the cherry-pick-0.46.0 label, making sure it lands in the right section on the pick.
| if hasattr(model, "peft_config") and os.path.isfile(modelopt_state_path): | ||
| modelopt_state = mto.load_modelopt_state(modelopt_state_path) |
There was a problem hiding this comment.
The trained quantizer state is loaded and then silently discarded on exactly the path this PR fixes. When the model is already converted (the QLoRA case), load_modelopt_state still reads the full modelopt_state_train.pth -- which includes modelopt_state_weights for the whole model -- and then nothing uses it. Cheap fix: fold the condition into the outer if so the file is never read:
if (
hasattr(model, "peft_config")
and not ModeloptStateManager.is_converted(model)
and os.path.isfile(modelopt_state_path)
):The more substantive half: is_converted only tells you the modes were restored, not that the quantizer amax/scales match what training saved. In today's flow they do -- _save_modelopt_state_with_weights (transformers_trainer.py:227-230) writes the state at trainer init from the already-quantized base, so it is the same state from_pretrained restored -- but that is an implicit invariant that a future change to when the trainer snapshots state would break silently. Either note it in the comment, or keep calling set_quantizer_state_dict unconditionally and only skip restore_from_modelopt_state.
There was a problem hiding this comment.
Taking the second half of this; pushing back on the first.
Folding os.path.isfile into the outer if: that guard is already gone. cc60139 removed it in response to your other comment, so the file is only read when the model is not converted — the QLoRA path never loads it. Current shape:
if hasattr(model, "peft_config") and not ModeloptStateManager.is_converted(model):
modelopt_state = mto.load_modelopt_state(f"{ckpt_path}/modelopt_state_train.pth")On the implicit invariant: you are right that is_converted proves the modes were restored, not that the quantizer values match what training saved. I went with your "note it in the comment" option rather than calling set_quantizer_state_dict unconditionally, because the unconditional version does not actually buy the robustness it looks like it does.
_save_modelopt_state_with_weights runs in QATTrainer.__init__, before training. So if a future change made training update amax, modelopt_state_train.pth would be just as stale as the base checkpoint — the adapter-only save omits those buffers, so neither source has the new values. Applying it unconditionally swaps one stale source for another while re-reading a file on the path you asked to keep clean. The failure would still be silent, just sourced differently.
Measured the current invariant rather than assuming it, on both checkpoint shapes: all 56 quantizer buffers match modelopt_state_train.pth exactly, and for fake-quant LoRA-QAT the 32 static quantizers all carry finite positive amax with none missing. set_quantizer_state_dict is re-applying values from_pretrained already loaded.
Comment added in 42c4671 (shortened in 4c6c36d):
# Skipping is safe only because QATTrainer writes modelopt_state_train.pth at trainer init,
# from that same base state.If you would rather have the belt-and-braces call anyway, say so and I will add it — it is a no-op today either way.
| # Restore modelopt quantizer state dict | ||
| # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this. | ||
| # For QLoRA the base checkpoint is quantized, so from_pretrained already restored the state. | ||
| modelopt_state_path = os.path.join(ckpt_path, "modelopt_state_train.pth") |
There was a problem hiding this comment.
The new os.path.isfile guard turns a loud failure into a silent one: a LoRA checkpoint with an unquantized base and no modelopt_state_train.pth now exports quietly as if nothing were quantized, where it previously raised FileNotFoundError. Worth an else: print_rank_0("[warn] no modelopt state found at ...") so that case is still visible in the log.
There was a problem hiding this comment.
Stale diff — that guard was already removed in cc60139, in response to the same point from the bot review. A LoRA checkpoint with an unquantized base and no modelopt_state_train.pth raises FileNotFoundError again rather than exporting quietly.
I added the guard for robustness and it was the wrong call for exactly the reason you give: a silently unquantized checkpoint is a worse outcome than a loud crash. Went with restoring the exception rather than a warning, since is_converted() is what handles the QLoRA case this PR fixes and the guard was buying nothing.
The import os it needed is gone too, so the pathlib.Path suggestion from your review body no longer applies — the line is back to the original f-string form.
| "base_layer.input_scale": "input_scale", | ||
| "base_layer.weight_scale": "weight_scale", | ||
| # NVFP4 global scale; the exported model cannot be dequantized without it. | ||
| "base_layer.weight_scale_2": "weight_scale_2", |
There was a problem hiding this comment.
Correct fix, and I checked the endswith matching cannot shadow it -- neither base_layer.weight nor base_layer.weight_scale is a suffix of ...base_layer.weight_scale_2, so the new entry is order-independent.
The concern is the shape of the surrounding code rather than this line: skip_keys.append("base_layer") below means any base_layer.* key that misses every replacement is dropped without a word, which is precisely how weight_scale_2 went missing. Adding one entry fixes NVFP4 and leaves the trap in place. Concretely still broken for AWQ: ...base_layer.input_quantizer._pre_quant_scale matches the earlier generic input_quantizer._pre_quant_scale -> pre_quant_scale entry, so the prefix survives and it exports as ...base_layer.pre_quant_scale. Latent today since only qlora_nvfp4.yaml ships, but it is the same bug class.
Suggestion: strip base_layer. generically for non-quantizer keys, or at minimum logger.warning when a base_layer.* key is dropped with no match, so the next missing scale surfaces at export time instead of at deployment.
There was a problem hiding this comment.
This one reads a stale diff — the generic strip you are asking for landed in cc60139, before this review. skip_keys.append("base_layer") and the enumerated map are both gone:
def _export_key(key: str) -> str:
return key.replace(".base_layer.", ".") if is_modelopt_qlora else keyapplied at both assignment sites. So the trap is removed rather than papered over, and the AWQ case you name is fixed too — I reproduced it first (...base_layer.input_quantizer._pre_quant_scale exported as l.base_layer.pre_quant_scale, prefix surviving exactly as you describe) and it now strips correctly.
Your bias point was the valuable one and I had missed it. It is not just latent: Qwen2 q/k/v are bias=True unconditionally, and Qwen2-7B-Instruct is the model in the originating NVBug. My tiny-Qwen3 e2e model has no biases, which is why the test did not catch it. Re-ran the full quantize -> QLoRA train -> export flow on a tiny Qwen3 built with attention_bias=True: 0 bias tensors exported before the fix, 8 after. Without this the PR would have shipped Qwen2 QLoRA checkpoints missing their attention biases.
Also confirmed the keys the blanket skip used to swallow are still dropped: weight_quantizer._scale / ._double_scale survive the strip but are removed by the existing RealQuantLinear.list_of_scale_tensors cleanup. test_postprocess_state_dict_qlora_strips_base_layer pins the exact surviving and dropped sets, including bias and AWQ pre_quant_scale.
Thanks for the endswith ordering check — matches what I found.
| # PEFT renames the quantized linear to `<name>.base_layer`, but `q_tensor_state` is | ||
| # keyed by the name it was saved with, so fall back to the stripped name. | ||
| key = name if name in q_tensor_state else name.removesuffix(".base_layer") |
There was a problem hiding this comment.
The fallback is one-directional: it handles "saved without .base_layer, loaded with it", but not the reverse. If the state was saved from a model compressed after adapters were attached (_quantize_model with compress: true), q_tensor_state keys carry .base_layer while a later load without adapters yields plain names -- no match, no re-wrap, and the same silent failure this PR is fixing. Normalizing the map once is symmetric and simpler:
q_tensor_state = {k.removesuffix(".base_layer"): v for k, v in q_tensor_state.items()}
...
key = name.removesuffix(".base_layer")
if key not in q_tensor_state:
continueSeparately, this only works because transformers 5.x injects adapters in place. transformers_trainer.py:210 carries a TODO to move to get_peft_model, which would prefix every module name with base_model.model. and break this again -- worth pinning that assumption in the comment.
There was a problem hiding this comment.
Correcting my earlier reply here, and the gap is now fixed in 42c4671.
I said the exact-name-first lookup meant "both key layouts are covered." That was wrong. It held only when the model is adapter-wrapped:
- state plain + model wrapped -> falls back to
removesuffix✅ - state suffixed + model wrapped -> exact match ✅
- state suffixed + model not wrapped -> both attempts yield the plain name, miss ❌
The third case is real, and I can name the flow that produces it, which the re-review couldn't find: mtq.compress runs after add_adapter in QATTrainer._quantize_model (transformers_trainer.py), so a library caller passing quant_args(compress=True, lora=True) saves q_tensor_state keyed with .base_layer. Loading that base without adapters then misses. Not reachable from the shipped examples/llm_qat/train.py, which never passes quant_args — so library-API-only, not a user-facing bug today.
Adopted @Edwardf0t1's symmetric normalization, which handles all three cases in one expression:
q_tensor_state = {k.removesuffix(".base_layer"): v for k, v in q_tensor_state.items()}
...
key = name.removesuffix(".base_layer")Also took the zero-match warning. It is gated on pending and not matched rather than plain matched == 0, so a model whose weights are all already wrapped (skipped by the QTensorWrapper guard) does not warn spuriously.
test_restore_qtensor_wrappers is now parametrized on both axes (4 cases). Verified the new [state_keyed_with_base_layer=True, wrap_in_lora=False] case fails against the old one-directional lookup, and added test_restore_qtensor_wrappers_warns_when_nothing_matches for the warning.
On pinning the get_peft_model assumption in a comment: I had it in, but the comments were cut back for length (4c6c36d). It is captured here in the thread instead.
| ) | ||
|
|
||
|
|
||
| def _run_export(ckpt_dir: str, export_dir: str): |
There was a problem hiding this comment.
Good to see the export step added to the e2e -- but note this is the only coverage for the _restore_qtensor_wrappers change, and it costs a full GPU QLoRA training run. grep finds no test anywhere referencing _restore_qtensor_wrappers or q_tensor_state, so that function has no unit coverage at all today.
A cheap unit test would guard the rename directly: a stub module tree containing a RealQuantLinear at ...q_proj.base_layer plus a hand-built q_tensor_state keyed without the suffix, asserting the weight comes back as a QTensorWrapper. That would also have caught the reverse-direction gap I flagged on transformers.py.
There was a problem hiding this comment.
That unit test exists — cc60139 added it, one commit before this review, and it is close to what you sketch: tests/unit/torch/opt/plugins/test_hf_patching.py::test_restore_qtensor_wrappers.
It builds a real compressed RealQuantLinear (NVFP4 quantize + compress both run on CPU, so no GPU needed), nests it under a _LoraLike stand-in for peft's lora.Linear, and asserts the weight comes back as a QTensorWrapper. 42c4671 then parametrized it on both key directions — which, as you predicted, is what caught the reverse-direction gap you flagged on transformers.py.
One trap worth recording for anyone extending it: the setup must del module._parameters["weight"] before assigning the plain Parameter. RealQuantParameterDict.__setitem__ re-wraps a same-shape plain Parameter back into a QTensorWrapper, so my first version of this test passed even with the fix reverted.
On the e2e cost — agreed it should not be the only coverage, which is why the unit test carries the real load. The export step stays because it is the only thing that exercises the is_converted guard on the fake-quant checkpoint shape, and it now compares exported scales against a direct PTQ export rather than just checking key presence (CodeRabbit's point).
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review #3 (head 776f47e5d). The single item left open in the previous round is now closed.
Previously flagged, all resolved:
postprocess_state_dict— enumeratedbase_layer.*rename map and the"base_layer"skip-key are gone, replaced by a generic.base_layer.→.strip applied at both assignment sites. I re-verified inquant_utils.pythat the keys the blanket skip used to swallow are still dropped:base_layer.weight_quantizer._amaxhits the_amaxskip key and matches no replacement suffix, andweight_quantizer._scale/._double_scalesurvive the strip only to be removed by theRealQuantLinear.list_of_scale_tensorscleanup.test_postprocess_state_dict_qlora_strips_base_layerpins the exact surviving/dropped sets includingbiasand AWQpre_quant_scale._restore_qtensor_wrappers— now covered by a real CPU unit test parametrized on PEFT-wrapped vs plain, with_parameters["weight"]deleted before reassignment so it genuinely fails with the fix reverted.examples/llm_qat/export.py— theos.path.isfileguard is gone (loud failure restored) and themodelopt_state_weightspop lives inside the branch that consumes it; net diff is a one-condition change.- New in this round: the
not ModeloptStateManager.is_converted(model)guard also changes the fake-quant LoRA-QAT path, which previously had no export coverage.test_qwen3_lora_qat_nvfp4now runs export and assertsquant_algo == NVFP4, that everyweight_scalehas a matchingweight_scale_2andinput_scale(both amax-derived, so their absence is what an uncalibrated fallback would look like), and that nobase_layer/_amaxkeys leak.export.pysetsis_qlora = hasattr(model, "peft_config"), so this checkpoint does take thebase_model/+is_modelopt_qlora=Truepath the test asserts on — the coverage is real, and the author's manual verification (base amaxes arrive via the base checkpoint's persistent buffers) matches the code.
Non-blocking leftovers, owner's call:
CHANGELOG.rstbug-fix entry under 0.47 was requested by Edwardf0t1 and is still absent; worth confirming it lands in the right section on thecherry-pick-0.46.0pick.- The optional hardening suggestions (warn when
_restore_qtensor_wrappersre-wraps zero modules; warn on export-key collision inpostprocess_state_dict) were not taken. Both are defensive-only — no reachable bug was demonstrated for either — but the zero-match warning is cheap and would turn the exact class of silent failure this PR fixes into a diagnosable message next time a wrapper renames modules. - The e2e assertions are key-presence rather than value comparisons against a trusted PTQ reference (CodeRabbit's point). Given the author verified byte-identical base exports out of band and a value-level oracle would double the GPU cost, key presence plus
quant_algoseems a reasonable trade-off here. - Reverse-direction lookup in
_restore_qtensor_wrappers(state keyed with.base_layer, live names without) is still unhandled despite the reply implying both layouts are covered; I couldn't find a reachable flow that produces it, so it stays a latent robustness gap rather than a bug.
No licensing changes (headers untouched, no vendored code) and no prompt-injection content in the PR metadata.
Complex PR: 1 existing test file modified or removed. Looping in a human for approval.
Re-review #3 leftovers — all four closed (
|
| Leftover | Resolution |
|---|---|
| CHANGELOG entry absent | Added. Filed under 0.46, not 0.47, since the PR carries cherry-pick-0.46.0 — so the pick applies cleanly instead of needing a hand-edit |
| Zero-match warning not taken | Taken, gated on pending and not matched so an all-already-wrapped model does not warn spuriously |
| Key-presence, not value comparison | LoRA-QAT test now compares every exported scale against a direct PTQ export |
| Reverse-direction lookup unhandled | Fixed via symmetric normalization of the .base_layer suffix on both sides |
Two corrections to things I said earlier:
-
The re-review is right that my reply on the reverse-direction thread implied both layouts were covered when they were not. The old lookup handled it only when the model was adapter-wrapped; state keyed with
.base_layerloaded onto plain module names missed. My claim was wrong, and the thread now has the full case table. -
The re-review could not find a reachable flow producing that layout. There is one:
mtq.compressruns afteradd_adapterinQATTrainer._quantize_model, so a library caller passingquant_args(compress=True, lora=True)saves keys with the suffix. It is not reachable from the shippedexamples/llm_qat/train.py, which never passesquant_args— library-API-only, not user-facing today, but more concrete than "latent".
Declined, with reasons in-thread: the postprocess_state_dict collision helper (the author of that suggestion posted a correction showing it is unreachable), and calling set_quantizer_state_dict unconditionally (it swaps one stale source for another rather than fixing the hypothetical it targets — the invariant is documented in a comment instead).
Verification: export suite 49 passed; tests/unit/torch/opt/plugins/ 30 passed / 1 skipped; both example tests pass (LoRA-QAT 4:07 with the added export, QLoRA 2:04). Each new test was checked to fail against the unfixed code. ruff/mypy/bandit clean.
One open item for whoever runs the pick: the changelog bullet sits under 0.46 on main, which is deliberate for the cherry-pick but does mean main carries a 0.46 entry for a fix merging after 0.46 was cut. Easy to move if that is not the convention here.
Trim the explanatory comments added in the previous commit to two lines each. No functional change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
…it test - postprocess_state_dict: strip ".base_layer" generically instead of enumerating renames. The enumerated map silently dropped every unlisted key, which lost linear biases (Qwen2 q/k/v have them) and left "base_layer" in the exported AWQ pre_quant_scale key. - export.py: drop the os.path.isfile guard so a LoRA checkpoint with no modelopt state still fails loudly instead of exporting unquantized, and pop the quantizer weights inside the branch that uses them. - Add a CPU unit test for the .base_layer fallback in _restore_qtensor_wrappers, and cover bias / pre_quant_scale in the export rename test. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
The is_converted guard also changes the fake-quant LoRA-QAT path, where the calibrated amaxes rather than packed weights are what must survive the load. Verified they do: the base PTQ safetensors carries all 32 _amax buffers, so from_pretrained restores them and the export matches a direct PTQ export byte-for-byte. Extend the test so that path stops being untested. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
…value oracle - _restore_qtensor_wrappers: normalize the .base_layer suffix on both the saved q_tensor_state keys and the module names, so the lookup works whether the model was compressed before adapters were attached (quantize.py --compress) or after (QATTrainer._quantize_model). Warn when saved weights match no module at all, which is the silent failure this PR set out to fix. - export.py: document why skipping the restore is safe (QATTrainer snapshots the state at trainer init from the base state from_pretrained already restored). - LoRA-QAT example test: compare exported scales against a direct PTQ export rather than only asserting key presence. - Unit tests: cover the reverse key direction and the no-match warning. - CHANGELOG: add the 0.47 Bug Fixes entry. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
The PR is labeled cherry-pick-0.46.0, so the entry belongs in the 0.46 section. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Trim the added comments and test docstrings to two lines each. No code change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
4c6c36d to
46596f0
Compare
| - Fix ``examples/vllm_serve`` serving shared experts uncalibrated: their ``gate_proj``/``up_proj`` quantizer keys were not merged into ``gate_up_proj`` on reload, so they matched no module and were dropped. | ||
| - Fix Qwen3-VL MoE PTQ failing on ``transformers>=5.12`` with ``AttributeError: 'QuantQwen3VLMoeTextExperts' object has no attribute 'hidden_size'`` (NVBug 6518551). transformers 5.12 moved ``Qwen3VLMoeTextExperts`` onto the standard ``@use_experts_implementation`` fused layout (``hidden_size``/``expert_dim`` renamed to ``hidden_dim``/``intermediate_dim``, ``gate_up_proj`` transposed to ``(num_experts, 2*intermediate_dim, hidden_dim)``, two ``F.linear`` calls per expert), but the legacy ``_QuantQwen3VLMoeTextExperts`` wrapper stayed statically registered and shadowed on-the-fly detection. The new layout is now left to ``register_fused_experts_on_the_fly``, which claims it with the generic ``_QuantFusedExperts``; the legacy wrapper is still registered on ``transformers<5.12``, whose ``torch.bmm``-based forward the generic wrapper cannot intercept. | ||
| - Fix ``examples/hf_ptq`` multi-node FSDP2 export (``--use_fsdp2``) failing with ``RuntimeError: Cannot set version_counter for inference tensor``. ``export_quantized`` wrapped its whole body in ``torch.inference_mode()``, so the full params gathered by ``get_model_state_dict(full_state_dict=True)`` were inference tensors and the subsequent ``state_dict()`` -> ``param.detach()`` could not set their version counter. The export context is now ``torch.no_grad()``, which still disables autograd but keeps the gathered params as normal tensors. | ||
| - Fix QLoRA export in ``examples/llm_qat/export.py`` failing with ``AssertionError: Model already has modelopt state!`` (NVBug 6542481). The QLoRA training output is an adapter-only checkpoint, so ``from_pretrained`` resolves the quantized base model from ``adapter_config.json`` and ``enable_huggingface_checkpointing`` already restores its ModelOpt state; the export then restored a second time. It now restores only when the loaded model is not already converted. Two further breakages on the same path are also fixed: ``_restore_qtensor_wrappers`` matched no modules because PEFT re-parents the quantized linear as ``<name>.base_layer`` while ``q_tensor_state`` is keyed by the name it was saved with (the packed NVFP4 weight then reached ``F.linear`` and raised a shape error), and ``postprocess_state_dict`` silently dropped every ``base_layer.*`` key missing from a hand-maintained rename map — losing the NVFP4 ``weight_scale_2`` global scale and any linear ``bias`` (Qwen2-style q/k/v biases), and leaving ``base_layer`` in the exported AWQ ``pre_quant_scale`` key. The rename is now a generic ``.base_layer.`` strip. |
4 participants
What does this PR do?
Type of change: Bug fix
Fixes
AssertionError: Model already has modelopt state!when exporting a QLoRA checkpoint(NVBug 6542481). The QLoRA output is adapter-only, so
from_pretrainedresolves the quantizedbase model and already restores the ModelOpt state;
export.pythen restored a second time.Fixing that exposed two more breakages on the same path, also fixed here:
_restore_qtensor_wrappersmissed every module — PEFT renames the compressed linears to<name>.base_layer, so no weight got re-wrapped and the packed NVFP4 weight hit a shape error.postprocess_state_dictdroppedweight_scale_2(missing from the QLoRA rename map), leavingthe exported checkpoint impossible to dequantize.
Usage
No API change —
examples/llm_qat/export.py --pyt_ckpt_path <qlora_ckpt> --export_path <out>now completes on the documented quantize → train → export flow.
Testing
Reproduced in the reported environment (TRT-LLM 1.3.0rc22, transformers 5.5.4, NVFP4).
test_qwen3_qlora_nvfp4and a unit test for the QLoRAbase_layerrename; both fail without the fix.the bf16 original (worst rel. error 0.10).
tests/gpu/torch/export/test_export.py(49 passed), save/load plugin tests.Before your PR is "Ready for review"
Additional Information
Fixes NVBug 6542481.
Summary by CodeRabbit
Bug Fixes
Tests