快速入门指南¶

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

pip install torchao

如果您想使用 nightly 版本，则可以使用以下命令安装 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 (
    benchmark_model,
    unwrap_tensor_subclass,
)

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 导出量化¶

PyTorch 2 导出量化是一个完整的图量化工作流，主要用于静态量化。它针对需要量化输入和输出激活以及权重的硬件，并依赖于识别算子模式来做出量化决策（例如，线性 - relu）。PT2E 量化会生成一个带有在算子周围插入量化和反量化操作的模式，并在降低过程中，量化算子模式将被融合到实际的量化算子中。目前有两种典型的降低路径：1. torch.compile 通过 inductor 降低 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 导出量化

建模用户

后端开发人员（请同时查看所有建模用户文档）

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

下一步¶

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

最后，如果您想为 torchao 做贡献，请不要忘记查看我们的贡献者指南和 Github 上的 “第一个好问题”列表！

快速入门指南¶

第一个量化示例¶

PyTorch 2 导出量化¶

下一步¶

文档

教程

资源