【大模型学习 | 量化】pytorch量化基础知识（1）

pytorch量化

[!note]

• 官方定义：performing computations and storing tensors at lower bitwidths than floating point precision.
• 支持INT8量化，可以降低4倍的模型大小以及显存需求，加速2-4倍的推理速度
• 通俗理解：降低权重和激活值的精度（FP32→INT8），从而提高模型大小以及显存需求。

一、前置知识

1.1 算子融合

​ 将多个连续层的计算操作合并为单个复合算子，减少对内存的访问次数

e.g. 例如将Conv → BN → ReLU, 融合为ConvBnReLU

​ NVIDA GPU:

// 未融合：多次启动核函数
conv_kernel<<<...>>>(input, weight, temp1);
bias_kernel<<<...>>>(temp1, bias, temp2);
relu_kernel<<<...>>>(temp2, output);

// 已融合：单核函数完成所有操作
fused_kernel<<<...>>>(input, weight, bias, output) {
    float val = conv2d(input, weight);
    val += bias;
    output = max(val, 0.0f);
}

二、量化知识

2.1 对称量化 & 非对称量化

⚙️ 区别

• 对称量化（Symmetric Quantization）

$$X_{int}=round(\frac{X_{float}}{scale}),scale=\frac{max(|X|)}{2^{n-1}-1}$$

• 非对称量化（Affine Quantization）

$$X_{int}=round(\frac{X_{float}}{scale})+zero\_point,scale=\frac{ma{x}_x-mi{n}_{)}x}{2^n-1}$$

$$zero\_point=round(\frac{-min(x)}{scale})$$

🤖 工程实现角度：为什么 PTQ 常用非对称，QAT 用对称？

2.2 PTQ & QAT

[!note]

PTQ 是直接对训练后的模型参数进行量化，因此适合于快速部署；QAT是通过插入伪量化节点，在训练过程中模拟量化误差以达到更高的精度，因此需要重新训练。

⚙️ 区别

[!tip]

QAT伪量化节点

• 作用：在训练时模拟量化的误差。在每一层训练时，权重、激活值依然是FT32，但在每一层的传播中，值被“量化再还原”，模拟了量化过程。
• 由于量化过程有round函数，是不可微的，因此需要Straight-Through Estimator（STE）近似梯度的 FakeQuant 模块

三、Pytorch实现量化的三种方式

参考链接：Quantization — PyTorch 2.7 documentation

[!note]

无论是PTQ 还是 QAT ， 每一种实现方式都需要 prepare_fx 和 convert_fx

model_prepared = quantize_fx.prepare_fx(model, qconfig_mapping, example_inputs)
model_quantized = quantize_fx.convert_fx(model_prepared)

🎯 核心功能：在模型的每一个 qconfig_mapping 指定的量化位置（如 Conv2d、Linear）处，插入对应的 observer 或 fake_quant 节点。

📦 插入两类模块：

类型	对应 prepare 的用途	说明
Observer	用于 PTQ	统计 min/max 用来 校准计算 scale 和 zero_point
FakeQuantize	用于 QAT	模拟量化误差，保留梯度流动，支持训练

3.1 Eager Mode Quantization

import torch

# define a floating point model where some layers could benefit from QAT
class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # QuantStub converts tensors from floating point to quantized
        self.quant = torch.ao.quantization.QuantStub()
        self.conv = torch.nn.Conv2d(1, 1, 1)
        self.bn = torch.nn.BatchNorm2d(1)
        self.relu = torch.nn.ReLU()
        # DeQuantStub converts tensors from quantized to floating point
        self.dequant = torch.ao.quantization.DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.dequant(x)
        return x

# create a model instance
model_fp32 = M()

# model must be set to eval for fusion to work
model_fp32.eval()

# attach a global qconfig, which contains information about what kind
# of observers to attach. Use 'x86' for server inference and 'qnnpack'
# for mobile inference. Other quantization configurations such as selecting
# symmetric or asymmetric quantization and MinMax or L2Norm calibration techniques
# can be specified here.
# Note: the old 'fbgemm' is still available but 'x86' is the recommended default
# for server inference.
# model_fp32.qconfig = torch.ao.quantization.get_default_qconfig('fbgemm')
model_fp32.qconfig = torch.ao.quantization.get_default_qat_qconfig('x86')

# fuse the activations to preceding layers, where applicable
# this needs to be done manually depending on the model architecture
model_fp32_fused = torch.ao.quantization.fuse_modules(model_fp32,
    [['conv', 'bn', 'relu']])

# Prepare the model for QAT. This inserts observers and fake_quants in
# the model needs to be set to train for QAT logic to work
# the model that will observe weight and activation tensors during calibration.
model_fp32_prepared = torch.ao.quantization.prepare_qat(model_fp32_fused.train())

# run the training loop (not shown)
training_loop(model_fp32_prepared)

# Convert the observed model to a quantized model. This does several things:
# quantizes the weights, computes and stores the scale and bias value to be
# used with each activation tensor, fuses modules where appropriate,
# and replaces key operators with quantized implementations.
model_fp32_prepared.eval()
model_int8 = torch.ao.quantization.convert(model_fp32_prepared)

# run the model, relevant calculations will happen in int8
res = model_int8(input_fp32)

3.2 FX Graph Mode Quantization （maintenance）

import torch
from torch.ao.quantization import (
  get_default_qconfig_mapping,
  get_default_qat_qconfig_mapping,
  QConfigMapping,
)
import torch.ao.quantization.quantize_fx as quantize_fx
import copy

model_fp = UserModel()

#
# post training dynamic/weight_only quantization
#

# we need to deepcopy if we still want to keep model_fp unchanged after quantization since quantization apis change the input model
model_to_quantize = copy.deepcopy(model_fp)
model_to_quantize.eval()
qconfig_mapping = QConfigMapping().set_global(torch.ao.quantization.default_dynamic_qconfig)
# a tuple of one or more example inputs are needed to trace the model
example_inputs = (input_fp32)
# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, example_inputs)
# no calibration needed when we only have dynamic/weight_only quantization
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# post training static quantization
#

model_to_quantize = copy.deepcopy(model_fp)
qconfig_mapping = get_default_qconfig_mapping("qnnpack")
model_to_quantize.eval()
# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, example_inputs)
# calibrate (not shown)
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# quantization aware training for static quantization
#

model_to_quantize = copy.deepcopy(model_fp)
qconfig_mapping = get_default_qat_qconfig_mapping("qnnpack")
model_to_quantize.train()
# prepare
model_prepared = quantize_fx.prepare_qat_fx(model_to_quantize, qconfig_mapping, example_inputs)
# training loop (not shown)
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# fusion
#
model_to_quantize = copy.deepcopy(model_fp)
model_fused = quantize_fx.fuse_fx(model_to_quantize)

3.3 PyTorch 2 Export Quantization

import torch
from torch.ao.quantization.quantize_pt2e import prepare_pt2e
from torch.export import export_for_training
from torch.ao.quantization.quantizer import (
    XNNPACKQuantizer,
    get_symmetric_quantization_config,
)

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(5, 10)

   def forward(self, x):
       return self.linear(x)

# initialize a floating point model
float_model = M().eval()

# define calibration function
def calibrate(model, data_loader):
    model.eval()
    with torch.no_grad():
        for image, target in data_loader:
            model(image)

# Step 1. program capture
# NOTE: this API will be updated to torch.export API in the future, but the captured
# result should mostly stay the same
m = export_for_training(m, *example_inputs).module()
# we get a model with aten ops

# Step 2. quantization
# backend developer will write their own Quantizer and expose methods to allow
# users to express how they
# want the model to be quantized
quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config())
# or prepare_qat_pt2e for Quantization Aware Training
m = prepare_pt2e(m, quantizer)

# run calibration
# calibrate(m, sample_inference_data)
m = convert_pt2e(m)