mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-07 08:54:02 +08:00
Frida Hou
18d9b1eea4
## What does this PR do?
**Type of change:** ? <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->
new feature
**Overview:** ?
Support static block-wise MSE for NVFP4 weight quantization.
Add a FP4 triton kernel that take in scales for each block. It also
quantizes the scales to FP8.
This PR does the following:
1. Enable static NVFP4 implementation, i.e. block scales for weights are
calculated during calibration and feed into fake quant kernels
2.Extend mse_calibrate to support static NVFP4 with block scales
searching by MSE and global scale set as MAX
3.Refinements: calibrate weight quantizers only once during MSE
calibration
## Usage
<!-- You can potentially add a usage example below. -->
Example config:
```python
NVFP4_WEIGHT_MSE_CFG = {
"quant_cfg": {
"*weight_quantizer": {
"num_bits": (2, 1),
"block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)},
"axis": None,
"enable": True,
},
"*input_quantizer": {
"enable": False,
},
**_default_disabled_quantizer_cfg,
},
"algorithm": {
"method": "mse",
"step_size": 0.25,
"start_multiplier": 0.25,
"stop_multiplier": 2.0,
},
}
NVFP4_WEIGHT_ACT_MSE_CFG = {
"quant_cfg": {
"*weight_quantizer": {
"num_bits": (2, 1),
"block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)},
"axis": None,
"enable": True,
},
"*input_quantizer": {
"num_bits": (2, 1),
"block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
"axis": None,
"enable": True,
},
**_default_disabled_quantizer_cfg,
},
"algorithm": {
"method": "mse",
"step_size": 0.25,
"start_multiplier": 0.25,
"stop_multiplier": 2.0,
},
}
```
## Testing
<!-- Mention how have you tested your change if applicable. -->
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->
## Additional Information
<!-- E.g. related issue. -->
---------
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
115 lines
3.7 KiB
Python
115 lines
3.7 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""High-level tests for quantization."""
|
|
|
|
import pytest
|
|
from _test_utils.torch.quantization.models import SimpleConv, SimpleConvLinear, SimpleLinear
|
|
from _test_utils.torch.quantization.quantize_common import (
|
|
FP4_SVDQUANT_CFG,
|
|
INT4_AWQ_CLIP_CFG,
|
|
INT4_AWQ_FULL_CFG,
|
|
quantize_model_and_forward,
|
|
save_restore_test,
|
|
)
|
|
|
|
import modelopt.torch.quantization as mtq
|
|
from modelopt.torch.quantization.extensions import get_cuda_ext_mx
|
|
|
|
NVFP4_WEIGHT_ACT_MSE_CFG = {
|
|
"quant_cfg": {
|
|
"*weight_quantizer": {
|
|
"num_bits": (2, 1),
|
|
"block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)},
|
|
"axis": None,
|
|
"enable": True,
|
|
},
|
|
"*input_quantizer": {
|
|
"num_bits": (2, 1),
|
|
"block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
|
|
"axis": None,
|
|
"enable": True,
|
|
},
|
|
},
|
|
"algorithm": "mse",
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("model_cls", [SimpleLinear, SimpleConv, SimpleConvLinear])
|
|
@pytest.mark.parametrize(
|
|
"config",
|
|
[
|
|
mtq.INT8_DEFAULT_CFG,
|
|
mtq.FP8_DEFAULT_CFG,
|
|
mtq.W4A8_AWQ_BETA_CFG,
|
|
mtq.INT8_SMOOTHQUANT_CFG,
|
|
mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG,
|
|
mtq.INT4_AWQ_CFG,
|
|
INT4_AWQ_CLIP_CFG,
|
|
INT4_AWQ_FULL_CFG,
|
|
mtq.NVFP4_DEFAULT_CFG,
|
|
FP4_SVDQUANT_CFG,
|
|
mtq.NVFP4_AWQ_LITE_CFG,
|
|
mtq.NVFP4_AWQ_CLIP_CFG,
|
|
mtq.NVFP4_AWQ_FULL_CFG,
|
|
mtq.MXFP8_DEFAULT_CFG,
|
|
mtq.MXFP6_DEFAULT_CFG,
|
|
mtq.MXFP4_DEFAULT_CFG,
|
|
mtq.MXINT8_DEFAULT_CFG,
|
|
mtq.NVFP4_KV_ROTATE_CFG,
|
|
mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG,
|
|
NVFP4_WEIGHT_ACT_MSE_CFG,
|
|
],
|
|
)
|
|
def test_quantize(model_cls, config):
|
|
"""Test quantize function can run without problems."""
|
|
if config in [
|
|
mtq.NVFP4_DEFAULT_CFG,
|
|
FP4_SVDQUANT_CFG,
|
|
mtq.NVFP4_AWQ_LITE_CFG,
|
|
mtq.NVFP4_AWQ_CLIP_CFG,
|
|
mtq.NVFP4_AWQ_FULL_CFG,
|
|
mtq.MXFP8_DEFAULT_CFG,
|
|
mtq.MXFP6_DEFAULT_CFG,
|
|
mtq.MXFP4_DEFAULT_CFG,
|
|
mtq.MXINT8_DEFAULT_CFG,
|
|
mtq.NVFP4_KV_ROTATE_CFG,
|
|
mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG,
|
|
NVFP4_WEIGHT_ACT_MSE_CFG,
|
|
]:
|
|
if get_cuda_ext_mx() is None:
|
|
pytest.skip("cuda_ext_mx is not available")
|
|
if model_cls in [SimpleConv, SimpleConvLinear]:
|
|
pytest.skip("Conv weight quantization will fail as the kernel_size < FP4 blocksize")
|
|
|
|
if config == mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG:
|
|
# reduce block sizes for simple testing models
|
|
config["quant_cfg"]["*weight_quantizer"]["block_sizes"] = {-1: 8, -2: 8}
|
|
model = model_cls().cuda()
|
|
calib_data = [model.get_input().cuda() for _ in range(8)]
|
|
quantize_model_and_forward(model, config, calib_data)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("model_cls", "quant_config"),
|
|
[
|
|
(SimpleLinear, mtq.INT8_SMOOTHQUANT_CFG),
|
|
(SimpleLinear, mtq.W4A8_AWQ_BETA_CFG),
|
|
(SimpleConvLinear, mtq.INT8_DEFAULT_CFG),
|
|
],
|
|
)
|
|
def test_save_restore(model_cls, quant_config):
|
|
save_restore_test(model_cls, "cuda", quant_config)
|