静态量化#
参考：静态量化
静态量化是指在推理或生成过程中使用固定的量化范围。与动态量化不同，动态量化会为每个新的输入批次动态计算新的量化范围，静态量化通常会导致更高效的计算，但可能会牺牲量化精度，因为无法实时适应输入分布的变化。
在静态量化中，这个固定的量化范围通常会在量化模型之前，在类似输入上进行校准。在校准阶段，首先在模型中插入观察者，以“观察”要量化的输入的分布，并使用该分布来决定在量化模型时最终使用哪些缩放因子和零点。
通过示例来说明如何在 torchao 中实现这一点。所有代码都可以在示例脚本中找到。从玩具线性模型开始：
import copy
import torch
class ToyLinearModel(torch.nn.Module):
def __init__(self, m=64, n=32, k=64):
super().__init__()
self.linear1 = torch.nn.Linear(m, k, bias=False)
self.linear2 = torch.nn.Linear(k, n, bias=False)
def example_inputs(self, batch_size=1, dtype=torch.float32, device="cpu"):
return (
torch.randn(
batch_size, self.linear1.in_features, dtype=dtype, device=device
),
)
def forward(self, x):
x = self.linear1(x)
x = self.linear2(x)
return x
dtype = torch.bfloat16
m = ToyLinearModel().eval().to(dtype).to("cuda")
m = torch.compile(m, mode="max-autotune")
校准阶段#
torchao 随附了简单的观察者实现 AffineQuantizedMinMaxObserver，该观察者在校准阶段记录了流经观察者的最小值和最大值。用户可以实现自己所需的更高级的观察技术，例如依赖于移动平均或直方图的方法，并且这些方法将来可以添加到 torchao 中。
from torchao.quantization.granularity import PerAxis, PerTensor
from torchao.quantization.observer import AffineQuantizedMinMaxObserver
from torchao.quantization.quant_primitives import MappingType
# per tensor input activation asymmetric quantization
act_obs = AffineQuantizedMinMaxObserver(
MappingType.ASYMMETRIC,
torch.uint8,
granularity=PerTensor(),
eps=torch.finfo(torch.float32).eps,
scale_dtype=torch.float32,
zero_point_dtype=torch.float32,
)
# per channel weight asymmetric quantization
weight_obs = AffineQuantizedMinMaxObserver(
MappingType.ASYMMETRIC,
torch.uint8,
granularity=PerAxis(axis=0),
eps=torch.finfo(torch.float32).eps,
scale_dtype=torch.float32,
zero_point_dtype=torch.float32,
)
接下来，定义被观测的线性层，将用它替换 torch.nn.Linear。这是高精度（例如 fp32）的线性模块，并且在此模块中插入了上述观察者，用于在校准期间记录输入激活值和权重值。
import torch.nn.functional as F
class ObservedLinear(torch.nn.Linear):
def __init__(
self,
in_features: int,
out_features: int,
act_obs: torch.nn.Module,
weight_obs: torch.nn.Module,
bias: bool = True,
device=None,
dtype=None,
):
super().__init__(in_features, out_features, bias, device, dtype)
self.act_obs = act_obs
self.weight_obs = weight_obs
def forward(self, input: torch.Tensor):
observed_input = self.act_obs(input)
observed_weight = self.weight_obs(self.weight)
return F.linear(observed_input, observed_weight, self.bias)
@classmethod
def from_float(cls, float_linear, act_obs, weight_obs):
observed_linear = cls(
float_linear.in_features,
float_linear.out_features,
act_obs,
weight_obs,
False,
device=float_linear.weight.device,
dtype=float_linear.weight.dtype,
)
observed_linear.weight = float_linear.weight
observed_linear.bias = float_linear.bias
return observed_linear
要在玩具模型中实际插入这些观察者：
from torchao.quantization.quant_api import (
_replace_with_custom_fn_if_matches_filter,
)
def insert_observers_(model, act_obs, weight_obs):
_is_linear = lambda m, fqn: isinstance(m, torch.nn.Linear)
def replacement_fn(m):
copied_act_obs = copy.deepcopy(act_obs)
copied_weight_obs = copy.deepcopy(weight_obs)
return ObservedLinear.from_float(m, copied_act_obs, copied_weight_obs)
_replace_with_custom_fn_if_matches_filter(model, replacement_fn, _is_linear)
insert_observers_(m, act_obs, weight_obs)
现在校准模型，这将在校准过程中记录统计数据的观察器填充数据。可以通过向“观察”模型输入一些示例数据来完成这一过程：
for _ in range(10):
example_inputs = m.example_inputs(dtype=dtype, device="cuda")
m(*example_inputs)
/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/torch/_inductor/cudagraph_trees.py:2442: UserWarning: Unable to hit fast path of CUDAGraphs because of pending, uninvoked backwards. Consider running with torch.no_grad() or using torch.compiler.cudagraph_mark_step_begin() before each model invocation
warnings.warn(
量化阶段#
有多种方式进行模型量化。介绍一种更简单的替代方法，即定义 QuantizedLinear 类，将用这个类替换掉 ObservedLinear。定义这个类并不是必要的。对于另一种方法，可以使用 torch.nn.Linear 类，请参阅示例类。
from torchao.dtypes import to_affine_quantized_intx_static
class QuantizedLinear(torch.nn.Module):
def __init__(
self,
in_features: int,
out_features: int,
act_obs: torch.nn.Module,
weight_obs: torch.nn.Module,
weight: torch.Tensor,
bias: torch.Tensor,
target_dtype: torch.dtype,
):
super().__init__()
self.act_scale, self.act_zero_point = act_obs.calculate_qparams()
weight_scale, weight_zero_point = weight_obs.calculate_qparams()
assert weight.dim() == 2
block_size = (1, weight.shape[1])
self.target_dtype = target_dtype
self.bias = bias
self.qweight = to_affine_quantized_intx_static(
weight, weight_scale, weight_zero_point, block_size, self.target_dtype
)
def forward(self, input: torch.Tensor):
block_size = input.shape
qinput = to_affine_quantized_intx_static(
input,
self.act_scale,
self.act_zero_point,
block_size,
self.target_dtype,
)
return F.linear(qinput, self.qweight, self.bias)
@classmethod
def from_observed(cls, observed_linear, target_dtype):
quantized_linear = cls(
observed_linear.in_features,
observed_linear.out_features,
observed_linear.act_obs,
observed_linear.weight_obs,
observed_linear.weight,
observed_linear.bias,
target_dtype,
)
return quantized_linear
这个线性类在开始时计算输入激活和权重的缩放因子和零点，从而在未来的前向计算中固定量化范围。现在，可以定义以下配置并将其传递给 torchao 的 主 quantize_ API，以实际对模型进行量化：
from dataclasses import dataclass
from torchao.core.config import AOBaseConfig
from torchao.quantization import quantize_
from torchao.quantization.transform_module import (
register_quantize_module_handler,
)
@dataclass
class StaticQuantConfig(AOBaseConfig):
target_dtype: torch.dtype
@register_quantize_module_handler(StaticQuantConfig)
def _apply_static_quant(
module: torch.nn.Module,
config: StaticQuantConfig,
):
"""
Define a transformation associated with `StaticQuantConfig`.
This is called by `quantize_`, not by the user directly.
"""
return QuantizedLinear.from_observed(module, config.target_dtype)
# filter function to identify which modules to swap
is_observed_linear = lambda m, fqn: isinstance(m, ObservedLinear)
# perform static quantization
quantize_(m, StaticQuantConfig(torch.uint8), is_observed_linear)
现在，将看到模型中的线性层被替换为 QuantizedLinear 类，并且输入激活的量化缩放因子和量化权重都是固定的
m
OptimizedModule(
(_orig_mod): ToyLinearModel(
(linear1): QuantizedLinear()
(linear2): QuantizedLinear()
)
)
m.linear1.act_scale
tensor([0.0221], device='cuda:0')
m.linear1.qweight
AffineQuantizedTensor(tensor_impl=PlainAQTTensorImpl(data=tensor([[ 6, 216, 194, ..., 254, 3, 124],
[ 60, 58, 240, ..., 76, 104, 209],
[ 63, 47, 2, ..., 73, 181, 199],
...,
[127, 111, 94, ..., 71, 253, 36],
[186, 252, 137, ..., 200, 29, 55],
[242, 217, 169, ..., 162, 119, 186]], device='cuda:0',
dtype=torch.uint8)... , scale=tensor([0.0009, 0.0010, 0.0010, 0.0010, 0.0009, 0.0010, 0.0009, 0.0010, 0.0010,
0.0010, 0.0009, 0.0009, 0.0010, 0.0009, 0.0009, 0.0010, 0.0010, 0.0009,
0.0009, 0.0010, 0.0009, 0.0010, 0.0009, 0.0010, 0.0009, 0.0010, 0.0009,
0.0009, 0.0010, 0.0009, 0.0010, 0.0009, 0.0009, 0.0010, 0.0009, 0.0010,
0.0009, 0.0009, 0.0010, 0.0009, 0.0010, 0.0010, 0.0010, 0.0009, 0.0010,
0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010,
0.0009, 0.0010, 0.0009, 0.0010, 0.0010, 0.0009, 0.0010, 0.0010, 0.0010,
0.0010], device='cuda:0')... , zero_point=tensor([121., 131., 127., 126., 124., 127., 129., 129., 126., 125., 132., 131.,
128., 127., 126., 129., 128., 131., 127., 128., 123., 129., 132., 124.,
132., 129., 131., 126., 129., 126., 129., 125., 126., 129., 128., 128.,
130., 125., 126., 129., 128., 128., 127., 124., 129., 127., 129., 128.,
128., 129., 126., 128., 129., 126., 134., 129., 124., 128., 130., 124.,
128., 129., 128., 128.], device='cuda:0')... , _layout=PlainLayout()), block_size=(1, 64), shape=torch.Size([64, 64]), device=cuda:0, dtype=torch.bfloat16, requires_grad=False)