Add LSQ NVFP4 export support - #1847
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
realAsma
left a comment
realAsma
left a comment
There was a problem hiding this comment.
BB: Rebase this PR to the latest main, remove the unnecessary changes.
🐝 I’m on it — I’ll rebase onto the latest |
Blocked before rebase: this PR still targets |
|
BB: Can you rebase this PR to latest main and only add the changes required to enable export with LSQ? (btw, I have renamed LAQ to LSQ, see #1884) |
|
BB: Please update PR description and PR title after addressing #1847 (comment) |
🐝 I’m on it — I’ll rebase onto the latest |
🐝 Acknowledged this instruction. I’ll update the PR title and description after addressing the linked feedback. |
🐝 Got it — I’ll update the PR title and description to match the narrowed LSQ-export scope after the branch changes are prepared. |
Blocked for now: the PR still declares |
Blocked outcome for the title/description request: the guarded pass paused because the PR still targets |
|
BB: Yes please retarget this PR to latest main. then address #1847 (comment) |
🐝 Acknowledged — I’ll retarget this PR to the latest |
Blocked after the requested retarget attempt: I safely cleared the unpublished wrong-base rebase and restored |
|
BB: Done, I manually rebase to main. do the rest. |
Update: the live PR base is now |
🐝 Got it — I’m taking over from the manual rebase and will finish the LSQ export work, validation, and PR metadata updates. |
I verified the PR now targets |
30e8b91 to
448e58e
Compare
🐝 Updated this PR to the requested LSQ-export-only scope.
The PR remains draft; no ready-for-review promotion or reviewer notification was made. |
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1847 +/- ##
==========================================
+ Coverage 60.21% 60.22% +0.01%
==========================================
Files 519 519
Lines 59377 59385 +8
==========================================
+ Hits 35755 35767 +12
+ Misses 23622 23618 -4
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:
|
realAsma
left a comment
realAsma
left a comment
There was a problem hiding this comment.
BB: For now I want to support export only for LSQ tied variant. Can you update this PR for that? I think with LSQ ted variant only support, we can significantly simplify the PR (because its looks similar to the NVFP4 with MSE/four-over-six weights?
You can raise an error if Dual LSQ is tried to export.
🐝 Got it — I’ll update the PR to support export only for tied LSQ, simplify the implementation, and reject Dual LSQ export. |
🐝 I triaged this as a valid fix: tied LSQ export only, with Dual LSQ export rejected. The PR is currently five commits behind |
Signed-off-by: realAsma <akuriparambi@nvidia.com>
Signed-off-by: realAsma <akuriparambi@nvidia.com>
448e58e to
540d2fc
Compare
Published the review fix at |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope
Full review of all 4 changed files (+118/-1): quant_utils.py, unified_export_hf.py, nvfp4_tensor.py, and the CPU test. Small PR — reviewed in full, nothing deprioritized.
Findings
- CRITICAL: 0
- IMPORTANT: 0
- SUGGESTION: 1 (guard reachability from non-HF export entry points — non-blocking)
Assessment
The change is tightly scoped and correct.
Algorithm correctness (verified end-to-end): For the only exportable variant — tied-amax LSQ with quantize_scales=True and quantize_pre_scale=True — the training forward (_fake_quantize) uses scale_pre == scale_post == _block_scale_from_amax(amax_post, quantize=True). The export path reconstructs the identical fp8-quantized block scale via weight_scale * weight_scale_2, and the new test asserts exactly this equality against _block_scale_from_amax(amax_post, True). The rejected set (tied_amax=False, quantize_pre_scale=False, quantize_scales=False) is precisely where scale_pre != scale_post or scales aren't fp8 — i.e. the configs that could not round-trip — so the NotImplementedError guard draws the boundary in the right place.
State/compat: No modelopt_state schema change, no public API signature change. The nvfp4_tensor.py fallback (getattr(_amax, None) then .amax) preserves existing static-quantizer behavior (they retain _amax) and only diverges for tied LSQ where enable_lsq deleted _amax. Backward compatible.
Test coverage: Good — parametrized over learnable-amax sets, an end-to-end _export_quantized_weight registration check, and rejection of all three unsupported variants.
Low risk. The single suggestion is a defense-in-depth nicety about which export entry points route through the guard, not a correctness issue for the current HF-export scope.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
1 participant
Summary
Adds LSQ NVFP4 support to unified Hugging Face export.
Testing
pre-commit run --files modelopt/torch/export/quant_utils.py modelopt/torch/export/unified_export_hf.py tests/unit/torch/quantization/test_nvfp4_static_export_cpu.pypytest_pwd tests/unit/torch/quantization/test_lsq.py tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py -q -x— 66 passedpytest_pwd tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py -q -x— 18 passedChecklist