贡献者指南¶

扩展 torchao 的通用指南¶

对于新的用例，例如训练数据类型 (如 fp4 训练)，可以先在 prototype 文件夹中添加新的张量子类 torchao/prototype，但你也可以查看 AffineQuantizedTensor，如果那里大部分功能已经支持你想要做的事情，例如为完全相同的仿射量化添加 int3 kernel。如果您对特定新用例的操作有疑问，请随时提交 Issue。更多详情，请参考我们的量化概述页面。

贡献现有代码库

向 AffineQuantizedTensor 添加功能，例如使其可训练、添加张量并行支持等：torchao/dtypes/affine_quantized_tensor.py
添加新的量化 API：torchao/quantization/quant_api.py
添加新的量化原始操作，例如现有量化原始操作的微小变体：torchao/quantization/quant_primitives.py
添加新的自动调优 Triton kernel：torchao/kernel
添加新的自定义 CPU/CUDA/MPS kernel：torchao/csrc
将自定义 kernel 与 AffineQuantizedTensor 集成 (可能还需要新的 layout)：例如，添加稀疏 marlin AQT layout #621。我们还没有决定是否将 AffineQuantizedTensor 拆分为更多的张量子类。

添加高效 kernel¶

自定义 Triton kernel¶

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

您可能还需要定义自己的自动调优器。

自定义手写 kernel¶

CPU/CUDA/MPS 的自定义 kernel (实现) 可以通过 torchao/csrc 实现，例如 int4 cuda，并通过 torch.ops.my_custom_op 访问

调度¶

为了调度到针对 CPU/CUDA/MPS 设备优化的 kernel，我们可以在 __torch_function__ 或 __torch_dispatch__ 中检查调度条件并调度到目标操作符，例如，bfloat16 激活和 uint4 权重 kernel 的条件可以在这里找到。

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

Layout/TensorImpl¶

有时需要对量化权重进行打包以获得最佳性能。这可以使用 layout 进行抽象。完整示例请参见这里。

流程¶

张量子类实现后，我们还可以将其封装到工厂函数中，例如：: # convert from floating point tensor to my dtype tensor subclass 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，并移除任何不必要的图中断 (graph breaks)。您可以在运行脚本时添加 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 文件夹中放入更多涵盖以下用例的示例。

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

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

Kernel 微基准测试¶

在对模型进行性能测试之前，我们还可以对单个线性操作符 (或其他计算密集型/内存密集型) 操作符在不同输入维度下进行微基准测试，以了解加速效果。对于您想要测试的特定 kernel，您可以创建一个类似 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 是我们用于计算密集型模型的代表模型。

Llama * 基准测试 * 评估
SAM * 基准测试和评估

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

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