贡献者指南¶

扩展 torchao 的通用指南¶

为现有代码库做贡献

添加新的 Tensor：torchao/quantization/quantize_/workflows
添加新的量化 API：torchao/quantization/quant_api.py
为现有的 Tensor 子类（如 Float8Tensor）添加功能，例如添加新的算子支持、使其可训练、添加张量并行支持等，请参考张量子类、测试
添加新的量化原语操作，例如现有量化原语操作的微小变体：torchao/quantization/quant_primitives.py
添加新的自动调优 Triton 内核：torchao/kernel
添加新的自定义 CPU/CUDA/MPS 内核：torchao/csrc

添加新的张量子类¶

torchao 张量子类由派生 dtype和打包格式进行组织，请查阅量化概述页面以理解这些概念。如果您的用例需要新的张量子类，即新的 dtype 或新的打包格式，而这些尚不存在，我们可以定义一个新的张量。

要理解如何在量化的上下文中实现张量子类，请参阅编写您自己的量化张量。

我们有一个实用基类：torchao.utils.TorchAOBaseTensor，它可以通过指定张量子类中张量和非张量属性的名称，帮助定义常见的实用函数和方法。例如：

class MyTensor(TorchAOBaseTensor):
    tensor_data_names = ["qdata", "scale"]
    tensor_attribute_names = ["device", "dtype"]

有了上述内容，我们将拥有多个方法和函数可供此张量使用，更多详情请参阅TorchAOBaseTensor的文档。

注意

torchao 中的许多现有用例仍然使用 AffineQuantizedTensor，但我们计划逐步淘汰它，以减少抽象层级，让更多人能够轻松地为 torchao 做贡献。

添加高效内核¶

自定义 Triton 内核¶

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

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

自定义手动编写的内核¶

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

在张量子类中使用手动编写的内核¶

为了调用优化后的内核，张量子类中有implements，例如，如果我们想调用一个新的自定义操作：torch.ops.torchao.my_mm_for_mps。

class Float8Tensor(TorchAOBaseTensor):
    ...

implements = Float8Tensor.implements

@implements([torch.nn.functional.linear, aten.linear.default])
def _(func, types, args, kwargs):
    ...
    # call into the custom op
    res = torch.ops.torchao.my_mm_for_mps(input_tensor.qdata, weight_tensor.qdata, input_tensor.scale, weight_tensor.scale)
    return res

KernelPreference¶

对于某些张量子类，量化和 mm 等操作可能有多个内核选择。在 torchao 张量子类中处理此问题的推荐方法是使用KernelPreference，它代表了我们希望用于量化、mm、group_mm 等操作的内核组。我们可以使用KernelPreference.AUTO作为默认选项，让开发者选择我们认为在不同条件下对用户最快的选项，这样用户就不需要担心细节，并且我们可以提供其他更具体的内核选项用于调试目的。

Float8Tensor 例如，具有

KernelPreference.AUTO，它将根据硬件（H100 SM89 或 SM90+）、库的可用性（是否安装了fbgemm_gpu_genai）、粒度（每行或每张量）选择最高效的量化和 mm 内核。
KernelPreference.TORCH 将使用 torchao 量化操作（_choose_scale_float8 和_quantize_affine_float8）和_scaled_mm。
Kerenel.FBGEMM 使用 fbgemm 量化和 mm 操作（torch.ops.fbgemm.f8f8bf16_rowwise）。

流程¶

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

请参阅量化算法/流程部分，了解仅权重/动态量化和其他模型级别 API 的示例。

使用 torch.compile 提高性能¶

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

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

序列化¶

为了支持序列化（使用张量子类作为权重的 torch.save 和 torch.load），我们需要将张量子类和相关对象添加到 safe globals（在 torch 2.5 之后可用），例如：: torch.serialization.add_safe_globals([Float8Tensor, QuantizeTensorToFloat8Kwargs])

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

注意

我们已集成了 Huggingface Transformer，并通过 Huggingface 的save_pretrained、push_to_hub和from_pretrained API 支持序列化和反序列化。我们还有Diffuser 模型的序列化示例。

其他功能支持¶

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

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

我们还在开发测试套件，以测试张量子类功能以及与 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，目前可以任选其一使用，但将来我们可能只会保留一个。

模型基准测试和评估¶

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

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

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

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

另请参阅基准测试用户指南和基准测试 API 指南，以了解如何使用我们的基准测试框架。