贡献者指南

扩展 torchao 的通用指南

对于新的用例，例如训练数据类型（如 fp4 训练），可以从在原型文件夹 torchao/prototype 中添加新的张量子类开始，但如果你想做的大部分事情都已在那里支持，例如为完全相同的仿射量化添加 int3 内核，也可以查看 AffineQuantizedTensor。欢迎随时提出问题，如果你对特定新用例如何操作有疑问。更多详细信息，请参阅我们的量化概述页面。

贡献现有代码库

• 为 AffineQuantizedTensor 添加功能，例如使其可训练，添加张量并行支持等：torchao/dtypes/affine_quantized_tensor.py

• 添加新的量化 API：torchao/quantization/quant_api.py

• 添加新的量化原始操作，例如现有量化原始操作的微小变体：torchao/quantization/quant_primitives.py

• 添加新的自动调整 Triton 内核：torchao/kernel

• 添加新的自定义 cpu/cuda/mps 内核：torchao/csrc

• 将自定义内核与 AffineQuantizedTensor 集成（也可能是一种新的布局）：以添加稀疏 Marlin AQT 布局 #621 为例。我们尚未决定是否将 AffineQuantizedTensor 拆分为更多张量子类。

添加高效内核

自定义 Triton 内核

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

• 实现示例

• 注册为自定义操作

你可能还需要定义自己的 自动调整器。

自定义手写内核

CPU/CUDA/MPS 的自定义内核（实现）可以通过 torchao/csrc 实现，例如 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)

序列化

更多详情请查看序列化文档。

注意

我们已与 Hugging Face Transformer 集成，并通过 Hugging Face 的 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 文件夹中放入更多示例，涵盖以下用例。

• 量化训练

• 量化张量的张量并行支持

• 与 executorch / torchchat 的兼容性

• [待办] FSDP

• [待办] QAT

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

我们还在开发测试套件，用于测试张量子类的功能以及与 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])

all linear shapes (M, K, N): [(10, 10, 10)]

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

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

模型基准测试和评估

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

注意：llama 模型（llama2/llama3）是我们内存密集型模型的代表模型，SAM 是我们计算密集型模型的代表模型。

• llama * 基准测试 * 评估

• SAM * 基准测试和评估

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

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