目录
  • 文档 >
  • torchao 贡献者指南
快捷方式

torchao 贡献者指南

目标

在本文档中，我们将讨论 (1). torchao 中不同优化技术的结构 (2). 如何为 torchao 做出贡献

注意：本文档目前主要关注推理，但我们计划在未来扩展到涵盖训练技术。

torchao 堆栈概述

首先，我们想要介绍 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 模拟）

注意：我们还将在量化张量部分讨论如何将稀疏性与量化结合使用

基本数据类型 (DTypes)

dtype 是一个有点过载的术语，基本数据类型 (dtype) 指的是在没有任何额外元数据的情况下有意义的数据类型（例如，当人们调用 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 已添加到 torch，如果 float4 和 float6 变得流行，我们也计划将其添加到 torch)

注意：以上某些类型目前仅为原型。当它们变得流行并获得硬件支持时，我们将考虑将其添加到 pytorch 核心。

当前支持

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

在 PyTorch 中添加占位符数据类型

正如 dev-discuss.pytorch.org/t/supporting-new-dtypes-in-pytorch/1833 中提到的，在 PyTorch 中添加数据类型的标准是它显示出广泛的应用。对于上述提到的基本数据类型，PyTorch 中支持的有

  • 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 操作：根据原始张量选择量化参数，通常用于动态量化，例如仿射量化的 scale 和 zero_point；quantize 操作：根据量化参数，将原始高精度张量量化为前一节中提到的数据类型的低精度张量；dequantize 操作：根据量化参数，将低精度张量反量化为高精度张量。

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

高效内核

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

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

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

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

在基本数据类型、量化原语运算符和高效内核之上，我们可以将所有内容粘合在一起，并通过子类化 torch.Tensor 构建量化（低精度）张量，该张量可以从高精度张量和一些可以配置用户想要的特定量化的参数构建，我们也可以称之为派生数据类型，因为它可以用基本数据类型的张量和一些额外的元数据（如 scale）表示。

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 如何与不同的运算符交互，例如，相同的数据表示在运行相同的运算符（例如 transpose、quantized_linear）时可能具有不同的实现，但运算符语义应保持不变。

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

量化算法/流程

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

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

注意：以下所有内容都用于解释概念，有关我们提供的实用程序和示例的更详细介绍，请参见 张量子类开发者指南 部分。

仅权重（Weight Only）量化

这是最简单的量化形式，并且很容易将仅权重量化应用于模型，特别是自从我们有了量化张量之后。我们所需要做的就是：

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 操作时，它将应用存储的 input_qunat_func 到激活，并重新分派到具有量化激活和权重的 F.linear。

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

但是，使用张量子类是首选方法，因为它更易于序列化/反序列化，如果我们使用张量子类来支持动态量化，那么我们可以直接加载量化权重，而无需为模型做进一步准备。否则，我们需要在加载量化权重之前先对模型进行模块交换或其他修改。

静态激活量化和权重量化

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

在高层面上，静态量化有三个步骤：(1) 插入观察者 (2) 校准 (3) 量化模型

插入观察者

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

如何定义观察者模块

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

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

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

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

  1. 使用张量子类 如果您唯一感兴趣量化的运算符是线性运算符，则可以使用 线性激活权重观察者，我们还有一个相应的 insert_observer_ API，用于处理修改线性的权重。

  2. 模块交换 或者，您也可以定义 ObservedLinear 模块（或其他模块类型），并将未观察的模块与观察的模块交换

校准

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

量化

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

或者，用户也可以进行 模块交换。

其他量化流程

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

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

训练

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

量化感知训练

待办事项

低位优化器

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

量化训练

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

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

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

为了将所有内容联系起来，以下是关于 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 在输入和打包权重上被调用

量化期间

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

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

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

  • TensorCoreTiledLayout：tensor core tiled 布局类型，存储打包格式的参数

  • TensorCoreTiledAQTTensorImpl: 张量核心平铺的 TensorImpl，存储用于高效 int4 权重专用内核（tinygemm 内核）的打包权重

模型执行期间

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

return F.linear(input, weight, bias)

其中 input 是一个 bfloat16 张量，weight 是一个 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 权重专用 dispatch_condition 检查输入是否为 bfloat16 张量，以及权重是否为 uint4 AffineQuantizedTensor。wint4 权重专用量化 内核实现 接受 bfloat16 输入张量和一个 int4 AffineQuantizedTensor，并使用输入张量和存储在 weight_tensor.tensor_impl 中的打包权重调用 torch.ops.aten._weight_int4pack_mm。

保存/加载期间

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

张量子类开发者指南

我们在上一节中介绍了高级概述以及所有内容如何连接在一起，本节将重点介绍张量子类，这是我们依赖的主要扩展点，以提供灵活性，从而在这些场景中支持使用低精度张量进行推理、训练和微调，以及与 torch.compile、autograd、分布式原语的可组合性。

先决条件

一些外部可用的张量子类资源

为什么选择张量子类？

人们可以通过多种方式实现量化技术或新的数据类型，我们推荐基于张量子类方法的主要动机有三点：（1）。将量化建模为数据类型转换是很自然的，因此使用张量子类实现它意味着我们没有引入新概念，而是重用 PyTorch 核心中已有的概念，如数据类型、布局（2）。由于张量子类在 torch function 或 aten ops 级别拦截计算，只要使用相同的函数/运算符，我们就能够量化模型。这使得使用原生模块变体（例如，nn.Linear 的略微修改版本）的模型仍然与量化兼容（3）。张量子类也是其他技术（如稀疏性和分布式）采用的方法，因此使用张量子类实现量化或数据类型转换将使其更容易与这些技术组合。

新数据类型的示例代码

请随时从教程开始，了解结合了我们讨论的所有内容的端到端工作示例，并返回文档以获取澄清和文档。

基本结构

张量子类需要定义一些基本方法：__new__、__init__、__tensor_flatten__、__tensor_unflatten__ 以及 torch 函数 __torch_function__ 和 aten ops __torch_dispatch__ 的调度函数。

这是一个基本结构的示例：

# 查看 https://github.com/pytorch/ao/blob/e283743b3cc4612bb641b88dca3670231724d396/torchao/utils.py#L437 中的文档 from torchao.utils import TorchAOBaseTensor

class MyDTypeLayout(TorchAOBaseTensor)

# 有关详细信息，请参阅教程代码 pass

class MyDtypeTensor(TorchAOBaseTensor)

“””我们需要定义 __new__ 以构造新的张量子类实例，并定义 __init__ 以初始化实例。这里对参数列表的外观没有要求，唯一的要求是 __new__ 必须通过调用 torch.Tensor._make_wrapper_subclass(cls, shape, ...) 返回一个 Tensor 实例 “”” @staticmethod def __new__(

cls, tensor_impl: MyDTypeLayout, shape: torch.Size, dtype: Optional[torch.dtype] = None,

):

… return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined]

def __init__(

self, tensor_impl: MyDTypeLayout, shape: torch.Size, …

):

self.tensor_impl = tensor_impl

“””__tensor_flatten__ 和 __tensor_unflatten__ 用于将张量解构为原生张量/属性，并从解构的张量和属性重建张量子类实例，这些是为 torch.compile 支持定义张量子类所必需的 “”” def __tensor_flatten__(self)

return [“tensor_impl”], [self.shape]

“””有关 outer_size 和 outer_stride 的文档，请参阅 https://github.com/pytorch/pytorch/blob/3bc2004f9123a32f381ef64202252d59109507f3/torch/utils/_python_dispatch.py#L289 “”” @classmethod def __tensor_unflatten__(

cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride

):
tensor_impl = tensor_data_dict[“tensor_impl”]

shape, = tensor_attributes return cls(

tensor_impl, shape if outer_size is None else outer_size,

)

“””从浮点张量 (fp32/fp16/bf16) 转换为当前数据类型的类方法 “””

@classmethod
def from_float(

cls, input_float: torch.Tensor,

):

mapping_type = MappingType.SYMMETRIC block_size = input_float.shape dtype = torch.int16 scale, _ = choose_qparams_affine(input_float, mapping_type, block_size, dtype) int_data = (input_float / scale).to(torch.int8) tensor_impl = MyDTypeLayout.from_plain(int_data, scale) return cls(tensor_impl, input_float.shape)

“””[可选] 请参阅 量化张量 部分下的 Layout/Packing 文档，以了解 layout_type 是什么 “”” @property def _layout(self) -> LayoutType

return self.tensor_impl._layout

“””有两种入口点可以修改 PyTorch 操作的行为：torch_function 和 torch_dispatch

__torch_function__: 每当在 Tensor 对象上调用 torch 级别函数时，都会调用它，例如：torch.nn.functional.linear、tensor.detach、tensor.reshape、tensor.t 等。

__torch_dispatch__: 当在 Tensor 对象上调用 aten 运算符时，将在 C++ 调度器中调用它，例如：aten.mm、aten.addmm、aten.detach.default、aten.t.default 等。您可以查看 https://github.com/pytorch/ao/blob/e283743b3cc4612bb641b88dca3670231724d396/torchao/utils.py#L361-L389 以了解 __torch_function__ 和 __torch_dispatch__ 正在做什么，但使用 TorchAoBaseTensor 用户可以直接使用一些辅助函数（请参阅下一节）

运算符支持

运算符支持有两种类型：torch function 和 aten ops。对于 torch function（例如，torch.nn.functional.linear），我们需要覆盖 Tensor 子类中的 __torch_function__ 回调，对于 aten ops（例如，torch.ops.aten.mm），我们需要覆盖 __torch_dispatch__ 回调函数。

对于新的数据类型，我们希望人们定义以下装饰器：

如果您的数据类型类继承自 torchao.utils.TorchAoBaseTensor，您可以执行

implements = my_dtype_tensor_cls.implements

我们可以使用以下方法实现运算符调度：

# torch.nn.functional.linear 的 torch_function 调度的示例 def _quantized_linear_op(input_tensor, weight_tensor, bias)

if isinstance(input_tensor, MyDtypeTensor)

input_tensor = input_tensor.dequantize()

if isinstance(weight_tensor, MyDtypeTensor)

weight_tensor = weight_tensor.dequantize()

return torch.nn.functional.linear(input_tensor, weight_tensor, bias)

@implements(torch.nn.functional.linear) def _(*args, **kwargs)

input_tensor, weight_tensor, bias = (

args[0], args[1], args[2] if len(args) > 2 else None,

) # 在这里使用 try/except，以便当 input_tensor/weight_tensor # 未被 _quantized_linear_op 中的任何调度路径选中时，我们可以有一个通用的回退，这使我们能够 # 使 _quantized_linear_op 中的分支更容易理解 try

return _quantized_linear_op(input_tensor, weight_tensor, bias)

except NotImplementedError
if isinstance(input_tensor, MyDtypeTensor)

input_tensor = input_tensor.dequantize()

if isinstance(weight_tensor, MyDtypeTensor)

weight_tensor = weight_tensor.dequantize()

return torch.nn.functional.linear(input_tensor, weight_tensor, bias)

# aten.detach.default 的 aten op 调度的示例 @implements(aten.detach.default) def _(func, *args, **kwargs)

# return_and_correct_aliasing 应由希望与 torch.compile 一起使用的包装器张量 __torch_dispatch__ 子类使用。它确保子类正确实现每个操作的别名行为， # 这对于 AOTAutograd 中的正确性是必需的。

# _apply_fn_to_data 仅将函数应用于 args[0] 中的张量数据，args[0] 是 my_dtype 的张量子类 # return return_and_correct_aliasing(

func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)

)

我们需要覆盖哪些操作？这取决于我们尝试量化的模型，通常覆盖的操作有：__torch_function__: torch.nn.functional.linear __torch_dispatch__: torch.ops.aten.addmm.default, torch.ops.aten.mm.default, torch.ops.aten.detach.default, torch.ops.aten.t.default

您还可以使用以下代码找到可以在 __torch_function__ 或 __torch_dispatch__ 中覆盖的操作，您可以从您想要优化的模型开始，首先仅覆盖重要的操作（如 linear），然后逐步扩展覆盖范围，直到测试运行并且您获得预期的优化生成代码（有关更多详细信息，请参阅“优化运算符”部分）：
class M(torch.nn.Module)
def __init__(self) -> None

super().__init__() self.linear = torch.nn.Linear(10, 10)

def forward(self, x: torch.Tensor) -> torch.Tensor

return self.linear(x) + x

from torch.overrides import TorchFunctionMode class TorchFunctionLoggingMode(TorchFunctionMode)

def __torch_function__(cls, func, types, args=(), kwargs=None)
if kwargs is None

kwargs = {}

print(f”TORCH_FUNC={str(func)}”) return func(*args, **kwargs)

with TorchFunctionLoggingMode()

m(*example_inputs)

## 示例输出 # TORCH_FUNC=<built-in function linear> # TORCH_FUNC=<method ‘add’ of ‘torch._C.TensorBase’ objects>

from torch.utils._python_dispatch import TorchDispatchMode class TorchDispatchLoggingMode(TorchDispatchMode)

def __torch_dispatch__(cls, func, types, args=(), kwargs=None)
if kwargs is None

kwargs = {}

print(f”ATEN_FUNC={str(func)}”) return func(*args, **kwargs)

with TorchDispatchLoggingMode()

m(*example_inputs)

## 示例输出 # ATEN_FUNC=aten.t.default # ATEN_FUNC=aten.addmm.default # ATEN_FUNC=aten.add.Tensor

# 或更完善的 torch_dispatch (aten) ops 日志记录：https://github.com/albanD/subclass_zoo/blob/main/logging_mode.py

或者，您可以运行测试示例（例如，将量化模型与张量并行、FSDP 等一起使用），并发现缺少的 ops 并添加它们，直到测试通过。

我们仍在制定一个表格，讨论对于每个功能，需要支持哪些运算符。

添加高效内核

自定义 triton 内核

自定义 triton 内核可以在 torchao/kernel 中实现和注册

您可能还需要定义自己的 autotuner。

自定义手写内核

可以通过 torchao/csrc 实现 cpu/cuda/mps 的自定义内核（实现），例如 int4 cuda，并通过 torch.ops.my_custom_op 访问

调度

为了调度到 cpu/cuda/mps 设备优化的内核，我们可以在 __torch_function__ 或 __torch_dispatch__ 中检查调度条件，并调度到目标运算符，例如，可以在此处找到 bfloat16 激活和 uint4 权重内核的条件。

特别是对于 AffineQuantizedTensor，我们还允许人们扩展量化线性以使用新的高效内核或通过定义两个函数来实现：dispatch_condition（定义调度到内核的条件）和 impl（获取激活、(量化) 权重、偏置张量并运行高效内核的实际实现），两者都以 input_tensor、weight_tensor、bias 作为参数，并且可以使用 register_aqt_quantized_linear_dispatch 注册到 AffineQuantizedTensor 中量化线性的调度。此处 是一个示例，展示了它的工作原理。

布局/TensorImpl

有时，量化权重必须打包才能产生最佳性能。这可以使用 layout 来抽象。有关完整示例，请参阅 此处。

流程

在实现张量子类之后，我们还可以将其包装到工厂函数中，例如：

# 从浮点张量转换为我的数据类型张量子类 to_my_dtype = MyDTypeTensor.from_float

对于模型级别 API，人们可以重用 torchao.quantize_，它允许人们将张量子类转换应用于线性的权重，并允许使用 过滤函数 来选择应将张量子类转换应用于哪个模块。

有关基于工厂函数的仅权重/动态量化/静态量化和其他类型的模型级别 API 的示例，请参阅“量化算法/流程”部分。

使用 torch.compile 提升性能

注意：对于 PyTorch 2.4 及更低版本，我们需要使用以下代码：

from torchao.utils import unwrap_tensor_subclass m_unwrapped = unwrap_tensor_subclass(m)

为了与 torch.compile 兼容，为了实现性能优化，我们应该首先使用 fullgraph=True 运行 torch.compile，并删除任何不必要的图中断。您可以在运行脚本时添加 TORCH_LOGS="output_code" 以查看 inductor 生成的代码。例如，TORCH_LOGS="output_code" python example.py：

model = torch.compile(model, mode=”max-autotune”, fullgraph=True)

序列化

有关更多详细信息，请查看序列化文档。

注意

我们已与 huggingface transformers 集成，并通过 huggingface save_pretrained/push_to_hub/from_pretrained API 支持序列化/反序列化：https://hugging-face.cn/docs/transformers/main/en/quantization/torchao

注意

在与 diffuser 集成中可以找到另一个示例：https://github.com/sayakpaul/diffusers-torchao/blob/main/inference/serialization_and_loading.md

其他功能支持

上面仅讨论了基本功能支持，我们还提供了关于如何通过扩展 MyDTypeTensor 来添加对训练、张量并行、FSDP 的支持的示例，我们将在 developer_api_guide 文件夹中放置更多示例，涵盖以下用例。

扩展 torchao 的通用指南

对于新的用例，例如，训练数据类型（如 fp4 训练），可以从在原型文件夹 torchao/prototype 中添加新的张量子类开始，但如果您的需求在很大程度上已在 AffineQuantizedTensor 中得到支持，例如为完全相同的仿射量化添加 int3 内核，您也可以查看它。如果您对特定新用例的操作有疑问，请随时打开 issue。

为现有代码库做贡献

张量子类功能/可组合性测试

我们还在开发测试套件，以测试张量子类的功能以及与 torch.compile、DTensor 等不同系统的可组合性（我们建议复制粘贴测试并进行调整，以测试您自己的张量子类）。

内核微基准测试

在我们测试模型性能之前，我们还可以对具有不同输入维度的单个线性运算符（或其他计算密集型/内存密集型）运算符进行一些微基准测试，以了解加速情况。对于您要进行基准测试的特定内核，您可以创建一个基准测试文件，如 benchmarks/benchmark_aq.py，并使用对目标模型重要的不同形状运行基准测试。快速获取线性运算和其他运算的相关形状的方法是使用 此 运行示例。

使用您感兴趣优化的模型更改模型，并运行以下命令：

python tutorials/developer_api_guide/print_op_and_shapes.py

示例输出：

TORCH_FUNC=<built-in function linear> (M, K, N): 10 10 10 TORCH_FUNC=<method ‘add’ of ‘torch._C.TensorBase’ objects> args[0] shape: torch.Size([10, 10])

所有线性形状 (M, K, N): [(10, 10, 10)]

所有线性形状的输出可以复制粘贴到 benchmarks/benchmark_your_kernel.py 下的微基准测试脚本代码中以进行基准测试。

对于基准测试辅助函数，现在我们有 1 和 2，现在可以随意使用其中任何一个，但我们将来可能会保留一个。

模型基准测试和评估

在您实现了量化流程之后，您可以在已经修改为对 torch.compile 友好的 llama（llama2/llama3）或 sam 模型上运行基准测试和评估，并与 torchao 中的现有技术进行比较。

注意：llama 模型（llama2/llama3）是我们内存受限模型的代表模型，而 sam 是我们计算受限模型的代表模型。

请查看每个脚本的 --help 选项以了解支持的选项，例如，您可以使用 --profile=profile_path 来获取运行的 chrome 跟踪，以了解详细的 chrome 跟踪。

如果有任何新的重要模型适合添加到 torchao 模型基准测试/评估文件夹中，请告知我们。

文档

访问 PyTorch 的综合开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得您的问题解答

查看资源

为了分析流量和优化您的体验，我们在本网站上提供 cookies。通过点击或导航，您同意允许我们使用 cookies。作为本网站的当前维护者，Facebook 的 Cookie 政策适用。了解更多信息，包括关于可用控件：Cookies 政策。