快速入門指南¶

在本快速入門指南中，我們將探索如何使用 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 匯出量化¶

下一步¶

文件

教程

資源