快速入门指南¶

在此快速入门指南中，我们将探讨如何使用 torchao 执行基本量化。首先，安装最新的稳定版 torchao

pip install torchao

如果您喜欢使用每夜发布版本，可以使用以下命令安装 torchao

pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121

torchao 与 PyTorch 最新的 3 个主要版本兼容，您也需要安装这些版本（详细说明）

pip install torch

第一个量化示例¶

torchao 中量化的主要入口点是 quantize_ API。此函数会就地修改您的模型，以根据用户配置插入自定义量化逻辑。本指南中的所有代码都可以在此示例脚本中找到。首先，让我们设置我们的玩具模型

import copy
import torch

class ToyLinearModel(torch.nn.Module):
    def __init__(self, m: int, n: int, k: int):
        super().__init__()
        self.linear1 = torch.nn.Linear(m, n, bias=False)
        self.linear2 = torch.nn.Linear(n, k, bias=False)

    def forward(self, x):
        x = self.linear1(x)
        x = self.linear2(x)
        return x

model = ToyLinearModel(1024, 1024, 1024).eval().to(torch.bfloat16).to("cuda")

# Optional: compile model for faster inference and generation
model = torch.compile(model, mode="max-autotune", fullgraph=True)
model_bf16 = copy.deepcopy(model)

现在我们调用主要的量化 API，将模型中的线性权重就地量化为 int4。更具体地说，这应用了 uint4 仅权重非对称每组量化，利用 tinygemm int4mm CUDA 内核实现高效的混合数据类型矩阵乘法

# torch 2.4+ only
from torchao.quantization import Int4WeightOnlyConfig, quantize_
quantize_(model, Int4WeightOnlyConfig(group_size=32))

量化模型现在可以使用了！请注意，量化逻辑是通过张量子类插入的，因此整体模型结构没有改变；只有权重张量被更新，但 nn.Linear 模块仍然是 nn.Linear 模块

>>> model.linear1
Linear(in_features=1024, out_features=1024, weight=AffineQuantizedTensor(shape=torch.Size([1024, 1024]), block_size=(1, 32), device=cuda:0, _layout=TensorCoreTiledLayout(inner_k_tiles=8), tensor_impl_dtype=torch.int32, quant_min=0, quant_max=15))

>>> model.linear2
Linear(in_features=1024, out_features=1024, weight=AffineQuantizedTensor(shape=torch.Size([1024, 1024]), block_size=(1, 32), device=cuda:0, _layout=TensorCoreTiledLayout(inner_k_tiles=8), tensor_impl_dtype=torch.int32, quant_min=0, quant_max=15))

首先，验证 int4 量化模型的大小是否大约是原始 bfloat16 模型的四分之一

>>> import os
>>> torch.save(model, "/tmp/int4_model.pt")
>>> torch.save(model_bf16, "/tmp/bfloat16_model.pt")
>>> int4_model_size_mb = os.path.getsize("/tmp/int4_model.pt") / 1024 / 1024
>>> bfloat16_model_size_mb = os.path.getsize("/tmp/bfloat16_model.pt") / 1024 / 1024

>>> print("int4 model size: %.2f MB" % int4_model_size_mb)
int4 model size: 1.25 MB

>>> print("bfloat16 model size: %.2f MB" % bfloat16_model_size_mb)
bfloat16 model size: 4.00 MB

接下来，我们证明量化模型不仅更小，而且速度更快！

from torchao.utils import (
    TORCH_VERSION_AT_LEAST_2_5,
    benchmark_model,
    unwrap_tensor_subclass,
)

# Temporary workaround for tensor subclass + torch.compile
# Only needed for torch version < 2.5
if not TORCH_VERSION_AT_LEAST_2_5:
    unwrap_tensor_subclass(model)

num_runs = 100
torch._dynamo.reset()
example_inputs = (torch.randn(1, 1024, dtype=torch.bfloat16, device="cuda"),)
bf16_time = benchmark_model(model_bf16, num_runs, example_inputs)
int4_time = benchmark_model(model, num_runs, example_inputs)

print("bf16 mean time: %0.3f ms" % bf16_time)
print("int4 mean time: %0.3f ms" % int4_time)
print("speedup: %0.1fx" % (bf16_time / int4_time))

在具有 80GB 内存的单个 A100 GPU 上，这将打印

bf16 mean time: 30.393 ms
int4 mean time: 4.410 ms
speedup: 6.9x

PyTorch 2 Export 量化¶

PyTorch 2 Export 量化是主要用于静态量化的全图量化工作流程。它针对需要输入、输出激活和权重都进行量化的硬件，并依赖于识别运算符模式来做出量化决策（例如线性 - relu）。PT2E 量化会生成一个在运算符周围插入量化和去量化操作的模式，在下沉过程中，量化运算符模式将被融合为实际的量化操作。目前有两种典型的下沉路径：1. 通过 inductor 下沉的 torch.compile 2. 通过委托的 ExecuTorch

这里我们展示一个 X86InductorQuantizer 的例子

API 示例

import torch
from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e
from torch.export import export
from torchao.quantization.pt2e.quantizer.x86_inductor_quantizer import (
    X86InductorQuantizer,
    get_default_x86_inductor_quantization_config,
)

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(5, 10)

   def forward(self, x):
       return self.linear(x)

# initialize a floating point model
float_model = M().eval()

# define calibration function
def calibrate(model, data_loader):
    model.eval()
    with torch.no_grad():
        for image, target in data_loader:
            model(image)

# Step 1. program capture
m = export(m, *example_inputs).module()
# we get a model with aten ops

# Step 2. quantization
# backend developer will write their own Quantizer and expose methods to allow
# users to express how they
# want the model to be quantized
quantizer = X86InductorQuantizer()
quantizer.set_global(xiq.get_default_x86_inductor_quantization_config())

# or prepare_qat_pt2e for Quantization Aware Training
m = prepare_pt2e(m, quantizer)

# run calibration
# calibrate(m, sample_inference_data)
m = convert_pt2e(m)

# Step 3. lowering
# lower to target backend

# Optional: using the C++ wrapper instead of default Python wrapper
import torch._inductor.config as config
config.cpp_wrapper = True

with torch.no_grad():
    optimized_model = torch.compile(converted_model)

    # Running some benchmark
    optimized_model(*example_inputs)

请遵循这些教程以开始使用 PyTorch 2 Export 量化

建模用户

后端开发者（也请查阅所有建模用户文档）

如何为 PyTorch 2 导出量化编写量化器

下一步¶

在此快速入门指南中，我们学习了如何使用 torchao 量化一个简单的模型。要了解更多关于 torchao 支持的不同工作流程，请参阅我们的主要 README。有关 torchao 中量化的更详细概述，请访问此页面。

最后，如果您想为 torchao 做出贡献，请不要忘记查看我们的贡献者指南和 Github 上我们的适合初学者的问题列表！

快速入门指南¶

第一个量化示例¶

PyTorch 2 Export 量化¶

下一步¶

文档

教程

资源