mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-20 05:59:32 +08:00
Wei-Ming Chen
bc96f1ce39
### What does this PR do?
Type of change: New feature <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->
<!-- Details about the change. -->
By default, FP8 KV cache quantization in hf_ptq.py now uses a constant
scale of 1.0 (amax=448.0) without a data-driven calibration pass. No KV
scales are written to the exported checkpoint — inference engines
(TRT-LLM, vLLM) use scale=1.0 when no scale is
present, so this is lossless. Pass --calibrate_kv_cache to opt into
data-driven per-tensor KV scale calibration (previous default behavior).
To support this cleanly in the quantization stack, a constant_amax field
is added to QuantizerAttributeConfig. Quantizers configured with
constant_amax skip calibration entirely (no forward pass needed), use
the fixed amax during fake-quant, and produce no
_amax buffer in the state dict.
### Usage
```python
# Quantize with default constant KV scale (no calibration pass for KV)
python hf_ptq.py --model ... --qformat fp8 --kv_cache_qformat fp8
# Opt into data-driven KV calibration
python hf_ptq.py --model ... --qformat fp8 --kv_cache_qformat fp8 --calibrate_kv_cache
# Use constant_amax in a custom quant config
quant_cfg = {
"quant_cfg": {
"*[kv]_bmm_quantizer": {"num_bits": (4, 3), "enable": True, "constant_amax": 448.0},
},
"algorithm": "max",
}
model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
```
### Testing
<!-- Mention how have you tested your change if applicable. -->
- All 8 existing GPU HF export tests pass
(tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py)
- Two new CPU unit tests added to TensorQuantizerTester:
test_constant_amax and test_constant_amax_skips_calibration
### Before your PR is "*Ready for review*"
Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).
Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).
- Is this change backward compatible?: ✅ (--calibrate_kv_cache flag
defaults to False; existing scripts without the flag now skip KV
calibration)
- 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?: ✅
### Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* New CLI flag --calibrate_kv_cache to optionally enable data-driven KV
cache calibration; default uses a fixed KV scale and omits KV scales
from exported checkpoints.
* Added constant_amax option to set fixed quantizer scales and skip
dynamic calibration for configured quantizers.
* **Bug Fixes**
* Removed forced flooring/clamp of KV cache scales; out-of-range
activations now emit a shorter warning.
* **Tests**
* Added tests for constant_amax behavior and calibration interaction.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Model Optimizer End-to-End Tests
This directory contains end-to-end tests for the Model Optimizer's examples.
Adding new tests
To add a test for a new example, create a new example directory in tests/examples following the existing examples as guidance.
Make sure to use as small models and less data as possible to keep the tests fast. Unless needed, have tests finish under 15 minutes.
Running the tests
To run a test, start from the recommended docker image from our installation docs.
Then mount your local modelopt directory to /workspace/Model-Optimizer and run this from the root of the repository.
cd /workspace/Model-Optimizer
pip install -e ".[all,dev-test]"
pytest tests/examples/$TEST
Environment variables
The following environment variables can be set to control the behavior of the tests:
MODELOPT_LOCAL_MODEL_ROOT: If set, the tests will use the local model directory instead of downloading the model from the internet. Default is not set, which means the model will be downloaded.