量化概述¶

首先，我们想介绍一下 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）仅权重（weight-only）量化流程（2）tinygemm bf16 激活 + int4 权重内核和量化原始操作（3）AffineQuantizedTensor张量子类，带有TensorCoreTiledLayout（4）torch.uint4 dtype（目前通过 quant_min/quant_max 模拟）

注意：我们还将在“量化张量”部分讨论如何将稀疏性与量化结合起来

基本数据类型¶

dtype 是一个被重载的术语，基本 dtype 指的是无需任何额外元数据即可理解的 dtype（例如，当人们调用 torch.empty(.., dtype) 时即可理解），欲了解更多详情，请查看：dev-discuss.pytorch.org/t/supporting-new-dtypes-in-pytorch/1833

无论我们正在进行何种量化，最终我们都将使用一些低精度 dtype 来表示量化数据，我们旨在 torchao 中支持的 dtype 是

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 核心中。

当前支持¶

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

在 PyTorch 中添加占位符 dtype¶

正如 dev-discuss.pytorch.org/t/supporting-new-dtypes-in-pytorch/1833 中所述，在 PyTorch 中添加 dtype 的标准是它显示出广泛的采用。对于上述基本 dtype，PyTorch 中支持的 dtype 包括

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

对于其他类型，我们计划等到有更多证据表明其广泛采用和硬件支持。

使用张量子类实现这些数据类型的张量操作¶

为此，要求是我们要确定一种“标准”的打包格式，并且希望这种格式易于高效实现，但对于 uintx 和 floatx，我们还没有集成足够的内核来决定这一点。因此，当前的打包实现尚未最终确定。在将更多 uintx、intx 和 floatx 内核集成到 torchao 中之后，我们可以重新审视。

将张量子类集成到 PyTorch 本机工厂函数中¶

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

量化原始操作¶

量化原始操作是指用于在低精度量化张量和高精度张量之间进行转换的操作符。我们主要有以下量化原始操作符：选择量化参数操作：根据原始张量选择量化参数，通常用于动态量化，例如仿射量化的比例和零点。量化操作：根据量化参数将原始高精度张量量化为前一节中提到的数据类型的低精度张量。反量化操作：根据量化参数将低精度张量反量化为高精度张量。

上述操作可能会有变体以适应特定的用例，例如对于静态量化，我们可能有 choose_qparams_affine_with_min_max，它将根据观察过程中获得的最小/最大值选择量化参数。

高效内核¶

我们还将拥有与低精度张量配合使用的高效内核，例如：

_weight_int4pack_mm tinygemm int4 核（bf16 激活 + int4 权重） int_matmul，接受两个 int8 张量并输出一个 int32 张量 int_scaled_matmul，执行矩阵乘法并对结果应用缩放。

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

量化张量（派生数据类型）¶

在基本数据类型、量化原始操作符和高效内核的基础上，我们可以将所有这些结合起来，通过对 torch.Tensor 进行子类化来构建一个量化（低精度）张量，该张量可以从高精度张量和一些可以配置用户所需特定量化的参数构建。我们也可以称之为派生数据类型，因为它可以用基本数据类型的张量和一些额外的元数据（如比例）来表示。

torchao 中现有的例子是 AffineQuantizedTensor，这意味着低精度张量通过仿射映射从高精度张量量化而来，即：low_precision_val = high_precision_val / scale + zero_point，其中 scale/zero_point 是可以通过量化原始操作或通过一些优化过程计算的量化参数。仿射量化是一种非常常见的量化类型，因为它直接将高精度值映射到低精度值，执行仿射变换（high_preicsion_val / scale + zero_point）。另一种常见的量化类型，尤其是对于更低的位宽（例如低于 4 位），是码本/查找表（look up table）量化。

布局和 TensorImpl¶

原生张量有一个硬编码的布局选择列表，最常见的是跨步布局（strided layout），它提供了存储的跨步多维视图，我们还有一些稀疏和 mkldnn 布局。

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

将张量打包成不同格式的想法与布局概念非常契合，这就是我们希望将其用于打包的原因。我们可以使用Layout表示不同类型的打包格式，并使用TensorImpl表示不同存储格式的实现。并且新的 TensorImpl，用于以打包格式存储张量，可以在 Python 级别的张量子类中添加，而无需修改 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 如何与不同操作符交互，例如，相同的数据表示在运行相同操作符时（例如，转置、量化线性）可以有不同的实现，但操作符语义应该保持不变。

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

量化算法/流程¶

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

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

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

仅权重（Weight Only）量化¶

这是最简单的量化形式，并且很容易将仅权重（weight only）量化应用于模型，特别是当我们拥有量化张量时。我们所需要做的就是：: linear_module.weight = torch.nn.Parameter(to_affine_quantized_intx(linear_module.weight, …), requires_grad=False))

将上述应用于模型中的所有线性模块，我们将得到一个仅权重（weight only）量化模型。

动态激活和权重（Dynamic Activation and Weight）量化¶

这以前被称为“动态量化”，但它意味着我们在运行时动态量化激活，并同时量化权重。与仅权重（weight only）量化相比，主要问题是如何将量化应用于激活。在 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。

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

但最好使用张量子类，因为它更容易序列化/反序列化。如果使用张量子类来支持动态量化，那么我们可以直接加载量化权重，而无需对模型进行进一步的准备。否则，我们需要先进行模块交换或对模型进行其他修改，然后才能加载量化权重。

静态激活量化和权重（Static Activation Quantization and Weight）量化¶

静态量化意味着激活是静态量化的，而不是在运行时动态量化的。就流程而言，静态量化需要使用样本数据进行校准，以便我们能够确定合适的量化参数。

在高级别上，静态量化有三个步骤：（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），参见示例。

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

其他量化流程¶

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

如果您正在开发新的量化算法/流程，但不确定如何以 PyTorch 原生方式实现它，请随时提出问题，描述您的算法如何工作，我们可以提供实施细节方面的建议。

训练¶

上述流程主要侧重于推理，但低位宽数据类型张量也可用于训练。

量化感知训练¶

待办事项

低位优化器¶

今天，我们有一些原型低位优化器：main/torchao/prototype/low_bit_optim，它实现了一种特定类型的 4 位、8 位和 float8，并且还可以与 FSDP 组合使用（带有查找表量化）。

量化训练¶

类似于低位优化器，我们在 main/torchao/prototype/quantized_training 中有量化训练原型，并且我们可以扩展 AffineQuantizedTensor 以支持训练，初步启用正在进行中，但仍需要大量后续工作，包括使其适用于不同的内核等。

您还可以查看量化训练的教程，其中介绍了如何使 dtype 张量子类可训练。

案例研究：int4 仅权重（weight only）量化在 torchao 中如何工作？¶

为了将所有内容联系起来，这里是对 int4 仅权重（weight only）量化在 torchao 中如何实现的更详细的说明。

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

发生什么：linear.weight = torch.nn.Parameter(to_affine_quantized_intx(linear.weight), requires_grad=False)
量化原始操作：调用 choose_qparams 和 quantize_affine 来量化张量
量化张量将是 AffineQuantizedTensor，一个具有派生数据类型（例如，带比例和零点的 int4）的量化张量
打包操作 _convert_weight_to_int4pack 用于打包量化权重以实现高效执行

模型执行期间：model(input)

在输入和打包权重上调用 torch.ops.aten._weight_int4pack_mm

量化期间¶

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

quantize_：模型级别 API，通过应用用户提供的转换函数（第二个参数）量化线性模块的权重
Int4WeightOnlyConfig：返回一个将线性权重转换为 int4 仅权重（weight only）量化权重的函数 * 调用 choose_qparams_affine 和 quantize_affine 等量化原始操作来量化张量
TensorCoreTiledLayout：张量核心平铺布局类型，存储打包格式的参数
TensorCoreTiledAQTTensorImpl：张量核心平铺 TensorImpl，存储打包权重以实现高效的 int4 仅权重（weight only）内核（tinygemm 内核）

模型执行期间¶

当我们运行量化模型 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 仅权重（weight only）分派条件检查输入是否为 bfloat16 张量以及权重是否为 uint4 AffineQuantizedTensor，int4 仅权重（weight only）量化内核实现接受 bfloat16 输入张量和 int4 AffineQuantizedTensor，并调用 torch.ops.aten._weight_int4pack_mm，输入张量和存储在 weight_tensor.tensor_impl 中的打包权重。

保存/加载期间¶

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