使用NVIDIA modelopt代替pt-quantization进行模型量化并使用TRT导出INT8模型

Model-opt与pytorch-quantizaiton的关系

pytorch-quantizaiton

pytorch-quantizaiton是包含在TensorRT内的一个附属工具。由于torch原生的量化方式一直无法得到工业上的认可，并且trt作为编译方决定了int8模型如何在GPU上运行。

因此使用trt本身提供的量化工具库可以让开发人员更方便的遵循trt的量化规则。pt-quantization功能很简单，提供基础layer的伪量化算子并进行替换，并支持多种校准(calibration)策略的选取，这些伪量化算子支持训练、推理以及ONNX导出。在导出ONNX后可以很方便的使用trt的explicit量化模式将ONNX编译为GPU上高效运行的engine。

之前用trt时用过它自带的隐式INT8量化，隐式量化可操控性弱、精度掉点严重，并且在校准时还需要将数据输入trt，比较麻烦。pytorch-quantization可以看作是trt的显式量化工具，其在pytorch侧便可以插入量化节点，无需手动操控ONNX，当导出的ONNX包含QDQ即量化节点时，trt便会默认进行显式量化。

Modelopt

TensorRT-Model-Optimizer是NV在2024年新建立的一个提供最新模型优化策略工具库，其不仅支持量化，还支持剪枝(pruning)、蒸馏(distillation)、稀疏(sparse)以及推理编码(speculative decoding)。

并且model opt优化的主要目标是近年来较火的LLM模型，优化后的模型支持trt-llm推理，支持多种针对transformer结构的优化算法(都是近五年比较新的算法)。在量化层面，其允许最基本的int8量化，同时还支持FP8、INT4以及weights only或weight INT4 Activation FP8的混合量化策略。同时在数月前的doc内其明确标有Model Optimizer是pytorch-quantization的超集字样(在9月份文档更新后删除了该字段，理由未知)。

modelopt可以配合nvfp4使用，不过它优化后的模型一样支持vllm部署，nv在huggingface上有自己的一个仓库用来存放了一大批使用modelopt量化+校准后的llm模型，可以一键在使用trt-llm进行推理。

总结

如果量化目标是只简单的卷积、全连接、循环模型，不是包含attention的类transformer模型，并且部署目标是NVIDIA平台且希望显式的在trt侧进行量化控制，使用pytorch-quantization。其余情况都应该考虑使用Modelopt，并且nv官方还在持续更新modelopt中。

大致来看很显然model opt的功能要比pt-quantization强大的多，pt-quantization能做到的东西model opt理论应该都可以做到。在上个单位工作时因为只需要对一些视觉模型进行优化，一般只使用pt-quantization，很熟悉这个库的功能以及量化基本原理。

后来24年底知道了modelopt后尝试用modelopt做了一些剪枝和稀疏的优化，发现在api调用上modelopt也比pt-quantization更符合程序员的编程习惯，相比来说用起来更舒服一些(其实也不太舒服)。再后来25年年初尝试用了一下model opt的量化，感觉在最简单的INT8量化pipeline上，modelopt的功能又有一些欠缺，至少从官方给的文档来看还缺少很多工程领域实际需要的功能。可能是因为modelopt侧重点还是llm吧，基于卷积的视觉模型或许真的要退出历史舞台了。

在年中opt一轮更新之后，我决定再探索一下针对规模较小的视觉模型，如何使用Model-opt对其进行纯粹的INT8量化(显示PTQ随后进行QAT)，以此来决定是否直接可以使用model opt代替pt-quantization

INT8 Quantization Pipeline

在工作中，当上游给到一个Pytorch模型(其实是由MMCV搭出来的)，我需要对该模型进行INT8的显式量化。由于上游给到的模型基本都是视觉(或者雷达多模态融合后)的以卷积为基础的小规模模型，模型内的大多数算子都有成熟的可替换的INT8量化算子，因此量化流程基本上包含以下几步：

  • 1; 测试模型基线性能：在X86平台使用Pytorch原生测试推理精度，在ARM平台(实际模型运行平台)导出FP16 engine测试推理速度；
  • 2; 粗放式的替换量化算子：对模型内的大部分算子(推荐替换的)进行替换，替换为INT8算子，进行校准后导出为ONNX；
  • 3; 观察INT8执行情况：粗放式的替换应该会导致模型在X86 PT原生侧就产生大幅度精度掉点，同时仅进行算子替换还是无法保证trt engine内该走INT8的部分都走INT8(此处的应该指的是对于替换了INT8算子的部分, 在trt graph内都应该全部走INT8精度, 不应该存在FP16精度混合);
  • 4; 添加额外INT8量化节点、修改模型forword链路：针对trt INT8 QDQ传播策略，修改short cut、split等通路forward链路，增加额外量化节点，观察trt graph，使其INT8传播至最广；
  • 5; 自动化测试敏感层：关闭所有量化节点，随后遍历式的打开单一节点进行精度测试，寻找影响精度最大的某些节点，同时可以对校准方式与校准集进行小幅度修改以寻找最佳策略；
  • 6; 重复测试模型性能：重复测试X86 PT侧量化后模型精度、ARM平台模型推理耗时，寻找一个最佳平衡点；
  • 7; 输出模型与测试：在ARM平台对模型进行大规模压测，若不满足精度要求则需回到步骤6再次寻找平衡点；
  • 8; 进行QAT：若时间允许，则可以对步骤7得到的量化后模型进行微调训练，进一步提高模型精度。

从上述的pipeline可以看到，真正使用量化工具库的部分只有替换算子的步骤，大部分耗时都在敏感层测试、修改forward链路上。因此，对于一个量化工具库，最好能满足：

  • 支持PT侧伪量化: 这样可以快捷方便的在PT侧进行精度测试并对量化层进行打开与关闭；
  • 支持ONNX QDQ输出: 在输出至ONNX时不能直接使用伪量化流程，应该输出为ONNX自带的Quantization与Dequantization节点；
  • 支持多种校准策略: 校准会严重影响量化精度，寻找合适的校准策略以及参数是量化过程中关键的一环，如果工具库可以提供多种校准策略或提供方便的校准参数修改API，可以大幅提高量化速度；
  • 支持PT侧STE(Straight-Through Estimator)训练: 直通估计是量化层进行后退传播的原理，量化工具库内置的量化算子应支持STE后退传播并可以直接进行QAT微调;
  • API使用方便快捷: 作为开发人员，当然希望使用的工具库API能够对pytorch/mmcv等框架库有较好支持。

事实上，pt-quantization基本可以满足上述所有的条件，抛开API使用便捷度不谈，其完全支持INT8的对称量化(TRT8版本仅支持INT8对称量化)，支持基于histogram的percentile/max/mse校准方式，支持伪量化节点到ONNX量化算子的导出，支持在校准后固定amax并进行QAT。理论上model-opt作为超集，也应该满足上述所有条件。

训练后量化INT8 PTQ

pytorch-quantization

pt-quantization官方文档给出的示例是使用quant_modules隐式的对模型进行量化, 在调用quant_modules.initialize()方法之后, pt-quantzation会动态的为后续加载的model进行量化替换(Dynamic module replacement using monkey patching.).
此处的量化替换是对模型内的算子进行了其对应的量化版本替换, 例如对于普通的conv2d, pt-quantization会将其替换为QuantConv2d, 对conv2d添加一个weight quantizer与input quantizer.

个人并不推荐使用文档给出的方法进行量化, 此种monkey patching的做法不太受程序员控制, 如果我们想控制什么时候我们可以得到量化后的模型, 就需要考虑自己创建一些简单的api. 参考代码部分参考了这篇CSDN博客。很神奇的是这篇文章里的替换量化算子的代码在上家工作的公司里见过几乎一模一样的，难以确定是前同事抄来的还是某个开源库内的，网上是搜不到其他类似的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def transfer_torch_to_quantization(nninstace : torch.nn.Module, quantmodule) -> torch.nn.Module:
    """
    Convert a input nn module to quantized version.
    e.g,    a conv2d -> quant_conv2d
            a full quant_conv2d include:                |-weight_quantizer_per_channel 
                                                        v
            inputs -> input_quantizer_per_tensor -> conv2d -> Relu ...
    so we need to determine whether the input module has weight or not.

    Args:
        nninstance: input torch module
        quantmodule: the replaced module
    Return:
        the quantized module
    """
    quant_instance = quantmodule.__new__(quantmodule)
    for k, val in vars(nninstace).items():
        setattr(quant_instance, k, val) 
    
    def __init__(self):
        if isinstance(self, quant_nn_utils.QuantInputMixin):
            quant_desc_input = quant_nn_utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True)
            self.init_quantizer(quant_desc_input)
        else:
            quant_desc_input, quant_desc_weight = quant_nn_utils.pop_quant_desc_in_kwargs(self.__class__)
            self.init_quantizer(quant_desc_input, quant_desc_weight)
    __init__(quant_instance)
    return quant_instance


def replace_to_quantization_module(model : torch.nn.Module) -> torch.nn.Module:
    """
    Recursively replace all the layer into quantized version based 
    on the quant_modules._DEFAULT_QUANT_MAP in a whole Module.

    Warning: All the layer in the quant_map would be replaced by the quantized version.

    Args:
        model: input nn module, could be a layer trees contains multiple submodules.
    Return:
        quantized model.
    """
    target = copy.deepcopy(model)
    module_dict = {}
    for entry in quant_modules._DEFAULT_QUANT_MAP:
        module = getattr(entry.orig_mod, entry.mod_name)
        module_dict[id(module)] = entry.replace_mod
    
    def recursive_and_replace_module(module, prefix=""):
        for name in module._modules:
            submodule = module._modules[name]
            path      = name if prefix == "" else prefix + "." + name
            recursive_and_replace_module(submodule, path)

            submodule_id = id(type(submodule))
            if submodule_id in module_dict:
                module._modules[name] = transfer_torch_to_quantization(submodule, module_dict[submodule_id])
        
    recursive_and_replace_module(target)
    return target

我们可以参考一下_DEFAULT_QUANT_MAP的内容:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
_DEFAULT_QUANT_MAP = [
    _quant_entry(torch.nn, "Conv1d", quant_nn.QuantConv1d),
    _quant_entry(torch.nn, "Conv2d", quant_nn.QuantConv2d),
    _quant_entry(torch.nn, "Conv3d", quant_nn.QuantConv3d),
    _quant_entry(torch.nn, "ConvTranspose1d", quant_nn.QuantConvTranspose1d),
    _quant_entry(torch.nn, "ConvTranspose2d", quant_nn.QuantConvTranspose2d),
    _quant_entry(torch.nn, "ConvTranspose3d", quant_nn.QuantConvTranspose3d),
    _quant_entry(torch.nn, "Linear", quant_nn.QuantLinear),
    _quant_entry(torch.nn, "LSTM", quant_nn.QuantLSTM),
    _quant_entry(torch.nn, "LSTMCell", quant_nn.QuantLSTMCell),
    _quant_entry(torch.nn, "AvgPool1d", quant_nn.QuantAvgPool1d),
    _quant_entry(torch.nn, "AvgPool2d", quant_nn.QuantAvgPool2d),
    _quant_entry(torch.nn, "AvgPool3d", quant_nn.QuantAvgPool3d),
    _quant_entry(torch.nn, "AdaptiveAvgPool1d", quant_nn.QuantAdaptiveAvgPool1d),
    _quant_entry(torch.nn, "AdaptiveAvgPool2d", quant_nn.QuantAdaptiveAvgPool2d),
    _quant_entry(torch.nn, "AdaptiveAvgPool3d", quant_nn.QuantAdaptiveAvgPool3d),
]

可见pt-quant为上述这些层做了量化映射，可见基本上没什么复杂的layer，都是一些比较基础的层。如果我们使用attention结构那pt-quant肯定是不会帮我们做量化的。如果我们自定义了某些量化层，也可以通过编辑这个MAP并使用上述API来达成替换的目的。

其中原博客提到的ConvTrans2d的量化问题在新版pt-quant内是解决了的，在使用反卷积时务必确保pt-quant的版本是最新的。

最后print一下量化结果:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
 (conv2): QuantConv2d(
        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
      )
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv3): QuantConv2d(
        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
      )
      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
    )
  )
  (layer2): Sequential(
    (0): Bottleneck(
      (conv1): QuantConv2d(
        256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False
...

Model-Opt

Model-opt的量化API更简单易懂, 但问题在于目前官方给出的文档内没有描述如何修改校准算法，也没有给出示例如何自定义在模型内添加quantizer(而不是简单的对某些算子进行量化版本的替换)
对于一个简单的模型, 直接使用default config对其内算子进行量化版本替换的代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

resnet50 = torchvision.models.resnet50(weights="ResNet50_Weights.IMAGENET1K_V1")
resnet50.eval()

torch.onnx.export(resnet50, torch.randn(1, 3, 224, 224), "resnet50.onnx")

data_loader = torch.utils.data.DataLoader(
    DummyDataset(num_samples=100, input_size=(3, 224, 224)),
    batch_size=1,
    shuffle=False,
)

config = mtq.INT8_DEFAULT_CFG.copy()

def fowrward_loop(model):
    for data, _ in data_loader:
        data = data.to("cuda:0")
        model(data)

resnet50.to("cuda:0")

quantized_model = mtq.quantize(resnet50, config, fowrward_loop)

可以看一下INT8_DEFAULT_CFG的内容:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
_default_disabled_quantizer_cfg = {
    "nn.BatchNorm1d": {"*": {"enable": False}},
    "nn.BatchNorm2d": {"*": {"enable": False}},
    "nn.BatchNorm3d": {"*": {"enable": False}},
    "nn.LeakyReLU": {"*": {"enable": False}},
    "*lm_head*": {"enable": False},
    "*proj_out.*": {"enable": False},  # In Whisper model, lm_head has key name proj_out
    "*block_sparse_moe.gate*": {"enable": False},  # Skip the MOE router
    "*router*": {"enable": False},  # Skip the MOE router
    "*mlp.gate.*": {"enable": False},  # Skip the MOE router
    "*mlp.shared_expert_gate.*": {"enable": False},  # Skip the MOE router
    "*output_layer*": {"enable": False},
    "output.*": {"enable": False},
    "default": {"enable": False},
}

INT8_DEFAULT_CFG = {
    "quant_cfg": {
        "*weight_quantizer": {"num_bits": 8, "axis": 0},
        "*input_quantizer": {"num_bits": 8, "axis": None},
        **_default_disabled_quantizer_cfg,
    },
    "algorithm": "max",
}

可见, 默认的cfg disable0了一些算子的量化, 但实际上还是会对算子进行量化版本替换(只不过不开启)。对于基本的CNN网络来说, convNd/transconvNd/liner/pooling层都会被替换为对应的quant版本并开启。

以resnet的初始层为例, 由Model-opt默认量化cfg得到的量化模型结构如下所示. conv和pooling都进行了量化，每个算子都有一个input侧的per tensor quantizer，而weight是是用了per channel的quantizer.
至于BN层，虽然也替换为了QuantBN但input quantizer并未开启。 至于BN层不开启的原因，在于BN层可以和其前置的conv层进行算子融合, 如果BN总是伴随着conv出现(CNN网络内大部分情况都是如此), 在inference的时候全都可以和conv进行算子融合从而导致网络中根本没有BN.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
ResNet(
  (conv1): QuantConv2d(
    3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
    (input_quantizer): TensorQuantizer(8 bit fake per-tensor amax=5.3670 calibrator=MaxCalibrator quant)
    (output_quantizer): TensorQuantizer(disabled)
    (weight_quantizer): TensorQuantizer(8 bit fake axis=0 amax=[0.0000, 0.7817](64) calibrator=MaxCalibrator quant)
  )
  (bn1): QuantBatchNorm2d(
    64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True
    (input_quantizer): TensorQuantizer(disabled)
    (output_quantizer): TensorQuantizer(disabled)
  )
  (relu): ReLU(inplace=True)
  (maxpool): QuantMaxPool2d(
    kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False
    (input_quantizer): TensorQuantizer(8 bit fake per-tensor amax=7.6515 calibrator=MaxCalibrator quant)
    (output_quantizer): TensorQuantizer(disabled)
  )

如果我们要自定义某些层不开启量化, 可以copy其cfg并覆写规则即可.

敏感层分析

量化感知训练INT8 QAT

在得到精度可以接受的PTQ模型后, 如果时间允许, 最好还是进行QAT. 使用一部分原始训练数据对量化后的模型进行微调, 可以显著提高模型的精度(尤其是对于视觉模型的部分B-box参数). 哪怕PTQ后的精度已经达标.
QAT之前, 可以考虑对PTQ模型对如下微调:

  • 固定BN层(freeze BN layer): 固定BN内的统计信息(即固定mean和bias, 但gamma和weight依旧在训练过程中更新), 可以加快微调收敛速度, 防止过拟合. 不过在数据量足够多、训练轮次足够多的条件下也不一定非要固定BN.
  • 固定quantizer的amax: 在学习过程中一般不开启amax的range动态学习功能, 这个取决于量化工具库如何实现量化算子.
  • 蒸馏: 条件允许可以考虑做蒸馏学习, 加快收敛速度.

在得到PTQ模型后, 先对BN层进行freeze, 代码如下(参考这篇博客):

1
2
3
4
5
6
7
8
def freeze_batch_norm_layer(input_model: torch.nn.Module):
    """ We just don't need bn to run statistics. The weight and gamma value in BN would be updated during the trainning stage."""
    for m in input_model.modules():
        if isinstance(m, (nn.BatchNorm2d)):
            m.eval()


freeze_batch_norm_layer(calibrated_resnet50_model)

对于使用pytorch-quantization或者model-opt库的我们来说, 这时候就可以直接开始训练了. 因为这两个库已经在训练过程中固定了amax, 不会在训练过程中对amax进行调整.

编程随笔 > AI相关
#Artificial Intelligence
使用NVIDIA modelopt代替pt-quantization进行模型量化并使用TRT导出INT8模型
https://blog.bakeneko-kuro.com/2025/12/27/hpc/using-model-opt-do-quantization/
作者
迷途黑猫
发布于
2025年12月27日
许可协议