量化概述

首先，我们想介绍一下 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_weight_only 量化使用：(1) 仅权重量化流程 (2) tinygemm bf16 激活 + int4 权重核 和 量化原语操作 (3) AffineQuantizedTensor 张量子类以及 TensorCoreTiledLayout (4) torch.uint4 dtype（目前通过 quant_min/quant_max 模拟）

注意：我们还将在量化张量 (Quantized Tensors) 部分讨论如何将稀疏性与量化结合

基本数据类型 (Basic DTypes)

dtype 有点过载，这里说的基本数据类型 (basic dtype) 是指那些没有额外元数据也能理解的数据类型（例如，当人们调用 torch.empty(.., dtype) 时能理解），更多详情请查阅：dev-discuss.pytorch.org/t/supporting-new-dtypes-in-pytorch/1833

无论我们进行何种量化，最终都将使用一些低精度数据类型 (low precision dtypes) 来表示量化数据，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 已添加到 torch，如果 float4 和 float6 流行起来，我们也计划将它们添加到 torch）

注意，上述部分类型目前仅为原型。当它们变得流行并获得硬件支持时，我们将考虑将其添加到 PyTorch 核心中。

当前支持 (Current Support)

在实际实现方面，有两个部分：1). 在 PyTorch 中，我们需要将 dtype 添加到 torch.dtype，例如 torch.uint2，示例：pytorch/pytorch#117208，但这只是占位符，以便我们可以使用 torch.uint2。2). 在 PyTorch 之外（例如在 torchao 中），我们使用张量子类实现这些 dtype 的张量操作，同时还需要一种标准的打包格式 (packing format)。

在 PyTorch 中添加占位符 dtype

如 dev-discuss.pytorch.org/t/supporting-new-dtypes-in-pytorch/1833 中所述，在 PyTorch 中添加 dtype 的标准是它已被广泛采用。对于上面提到的基本数据类型，PyTorch 中支持的包括

• torch.uint1 到 torch.uint8, torch.int1 到 torch.int8, torch.float8_e4m3fn, torch.float8_e5m2, torch.float8_e4m3fnuz, torch.float8_e5m2fnuz

对于其他类型，我们计划等到有更多证据表明其被广泛采用并获得硬件支持时再进行添加。

使用张量子类实现这些 dtype 的张量操作

为此，要求是我们确定一种“标准”打包格式 (packing format)，并且希望它易于高效实现，但对于 uintx 和 floatx，我们尚未集成足够的内核来决定这一点。因此，当前的 打包实现 尚非最终版本。在将更多 uintx、intx 和 floatx 内核集成到 torchao 后，我们可以重新讨论这一点。

将张量子类集成到 PyTorch 原生工厂函数中

之后，我们可以将工厂函数与张量子类连接起来，例如： torch.empty(..., dtype=torch.int4, ...) 可以创建一个 Int4Tensor 张量子类，其打包格式在之前的步骤中确定。

量化原语操作 (Quantization Primitive Ops)

量化原语操作 (Quantization primitive ops) 是指用于在低精度量化张量和高精度张量之间进行转换的操作符。我们将主要有以下量化原语操作符：choose_qparams ops：根据原始张量选择量化参数的操作符，通常用于动态量化，例如 affine 量化的 scale 和 zero_point；quantize op：根据量化参数将原始高精度张量量化为前一节提到的 dtype 的低精度张量；dequantize op：根据量化参数将低精度张量反量化回高精度张量。

上述操作可能存在变体以适应特定用例，例如对于静态量化，我们可能拥有 choose_qparams_affine_with_min_max，它将根据观察过程得出的 min/max 值来选择量化参数。

高效内核 (Efficient kernels)

我们还将拥有适用于低精度张量的高效内核，例如

_weight_int4pack_mm tinygemm int4 核 (bf16 激活 + int4 权重) int_matmul 接受两个 int8 张量并输出一个 int32 张量 int_scaled_matmul 执行 matmul 并对结果应用 scale。

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

量化张量（派生数据类型）(Quantized Tensors)

在基本数据类型、量化原语操作符和高效内核的基础上，我们可以将所有这些组合起来，通过继承 torch.Tensor 构建一个量化（低精度）张量 (Quantized Tensor)，该张量可以由高精度张量和一些用于配置用户所需特定量化的参数构造。我们也可以将此类称为派生数据类型 (derived dtypes)，因为它可以用基本数据类型的张量和一些额外的元数据（如 scale）表示。

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

布局和 TensorImpl (Layout and TensorImpl)

原生张量有一份硬编码的 布局 选择列表，最常见的是 strided 布局 (strided layout)，它提供了存储的分步式、多维视图；我们还有一些 sparse 和 mkldnn 布局。

以 sparse COO 张量 为例，它拥有 torch.sparse_coo 布局 (layout)，以及 SparseTensorImpl，后者改变了张量的存储方式。

将张量打包成不同格式的想法与布局 (layout) 概念非常吻合，这就是我们希望将其用于打包的原因。我们可以使用 Layout 来表示不同类型的打包格式，并使用 TensorImpl 来实现不同的存储格式。在 Python 级别的张量子类中，无需修改 C++ PyTorch 核心代码即可添加以打包格式存储张量的新 TensorImpl。

例如，对于 _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

注意，布局 (layout) 不仅是用于自定义数据表示的抽象，它也用于描述 TensorImpl 如何与不同的操作符交互，例如 transpose、quantized_linear，但操作符的语义应保持不变。

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

量化算法/流程 (Quantization Algorithms/Flows)

在栈的最顶层是最终的量化算法和量化流程。传统上我们有仅权重量化 (weight only quantization)、动态量化 (dynamic quantization) 和静态量化 (static quantization)，但现在我们也看到更多类型的量化正在出现。

为了演示目的，假设在之前的步骤之后我们定义了 AffineQuantizedTensor 和 to_affine_quantized 工厂函数。为简单起见，假设 to_affine_quantized 接受一个高精度浮点张量和一个 target_dtype（例如 torch.int8），并将其转换为具有相应 dtype 的 AffineQuantizedTensor。

注意：以下内容均为概念性解释，有关我们提供的工具和示例的更详细介绍，请参见 Tensor Subclass Developer Guide 部分。

仅权重量化 (Weight Only Quantization)

这是最简单的量化形式，很容易将仅权重量化应用于模型，特别是当我们拥有量化张量 (Quantized Tensor) 时。我们只需执行以下操作：

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

将上述操作应用于模型中的所有线性模块，即可获得仅权重量化模型。

动态激活和权重量化 (Dynamic Activation and Weight Quantization)

这之前被称为“动态量化”(dynamic quantization)，但它意味着我们在运行时动态量化激活，同时也量化权重。与仅权重量化相比，主要问题是如何将量化应用于激活。在 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 操作时，它将把存储的 input_qunat_func 应用于激活，并重新分派给使用量化激活和权重的 F.linear。

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

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

静态激活和权重量化 (Static Activation Quantization and Weight Quantization)

静态量化 (Static quantization) 意味着激活是静态量化的，而不是在运行时动态量化。在流程方面，静态量化需要使用样本数据进行校准 (calibration)，以便我们确定合适的量化参数。

从高层次上看，静态量化有三个步骤：(1) 插入观察者 (insert observers) (2) 校准 (calibration) (3) 量化模型 (quantize the model)

插入观察者 (Insert Observers)

在插入观察者这一步，我们需要向操作符的输入（和输出）激活以及权重添加观察者模块 (observer modules)，以收集张量的统计信息。因此，我们需要解决两个问题：如何定义观察者模块？如何将观察者模块添加到模型中？

如何定义观察者模块

观察者是特定的，取决于：(1) 量化类型（例如仿射量化 (affine quantization)，基于查找表的量化）(2) 我们希望跟踪的统计信息类型，例如 min max 观察者 (min max observer)，移动平均观察者 (moving average observer)。

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

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

如何将观察者模块添加到模型中

1. 使用张量子类 如果您只对量化线性操作符感兴趣，可以使用 linear activation weight observer，我们还有一个相应的 insert_observer_ API，用于处理修改线性层的权重。

2. 模块替换 (Module Swap) 另外，您也可以定义并使用 ObservedLinear 模块（或其他模块类型），并将未观察的模块替换为已观察的模块

校准 (Calibration)

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

量化 (Quantize)

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

另外，用户也可以进行 模块替换。

其他量化流程 (Other Quantization Flows)

对于不属于上述任何一种的其他量化流程/算法，我们也打算提供常见模式的示例。例如，类似 GPTQ 的量化流程 由 Autoround 采用，它使用 MultiTensor 和模块 hook 来优化模块。

如果您正在研究新的量化算法/流程，并且不确定如何在 PyTorch 原生方式下实现，请随时提出 issue，描述您的算法的工作方式，我们可以帮助提供实现细节方面的建议。

训练 (Training)

上述流程主要侧重于推理 (inference)，但低位宽数据类型张量 (low bit dtype Tensors) 也可以用于训练 (training)。

量化感知训练 (Quantization Aware Training)

待办 (TODO)

低位宽优化器 (Low Bit Optimizers)

目前我们有一些低位宽优化器 (low bit optimizers) 的原型：main/torchao/prototype/low_bit_optim，实现了特定类型的 4 位、8 位和 float8 优化器，并且可以与 FSDP 组合使用（带有查找表量化）。

量化训练 (Quantized Training)

与低位宽优化器类似，我们在 main/torchao/prototype/quantized_training 中提供了量化训练 (quantized training) 原型，我们可以扩展 AffineQuantizedTensor 以支持训练。初步支持正在进行中，但还需要大量后续工作，包括使其适用于不同的内核等。

您还可以查看 量化训练 (Quantized Training) 的教程，该教程介绍了如何使 dtype 张量子类可训练。

案例研究：torchao 中的 int4 仅权重量化是如何工作的？

为了将所有内容串联起来，这里更详细地讲解一下 int4 仅权重量化在 torchao 中的实现方式。

量化流程： quantize_(model, int4_weight_only())

• 发生的情况是： linear.weight = torch.nn.Parameter(to_affine_quantized_intx(linear.weight, …), requires_grad=False))

• 量化原语操作：调用 choose_qparams 和 quantize_affine 来量化张量

• 量化后的张量将是 AffineQuantizedTensor，这是一个带有派生数据类型（例如带有 scale 和 zero_point 的 int4）的量化张量

• 打包操作符 _convert_weight_to_int4pack，用于将量化权重打包以实现高效执行

模型执行期间： model(input)

• 对输入和打包的权重调用 torch.ops.aten._weight_int4pack_mm

量化期间 (During Quantization)

首先我们从 API 调用开始： quantize_(model, int4_weight_only())，它的作用是将模型中 nn.Linear 模块的权重转换为 int4 量化张量（一个 AffineQuantizedTensor，它是 int4 dtype，非对称，按组量化），使用 tinygemm 核的布局： tensor_core_tiled 布局。

• quantize_：模型级别的 API，通过应用用户提供的转换函数（第二个参数）来量化线性层的权重

• int4_weight_only：返回一个函数的函数，该函数将线性层的权重转换为 int4 仅权重量化权重 * 调用量化原语操作符（如 choose_qparams_affine 和 quantize_affine）来量化张量

• TensorCoreTiledLayout：Tensor Core 平铺布局类型，存储打包格式的参数

• TensorCoreTiledAQTTensorImpl：Tensor Core 平铺的 TensorImpl，存储打包的权重以用于高效的 int4 仅权重核 (tinygemm 核)

模型执行期间 (During Model Execution)

当我们运行量化模型 model(inputs) 时，我们将通过 nn.Linear 中的函数式线性操作符。

return F.linear(input, weight, bias)

其中输入是一个 bfloat16 张量，权重是一个 int4 AffineQuantizedTensor，它会调用 AffineQuantizedTensor 子类的 __torch_function__，当输入之一是 AffineQuantizedTensor 时，这将最终进入 F.linear 的实现，因此它会调用：

return weight_tensor._quantized_linear_op(input_tensor, weight_tensor, bias)

_quantized_linear_op 会遍历 _AQT_QLINEAR_DISPATCH_TABLE 并检查每个调度条件，如果调度条件满足，它将调用使用 input/weight/bias 的实现。请查看此文档以了解 dispatch_condition 和 impl 的解释。

int4 仅权重的调度条件检查输入是否为 bfloat16 张量且权重是否为 uint4 AffineQuantizedTensor。int4 仅权重量化的内核实现接受一个 bfloat16 输入张量和一个 int4 仿射量化张量，并调用 torch.ops.aten._weight_int4pack_mm，传入输入张量以及存储在 weight_tensor.tensor_impl 中的打包权重。

保存/加载期间

由于 AffineQuantizedTensor 权重仍然是一个 torch.Tensor，因此保存/加载的工作方式与原始高精度浮点模型相同。有关更多详细信息，请参阅序列化文档。