量化概述

量化概述#

首先概述 torchao 的组件栈：

Quantization Algorithms/Flows: weight only/dynamic/static quantization, hqq, awq, gptq etc.
---------------------------------------------------------------------------------------------
        Quantized Tensors (derived dtypes): AffineQuantizedTensor, CodebookQuantizedTensor
---------------------------------------------------------------------------------------------
  Quantization Primitive Ops/Efficient Kernels: matmul, quantize, dequantize
---------------------------------------------------------------------------------------------
            Basic dtypes: uint1-uint7, int1-int8, float3-float8

任何量化算法都会使用上述组件栈中的某些组件，例如 int4 权重量化使用：

  1. 仅权重量化流程

  2. tinygemm bf16 激活 + int4 权重内核和量化原语算子

  3. 具有 TensorCoreTiledLayout 的 AffineQuantizedTensor 张量子类

  4. torch.uint4 数据类型（目前通过 quant_min/quant_max 模拟）

基础数据类型#

数据类型是一个有点过载的术语，所谓基础数据类型，指的是无需任何额外元数据即可有意义的数据类型（例如，当人们调用 torch.empty(.., dtype) 时就有意义），更多详情请查看：dev-discuss.pytorch.org/t/supporting-new-dtypes-in-pytorch/1833

无论进行何种量化，最终都会使用一些低精度数据类型来表示量化后的数据，torchao 旨在支持的数据类型有：

  • torch.uint1 到 torch.uint8 在 pytorch 2.3 及更高版本中可用

  • torch.int1 到 torch.int8 在 pytorch 2.6 及更高版本中可用

  • torch.float3_e2_m0, torch.float4_e2_m1, torch.float4_e3_m0, torch.float5_e2_m2, torch.float5_e3_m1, torch.float6_e2_m3, torch.float6_e3_m2, torch.float8_e4m3fn, torch.float8_e5m2, torch.float8_e4m3fnuz, torch.float8_e5m2fnuz (float8 is added to torch, we also plan to add float4 and float6 to torch if they become popular) torch.float3_e2_m0 、 torch.float4_e2_m1 、 torch.float4_e3_m0 、 torch.float5_e2_m2 、 torch.float5_e3_m1 、 torch.float6_e2_m3 、 torch.float6_e3_m2 、 torch.float8_e4m3fn 、 torch.float8_e5m2 、 torch.float8_e4m3fnuz 、 torch.float8_e5m2fnuz （float8 已加入 torch，我们计划在 float4 和 float6 变得流行时也将它们加入 torch）

量化原语算子#

量化原语算子指的是用于在低精度量化张量和高精度张量之间进行转换的算子。主要包含以下量化原语算子：选择量化参数的算子（choose_qparams ops）：根据原语张量选择量化参数，通常用于动态量化，例如仿射量化的缩放因子和零点量化算子（quantize op）：根据量化参数将原始高精度张量量化为前文提到的低精度张量反量化操作（dequantize op）：根据量化参数将低精度张量去量化为高精度张量

可能会有一些变化来适应特定的使用场景，例如对于静态量化，可能会有 choose_qparams_affine_with_min_max ，它将根据观察过程中得出的最小/最大值来选择量化参数。

高效内核#

还将拥有与低精度张量协同工作的有效内核，例如

  • _weight_int4pack_mm：微型 gemm int4 内核（bf16 激活 + int4 权重）int_matmul：接受两个 int8 张量并输出 int32 张量的整数矩阵乘法 int_scaled_matmul：执行矩阵乘法并同时对结果应用缩放。

注意：还可以依赖 torch.compile 通过 triton 生成内核，例如当前的 int8 权重量化内核仅依赖 torch.compile 来获得加速。在这种情况下，没有特定的“有效内核”对应于量化的类型。

量化张量（衍生数据类型）#

在基本数据类型（dtypes）、量化原语运算符和高效内核的基础上，我们可以将所有内容整合起来，通过继承torch.Tensor构建一个量化（低精度）张量，该张量可以从高精度张量和一些配置特定量化需求的参数构建，我们也可以称其为衍生数据类型，因为它可以用基本数据类型的张量和一些额外的元数据（如缩放比例）来表示。

torchao 中现有的示例是 AffineQuantizedTensor ，这意味着低精度张量是通过仿射映射从高精度张量量化的，即： low_precision_val = high_precision_val / scale + zero_point ，其中 scale / zero_point 是可以通过量化原语操作或通过某些优化程序计算的量化参数。仿射量化是一种非常常见的量化类型，因为当我们尝试将高精度值映射到低精度值时，进行仿射变换（ high_preicsion_val / scale + zero_point ）是很直观的。另一种常见的量化类型，特别是对于低位宽（例如低于 4 位）的是基于码本/查找表的量化。

布局和 TensorImpl#

原生张量有硬编码的布局选择列表，最常见的布局是 strided 布局，它提供了对存储的 strided、多维视图，还有一些稀疏和 mkldnn 布局。

以稀疏 COO 张量为例，它具有 torch.sparse_coo 布局，以及 SparseTensorImpl，后者改变了张量的存储方式。

将张量打包到不同格式的想法与布局概念非常契合，这就是我们想要重用这个打包方式的原因。我们可以使用 Layout 来处理不同类型的打包格式，使用 TensorImpl 来处理不同的存储格式实现。并且，可以在 Python 层的张量子类中添加新的 TensorImpl 来以打包格式存储张量，而无需修改 C++ pytorch 核心代码。

例如，对于 _weight_int4pack_mm ，我们需要将权重打包成 Tensor Core 友好的格式，我们称之为 TensorCoreTiledLayout。我们添加一个 tensor_impl 来存储打包（或解包）的权重，并使用 layout 来存储与打包相关的不同参数：

class AffineQuantizedTensor(...):
  # tensor_impl is also implemented with tensor subclass
  tensor_impl: torch.Tensor

  # to not conflict with existing layout property, we use `_layout`
  @property
  def _layout(self) -> Layout:
      return self.tensor_impl._layout

请注意，布局不仅是一种抽象，用于自定义数据表示，还用于 TensorImpl 与不同算子交互的方式。例如，相同的数据表示在运行相同算子时可能具有不同的实现，如转置、量化线性，但算子的语义应保持不变。

通过布局抽象，Quantize + Sparse Tensor 也可以得到支持，例如 int4 权重量化 + 稀疏。我们还提供了一些常用工具，帮助人们为量化张量添加不同的布局，请查看下面的开发者指南获取代码示例。

量化算法/流程#

在堆栈的顶部将是最终的量化算法和量化流程。传统上我们有仅权重量化、动态量化和静态量化，但现在我们也看到了更多类型的量化正在出现。

为了演示目的，假设在之前的步骤中我们定义了 AffineQuantizedTensor 和 to_affine_quantized 工厂函数。为了简化，假设 to_affine_quantized 接收一个高精度浮点张量和一个目标数据类型（例如 torch.int8），并将其转换为具有相应数据类型的 AffineQuantizedTensor 。

注意：以下内容均为概念解释，关于我们提供的工具和示例的更详细介绍可以在 Tensor Subclass Developer Guide 部分找到。

仅权重量化#

这是最简单的量化形式，将权重量化应用于模型非常容易，特别是我们已经有量化张量。只需要做以下操作：

linear_module.weight = torch.nn.Parameter(to_affine_quantized_intx(linear_module.weight, ...), requires_grad=False))

将上述方法应用于模型中的所有线性模块，我们就能得到一个仅权重量化的模型。

动态激活和权重量化#

这被称为“动态量化”，但它的意思是我们在运行时动态量化激活值，同时也量化权重。与仅量化权重的相比，主要问题是我们如何将量化应用于激活值。在 torchao 中，我们常用的模式是在量化权重之上应用 to_linear_activation_quantized ：

quantized_weight = to_affine_quantized(linear_module.weight) activation_and_weight_quantized = to_linear_activation_quantized(quantized_weight) linear_module.weight = torch.nn.Parameter(activation_and_weight_quantized, requires_grad=False)

to_linear_activation_quantized 用于对激活值应用量化，它接受一个 input_quant_func 来量化激活值和原始权重，在运行时当遇到 F.linear 操作时，会应用存储的输入量化函数到激活值，并重新调度到 F.linear 使用量化的激活值和权重。

如果上述方法无效，用户也可以进行模块替换，或者使用 torch.fx.symbolic_trace() 获取一个可以修改的跟踪模块。

但使用张量子类更受推荐，因为这样更便于序列化/反序列化。如果我们使用张量子类来支持动态量化，那么可以直接加载量化后的权重，而无需对模型进行进一步准备。否则，在加载量化后的权重之前，我们需要先对模型进行模块替换或其他修改。

静态激活量化与权重量化#

静态量化是指激活值在运行时不是动态量化，而是静态量化的。从流程上看，静态量化需要使用样本数据进行校准，以便我们可以确定合适的量化参数。

从高层次来看，静态量化有三个步骤：(1)插入观察者 (2)校准 (3)量化模型

插入观察者#

在插入观察者步骤中，需要向算子的输入（和输出）激活值和权重添加观察者模块，以收集张量的统计信息。因此，需要解决两个问题：如何定义观察者模块？如何将观察者模块添加到模型中。

如何定义观察者模块#

观察者特定于：(1) 量化类型（例如仿射量化、基于查找表的量化）(2) 我们想要跟踪的统计类型，例如最小最大观察者、移动平均观察者。

通常，观察者模块应该定义 forward 和 calculate_qparams

对于仿射量化，定义了 AffineQuantizedMinMaxObserver，它根据仿射量化的粒度记录 min_val/max_val，并定义了如何根据记录的统计信息计算 qparams。

如何向模型添加观察者模块#

如果你们感兴趣的唯一算子是线性算子，可以使用线性激活权重观察器，我们还提供了一个相应的 insert_observer_ API 来处理修改线性权重。

模块替换是另一种方法，你们也可以定义一个 ObservedLinear 模块（或其他模块类型），并将未观察的模块与观察的模块进行替换。

校准#

校准步骤通常很简单，通常我们只需要将模型运行通过校准数据集。对于更复杂的校准（例如，我们记录所有输入并根据所有输入进行优化），我们将在下一节中介绍其中一些内容。

量化#

我们可以重用 quantize_ API，但提供一个不同的 apply_tensor_subclass 函数，将观察到的线性模块转换为具有量化权重和静态量化输入激活的线性模块，这可以与动态量化（使用 to_linear_activation_quantized ）以相同的方式进行，请参考示例。

目录