快速入門指南

在本快速入門指南中，我們將探討如何使用 torchao 執行基本的量化操作。首先，安裝最新的穩定版 torchao

pip install torchao

如果您更傾向於使用 nightly 版本，則可以使用以下命令進行安裝

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 僅權重（weight-only）非對稱每組（per-group）量化，利用 tinygemm int4mm CUDA kernel 實現高效的混合資料型別矩陣乘法

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

量化後的模型現在可以使用了！請注意，量化邏輯是透過 tensor 子類插入的，因此整體模型結構沒有改變；只有權重張量被更新了，而 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

後續步驟

在本快速入門指南中，我們學習瞭如何使用 torchao 量化一個簡單模型。要了解 torchao 支援的不同工作流程，請參閱我們的主要 README。要獲取 torchao 量化的更詳細概述，請訪問此頁面。

最後，如果您想為 torchao 貢獻力量，請不要忘記檢視我們的貢獻者指南以及 Github 上的適合初學者的問題列表！