MSE Observer Enhancement - #2950
Conversation
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review. Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new 'Four Over Six' (4/6) adaptive block scaling observer for NVFP4 quantization, adds a corresponding example script, and updates the MSE quantization search logic to support range expansion and custom global scale maximums. The reviewer feedback identifies several key improvements: ensuring idempotency in the new observer's parameter retrieval via caching, fixing a range-scaling bug when expanding positive minimum values, removing a leftover debugging statement, and caching dynamically created scale data classes to avoid overhead and compilation issues under torch.compile.
coderabbitai
Bot
added
enhancement
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
src/llmcompressor/modifiers/quantization/calibration.py (1)
147-147: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the commented-out debugger.
#import pdb; pdb.set_trace()is dead debug code in production source and is easy to re-enable accidentally. Remove it before merge.🤖 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 `@src/llmcompressor/modifiers/quantization/calibration.py` at line 147, Remove the commented-out debugger line from the calibration module, including the “import pdb; pdb.set_trace()” statement. Do not replace it with other debugging code.src/llmcompressor/observers/mse_quant.py (1)
29-30: 📐 Maintainability & Code Quality | 🔵 TrivialDocument the
expand/maxshrinkcoupling and add a docstring entry forglobal_scale_max.With default
maxshrink=0.20, grid=100.0, the shrink factorponly sweeps[0.8, 1.0]. If a caller setsexpand=2.0without also increasingmaxshrink, the search never gets close to the true observed range (would needp≈0.5). This coupling isn't documented, andglobal_scale_max(new parameter) has no docstring entry unlikeexpand.Also applies to: 50-53
🤖 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 `@src/llmcompressor/observers/mse_quant.py` around lines 29 - 30, Update the parameter documentation for expand, maxshrink, and global_scale_max in the relevant observer API. Document that increasing expand requires a corresponding increase to maxshrink so the shrink-factor search can reach the observed range (for example, expand=2.0 needs p near 0.5 rather than the default lower bound of 0.8). Add global_scale_max’s purpose and behavior alongside the existing expand entry.
🤖 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 `@examples/quantization_w4a4_fp4/llama3_fouroversix_example.py`:
- Around line 88-90: Update the quantization configuration around the observer
argument to use an observer name registered by the repository or
compressed_tensors, preserving the MSE scale-selection behavior where supported;
alternatively, add the required registration before this configuration is
constructed.
In `@src/llmcompressor/observers/base.py`:
- Around line 107-114: Update the global-scale setup in the observer method
containing gs_max so gparam_kwargs["scale_data"] uses an _CustomFP8ScaleData
instance initialized with gs_max as the configured maximum and its negation as
the minimum. Preserve the existing conditional behavior when global_scale_max is
absent.
- Around line 30-35: Remove the undefined class-level max and min assignments
from _CustomFP8ScaleData, allowing callers to provide these bounds explicitly
during instantiation as done by mse_quant.py. Keep the remaining FloatArgs
configuration unchanged and ensure importing the module no longer evaluates an
undefined gs_max reference.
---
Nitpick comments:
In `@src/llmcompressor/modifiers/quantization/calibration.py`:
- Line 147: Remove the commented-out debugger line from the calibration module,
including the “import pdb; pdb.set_trace()” statement. Do not replace it with
other debugging code.
In `@src/llmcompressor/observers/mse_quant.py`:
- Around line 29-30: Update the parameter documentation for expand, maxshrink,
and global_scale_max in the relevant observer API. Document that increasing
expand requires a corresponding increase to maxshrink so the shrink-factor
search can reach the observed range (for example, expand=2.0 needs p near 0.5
rather than the default lower bound of 0.8). Add global_scale_max’s purpose and
behavior alongside the existing expand entry.
🪄 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: Pro
Run ID: c3310b93-36f7-4e0f-ba69-26a5558b8e76
📒 Files selected for processing (13)
docs/scripts/zensical_gen_files.pyexamples/multimodal_vision/gemma4_example.pyexamples/quantization_w4a4_fp4/llama3_fouroversix_example.pyexamples/quantization_w8a8_fp8/qwen3_reranker_example.pyexamples/quantization_w8a8_int8/benchmark_smoothquant_ddp.pysrc/llmcompressor/modifiers/quantization/calibration.pysrc/llmcompressor/observers/__init__.pysrc/llmcompressor/observers/base.pysrc/llmcompressor/observers/mse.pysrc/llmcompressor/observers/mse_quant.pytests/llmcompressor/modifiers/autoround/test_base.pytests/llmcompressor/modifiers/quantization/test_sequential_observation.pytests/llmcompressor/transformers/compression/test_compression_ddp.py
🔗 Linked repositories identified
CodeRabbit considers these linked repositories for cross-repo context during reviews:
vllm-project/compressed-tensors(manual)
|
This pull request has merge conflicts that must be resolved before it can be |
ea69b1d to
7fd9a08
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
|
The quality checks have failed. Please run |
Merge Protections🟢 All 2 merge protections satisfied — ready to merge. Show 2 satisfied protections🟢 Require one maintainer reviewAll PRs must have at least one approving review from a maintainer before merging.
🟢 Require two reviewsPRs labelled "two-reviews" must have at least two approving reviews before merging.
|
brian-dellabetta
left a comment
brian-dellabetta
left a comment
There was a problem hiding this comment.
This looks good to me pending @HDCharles 's thoughts, since I know you've been working closely on the 4over6 stuff with him. usage of observer kwargs seems like the right place to plug this in
|
The quality checks have failed. Please run |
|
The quality checks have failed. Please run |
- Add `expand` parameter to MSE grid search (default 1.0, no behavior change) - Add `nvfp4_expanded_mse` observer with tuned defaults for NVFP4 range expansion - Add `expand` and `nvfp4_expanded_imatrix_mse` observer to imatrix - Skip scale_dtype rounding during grid search for NVFP4 observers - Add llama3 NVFP4 expanded MSE example Signed-off-by: Roderick Wu <roderick@neuralmagic.com>
a3c3885 to
10a566d
Compare
|
The quality checks have failed. Please run |
Signed-off-by: Roderick-Wu <roderickwu2003@gmail.com>
|
The quality checks have failed. Please run |
HDCharles
left a comment
HDCharles
left a comment
There was a problem hiding this comment.
good, the imatrix thing is an experiment so we add functionality (expand) but not the UX helper, once that's removed this will be ready to land
Signed-off-by: Roderick Wu <Roderick-Wu@h100-03.nemg-001.lab.rdu2.dc.redhat.com>
…e_example.py Co-authored-by: Kyle Sayers <kylesayrs@gmail.com> Signed-off-by: Roderick Wu <roderickwu2003@gmail.com>
|
The quality checks have failed. Please run |
4 participants
Improve the use of mse observer for NVFP4.
-- Grid search expansion factor added as observer kwargs -- there are settings that enable a search that is equivalent to a superset of fouroversix. By expanding the grid search range, we cover fouroversix as well as potentially better ranges.
-- Adds new nvfp4 observer with good defaults, which also disables rounding during grid search.
fouroversix (https://arxiv.org/abs/2512.02010) chooses between two quantization ranges: the full absmax or absmax*1.5. The reasoning is that in fp4 the largest values jump from 4 --> 6 which is a significatn gap compared to previous increments. By using absmax*1.5, we are limiting the largest values in a block to be quantized to 4. This uses fewer values and therefore loses granularity but in exchange ensures that rounding error for large values will be smaller. This is better for some blocks.
We can go beyond this and start the grid search at a higher max -- beyond the maximum observed value -- and gradually shrink the scale. This includes searching over the two points used in fouroversix, as well as any scale range between them.
We also include this expand kwarg for the imatrix observer. Using the same default mse params, we find that it performs slightly worse but within reasonable error. Tuning imatrix parameters with the expanded search can be done later. Both of them outperform gptq. Both imatrix and gptq use ultrachat dataset.
Justification
Starting at very high range and shrinking down, empirically observe all blocks follow this distribution:
Register new observer
"nvfp4_mse"with defaultsFouroversix does set the global scale beforehand. The mse observer performs grid search without a global scale. We consider using a global-scale prior guess calculated from observed absmax values. When we include good guesses for the global scale ahead of time we can also improve over fouroversix.
These are improvements over fouroversix and existing mse implementation. However, still worse than just dropping global scale from grid search and disabling roudning.
Code for tests + compression and eval scripts:
#3042