Fix prequant layernorm export without scales#1838
Merged
Merged
Conversation
Signed-off-by: realAsma <akuriparambi@nvidia.com>
|
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. |
Contributor
📝 WalkthroughWalkthroughAdds an early return in ChangesPre-quant LayerNorm fusion
PTQ int4 blockwise recipe
Estimated code review effort🎯 2 (Simple) | ⏱️ ~15 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Signed-off-by: realAsma <akuriparambi@nvidia.com>
Contributor
|
4 tasks
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1838 +/- ##
==========================================
- Coverage 77.29% 76.77% -0.53%
==========================================
Files 513 513
Lines 56920 56922 +2
==========================================
- Hits 43997 43700 -297
- Misses 12923 13222 +299
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:
|
meenchen
reviewed
Jun 26, 2026
| @@ -16,47 +16,19 @@ | |||
| metadata: | |||
Contributor
There was a problem hiding this comment.
@realAsma could you share why we need this int4 recipe in the modelopt_recipes/general/ptq/? I think we don't support export for int4 blockwise. cc @Edwardf0t1
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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 `@modelopt/torch/export/quant_utils.py`:
- Around line 1260-1263: The LayerNorm export logic is checking the private
TensorQuantizer buffer `_pre_quant_scale` directly, which can miss the case
where `TensorQuantizer.pre_quant_scale` is `None` and incorrectly fuse a stale
scale. Update the guard in `quant_utils` where the `modules[0].input_quantizer`
is handled to use the public `pre_quant_scale` property instead of `hasattr(...,
"_pre_quant_scale")`, and only proceed with the fusion path when that property
is present and non-None before moving it to `layernorm_module.weight.device`.
🪄 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: d4845bc5-56e1-4e34-bf8a-b5a0dad9aba3
📒 Files selected for processing (4)
modelopt/torch/export/quant_utils.pymodelopt_recipes/general/ptq/int4_blockwise_weight_only.yamltests/unit/recipe/test_loader.pytests/unit/torch/export/test_unified_export_hf.py
Comment on lines
+1260
to
+1263
| if not hasattr(modules[0].input_quantizer, "_pre_quant_scale"): | ||
| return | ||
|
|
||
| pre_quant_scale = modules[0].input_quantizer._pre_quant_scale.to(layernorm_module.weight.device) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Use pre_quant_scale instead of the private _pre_quant_scale buffer.
TensorQuantizer.pre_quant_scale is also None when pre-quant scaling is disabled, not just when the buffer is missing. This guard bypasses that contract, so export can still fuse a stale scale into the LayerNorm even though the forward path no longer applies it.
Suggested fix
- if not hasattr(modules[0].input_quantizer, "_pre_quant_scale"):
+ pre_quant_scale = modules[0].input_quantizer.pre_quant_scale
+ if pre_quant_scale is None:
return
-
- pre_quant_scale = modules[0].input_quantizer._pre_quant_scale.to(layernorm_module.weight.device)
+ pre_quant_scale = pre_quant_scale.to(layernorm_module.weight.device)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if not hasattr(modules[0].input_quantizer, "_pre_quant_scale"): | |
| return | |
| pre_quant_scale = modules[0].input_quantizer._pre_quant_scale.to(layernorm_module.weight.device) | |
| pre_quant_scale = modules[0].input_quantizer.pre_quant_scale | |
| if pre_quant_scale is None: | |
| return | |
| pre_quant_scale = pre_quant_scale.to(layernorm_module.weight.device) |
🤖 Prompt for 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.
In `@modelopt/torch/export/quant_utils.py` around lines 1260 - 1263, The LayerNorm
export logic is checking the private TensorQuantizer buffer `_pre_quant_scale`
directly, which can miss the case where `TensorQuantizer.pre_quant_scale` is
`None` and incorrectly fuse a stale scale. Update the guard in `quant_utils`
where the `modules[0].input_quantizer` is handled to use the public
`pre_quant_scale` property instead of `hasattr(..., "_pre_quant_scale")`, and
only proceed with the fusion path when that property is present and non-None
before moving it to `layernorm_module.weight.device`.
kevalmorabia97
pushed a commit
that referenced
this pull request
Jul 3, 2026
- Skip pre-quant LayerNorm fusion when the representative input quantizer has no `_pre_quant_scale` buffer. - Add CPU unit coverage for the no-scale weight-only path and the existing fusion/removal behavior when pre-quant scales are present. - Compose the public `general/ptq/int4_blockwise_weight_only` recipe from the shared recipe units so it stays aligned with `INT4_BLOCKWISE_WEIGHT_ONLY_CFG`. NVBug: https://nvbugspro.nvidia.com/bug/6311597 NVBug 6311597 reports a deterministic HF export crash for `general/ptq/int4_blockwise_weight_only` on Llama-3.1-8B. That weight-only recipe disables input quantization and skips calibration, so `_pre_quant_scale` is not registered, but export still reaches `fuse_prequant_layernorm` through the AWQ-like format path. This PR also carries the recipe-sync diff that was previously in draft PR #1836; PR #1836 has been closed after moving that diff here. - `pre-commit run --files modelopt/torch/export/quant_utils.py tests/unit/torch/export/test_unified_export_hf.py modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml tests/unit/recipe/test_loader.py` - `pytest_pwd tests/unit/torch/export/test_unified_export_hf.py::test_fuse_prequant_layernorm_skips_modules_without_pre_quant_scale tests/unit/recipe/test_loader.py::test_load_recipe_all_builtins tests/unit/recipe/test_loader.py::test_general_ptq_yaml_matches_config_dicts tests/unit/recipe/test_presets.py -q` -> 29 passed Full GB10/Llama repro was not run locally. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> * **Bug Fixes** * Improved handling of layer norm fusion so it safely skips cases where pre-quantization scale data is unavailable, avoiding unexpected errors. * Updated the layer norm fusion flow to correctly apply scaling when pre-quantization data is present. * **Tests** * Expanded coverage for export and recipe loading scenarios, including cases with and without pre-quantization scale data. * Added validation for the new int4 blockwise weight-only recipe. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: realAsma <akuriparambi@nvidia.com>
kevalmorabia97
pushed a commit
that referenced
this pull request
Jul 3, 2026
- Skip pre-quant LayerNorm fusion when the representative input quantizer has no `_pre_quant_scale` buffer. - Add CPU unit coverage for the no-scale weight-only path and the existing fusion/removal behavior when pre-quant scales are present. - Compose the public `general/ptq/int4_blockwise_weight_only` recipe from the shared recipe units so it stays aligned with `INT4_BLOCKWISE_WEIGHT_ONLY_CFG`. NVBug: https://nvbugspro.nvidia.com/bug/6311597 NVBug 6311597 reports a deterministic HF export crash for `general/ptq/int4_blockwise_weight_only` on Llama-3.1-8B. That weight-only recipe disables input quantization and skips calibration, so `_pre_quant_scale` is not registered, but export still reaches `fuse_prequant_layernorm` through the AWQ-like format path. This PR also carries the recipe-sync diff that was previously in draft PR #1836; PR #1836 has been closed after moving that diff here. - `pre-commit run --files modelopt/torch/export/quant_utils.py tests/unit/torch/export/test_unified_export_hf.py modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml tests/unit/recipe/test_loader.py` - `pytest_pwd tests/unit/torch/export/test_unified_export_hf.py::test_fuse_prequant_layernorm_skips_modules_without_pre_quant_scale tests/unit/recipe/test_loader.py::test_load_recipe_all_builtins tests/unit/recipe/test_loader.py::test_general_ptq_yaml_matches_config_dicts tests/unit/recipe/test_presets.py -q` -> 29 passed Full GB10/Llama repro was not run locally. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> * **Bug Fixes** * Improved handling of layer norm fusion so it safely skips cases where pre-quantization scale data is unavailable, avoiding unexpected errors. * Updated the layer norm fusion flow to correctly apply scaling when pre-quantization data is present. * **Tests** * Expanded coverage for export and recipe loading scenarios, including cases with and without pre-quantization scale data. * Added validation for the new int4 blockwise weight-only recipe. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: realAsma <akuriparambi@nvidia.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
added
the
cherry-pick-done
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
4 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.
Summary
_pre_quant_scalebuffer.general/ptq/int4_blockwise_weight_onlyrecipe from the shared recipe units so it stays aligned withINT4_BLOCKWISE_WEIGHT_ONLY_CFG.Context
NVBug: https://nvbugspro.nvidia.com/bug/6311597
NVBug 6311597 reports a deterministic HF export crash for
general/ptq/int4_blockwise_weight_onlyon Llama-3.1-8B. That weight-only recipe disables input quantization and skips calibration, so_pre_quant_scaleis not registered, but export still reachesfuse_prequant_layernormthrough the AWQ-like format path.This PR also carries the recipe-sync diff that was previously in draft PR #1836; PR #1836 has been closed after moving that diff here.
Tests
pre-commit run --files modelopt/torch/export/quant_utils.py tests/unit/torch/export/test_unified_export_hf.py modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml tests/unit/recipe/test_loader.pypytest_pwd tests/unit/torch/export/test_unified_export_hf.py::test_fuse_prequant_layernorm_skips_modules_without_pre_quant_scale tests/unit/recipe/test_loader.py::test_load_recipe_all_builtins tests/unit/recipe/test_loader.py::test_general_ptq_yaml_matches_config_dicts tests/unit/recipe/test_presets.py -q-> 29 passedFull GB10/Llama repro was not run locally.
Summary by CodeRabbit
Bug Fixes
Tests