torchao

torchao 是一个 PyTorch 架构优化库，支持自定义高性能数据类型、量化和稀疏性。它可以与 PyTorch 原生功能（如 torch.compile）组合使用，以实现更快的推理和训练。

请参阅下表了解 torchao 的其他功能。

功能	描述
量化感知训练 (QAT)	以极小的精度损失训练量化模型（参见 QAT README）
Float8 训练	使用 float8 格式进行高吞吐量训练（参见 torchtitan 和 Accelerate 文档）
稀疏性支持	用于加速推理的半结构化（2:4）稀疏性（参见使用半结构化（2:4）稀疏性加速神经网络训练博客文章）
优化器量化	使用 Adam 的 4 位和 8 位变体来减少优化器状态内存
KV 缓存量化	启用长上下文推理，降低内存占用（参见 KV Cache 量化）
自定义内核支持	使用您自己的与 `torch.compile` 兼容的操作
FSDP2	可与 FSDP2 组合用于训练

有关该库的更多详细信息，请参阅 torchao 的 README.md。

torchao 支持以下量化技术。

A16W8 Float8 动态量化
A16W8 Float8 仅权重量化
A8W8 Int8 动态量化
A16W8 Int8 仅权重量化
A16W4 Int4 仅权重量化
A16W4 Int4 仅权重量化 + 2:4 稀疏性
自动量化

torchao 还支持模块级别的配置，通过指定一个字典，其中包含模块的完全限定名称及其对应的量化配置。这使得可以跳过对某些层的量化，并为不同的模块使用不同的量化配置。

请查看下表，以确认您的硬件是否兼容。

组件	兼容性
CUDA 版本	✅ cu118、cu126、cu128
XPU 版本	✅ pytorch2.8
CPU	✅ 更改 `device_map="cpu"`（参见下文示例）

使用以下命令从 PyPi 或 PyTorch 索引安装 torchao。

PyPi

PyTorch 索引

需要 torchao >= 0.15.0。基于字符串的 API（例如 TorchAoConfig("int4_weight_only")）已被移除——请改用 AOBaseConfig 对象（参见下方示例）。

量化示例

TorchAO 提供了多种量化配置。每种配置都可以通过 group_size、scheme 和 layout 等参数进行进一步自定义，以针对特定的硬件和模型架构进行优化。

有关可用配置的完整列表，请参阅量化 API 文档。

您可以手动选择量化类型和设置，也可以自动选择量化类型。

创建一个 TorchAoConfig，并指定要量化的权重的量化类型和 group_size（适用于 int8 weight only 和 int4 weight only）。将 cache_implementation 设置为 "static"，以自动 torch.compile 前向方法。

我们将根据硬件（例如 A100 GPU、H100 GPU、CPU）展示推荐的量化方法示例。

如果将 cache_implementation="static" 设置为 "static"，torchao 会在首次推理时自动编译模型。每当 batch size 或 max_new_tokens 发生变化时，模型都会重新编译。在 generate() 中传入 disable_compile=True 即可在不进行编译的情况下执行量化。

H100 GPU

float8-dynamic-and-weight-only

int4-weight-only

A100 GPU

int8-dynamic-and-weight-only

int4-weight-only

Intel XPU

int8-dynamic-and-weight-only

int4-weight-only

CPU

int8-dynamic-and-weight-only

int4-weight-only

按模块量化

1. 跳过特定层的量化

使用 FqnToConfig，我们可以为所有层指定默认配置，同时跳过对某些层的量化。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig

model_id = "meta-llama/Llama-3.1-8B-Instruct"

from torchao.quantization import Int4WeightOnlyConfig, FqnToConfig
config = Int4WeightOnlyConfig(group_size=128)

# set default to int4 (for linears), and skip quantizing `model.layers.0.self_attn.q_proj`
quant_config = FqnToConfig({"_default": config, "model.layers.0.self_attn.q_proj": None})
quantization_config = TorchAoConfig(quant_type=quant_config)
quantized_model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", dtype=torch.bfloat16, quantization_config=quantization_config)
# lm_head is not quantized and model.layers.0.self_attn.q_proj is not quantized
print("quantized model:", quantized_model)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Manual Testing
prompt = "Hey, are you conscious? Can you talk to me?"
inputs = tokenizer(prompt, return_tensors="pt").to(quantized_model.device, quantized_model.dtype)
generated_ids = quantized_model.generate(**inputs, max_new_tokens=128)
output_text = tokenizer.batch_decode(
    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print(output_text)

2. 使用不同的量化配置对不同的层进行量化（不使用正则表达式）

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig

model_id = "facebook/opt-125m"

from torchao.quantization import Int4WeightOnlyConfig, FqnToConfig, Int8DynamicActivationInt4WeightConfig, IntxWeightOnlyConfig, PerAxis, MappingType

weight_dtype = torch.int8
granularity = PerAxis(0)
mapping_type = MappingType.ASYMMETRIC
embedding_config = IntxWeightOnlyConfig(
    weight_dtype=weight_dtype,
    granularity=granularity,
    mapping_type=mapping_type,
)
linear_config = Int8DynamicActivationInt4WeightConfig(group_size=128)
quant_config = FqnToConfig({"_default": linear_config, "model.decoder.embed_tokens": embedding_config, "model.decoder.embed_positions": None})
# set `include_embedding` to True in order to include embedding in quantization
# when `include_embedding` is True, we'll remove input embedding from `modules_not_to_convert` as well
quantization_config = TorchAoConfig(quant_type=quant_config, include_embedding=True)
quantized_model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu", dtype=torch.bfloat16, quantization_config=quantization_config)
print("quantized model:", quantized_model)
# make sure embedding is quantized
print("embed_tokens weight:", quantized_model.model.decoder.embed_tokens.weight)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Manual Testing
prompt = "Hey, are you conscious? Can you talk to me?"
inputs = tokenizer(prompt, return_tensors="pt").to("cpu", quantized_model.dtype)
generated_ids = quantized_model.generate(**inputs, max_new_tokens=128, cache_implementation="static")
output_text = tokenizer.batch_decode(
    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print(output_text)

3. 使用不同的量化配置（通过正则表达式）对不同的层进行量化

我们同样可以使用正则表达式来指定所有 module_fqn 与正则表达式匹配的模块的配置，所有正则表达式都应以 re: 开头，例如 re:layers\..*\.gate_proj 将匹配所有类似 layers.0.gate_proj 的层。请参阅此处获取文档。

import logging

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig

# Configure logging to see warnings and debug information
logging.basicConfig(
    level=logging.INFO, format="%(name)s - %(levelname)s - %(message)s"
)

# Enable specific loggers that might contain the serialization warnings
logging.getLogger("transformers").setLevel(logging.INFO)
logging.getLogger("torchao").setLevel(logging.INFO)
logging.getLogger("safetensors").setLevel(logging.INFO)
logging.getLogger("huggingface_hub").setLevel(logging.INFO)

model_id = "facebook/opt-125m"

from torchao.quantization import (
    Float8DynamicActivationFloat8WeightConfig,
    Int4WeightOnlyConfig,
    IntxWeightOnlyConfig,
    PerRow,
    PerAxis,
    FqnToConfig,
    Float8Tensor,
    Int4TilePackedTo4dTensor,
    IntxUnpackedToInt8Tensor,
)

float8dyn = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
int4wo = Int4WeightOnlyConfig(int4_packing_format="tile_packed_to_4d")
intxwo = IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerAxis(0))

qconfig_dict = {
    # highest priority
    "model.decoder.layers.3.self_attn.q_proj": int4wo,
    "model.decoder.layers.3.self_attn.k_proj": int4wo,
    "model.decoder.layers.3.self_attn.v_proj": int4wo,
    # vllm
    "model.decoder.layers.3.self_attn.qkv_proj": int4wo,

    "re:model\.decoder\.layers\..+\.self_attn\.q_proj": float8dyn,
    "re:model\.decoder\.layers\..+\.self_attn\.k_proj": float8dyn,
    "re:model\.decoder\.layers\..+\.self_attn\.v_proj": float8dyn,
    # this should not take effect and we'll fallback to _default
    # since no full mach (missing `j` in the end)
    "re:model\.decoder\.layers\..+\.self_attn\.out_pro": float8dyn,
    # vllm
    "re:model\.decoder\.layers\..+\.self_attn\.qkv_proj": float8dyn,

    "_default": intxwo,
}
quant_config = FqnToConfig(qconfig_dict)
quantization_config = TorchAoConfig(quant_type=quant_config)
quantized_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    quantization_config=quantization_config,
)
print("quantized model:", quantized_model)
tokenizer = AutoTokenizer.from_pretrained(model_id)
for i in range(12):
    if i == 3:
        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.q_proj.weight, Int4TilePackedTo4dTensor)
        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Int4TilePackedTo4dTensor)
        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Int4TilePackedTo4dTensor)
    else:
        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.q_proj.weight, Float8Tensor)
        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Float8Tensor)
        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Float8Tensor)
    assert isinstance(quantized_model.model.decoder.layers[i].self_attn.out_proj.weight, IntxUnpackedToInt8Tensor)

# Manual Testing
prompt = "What are we having for dinner?"
print("Prompt:", prompt)
inputs = tokenizer(
    prompt,
    return_tensors="pt",
).to(quantized_model.device, quantized_model.dtype)
# setting temperature to 0 to make sure result deterministic
generated_ids = quantized_model.generate(**inputs, max_new_tokens=128, temperature=0)

correct_output_text = tokenizer.batch_decode(
    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print("Response:", correct_output_text[0][len(prompt) :])


# Load model from saved checkpoint
reloaded_model = AutoModelForCausalLM.from_pretrained(
    save_to,
    device_map=torch.accelerator.current_accelerator(),
    torch_dtype=torch.bfloat16,
    # quantization_config=quantization_config,
)

generated_ids = reloaded_model.generate(**inputs, max_new_tokens=128, temperature=0)
output_text = tokenizer.batch_decode(
    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print("Response:", output_text[0][len(prompt) :])

assert(correct_output_text == output_text)

序列化

使用 save_pretrained（以 safetensors 格式）保存量化模型仅在 torchao >= v0.15 中受支持。对于更低版本，只能通过 torch.save 手动保存为不安全的 .bin 检查点。

save-locally

push-to-huggingface-hub

加载量化模型

加载量化模型取决于量化方案。对于 int8 和 float8 等量化方案，您可以在任意设备上对模型进行量化，并在任意设备上加载它。下面的示例演示了在 CPU 上量化模型，然后在 CUDA 或 XPU 上加载该模型。

import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int8WeightOnlyConfig

quant_config = Int8WeightOnlyConfig(group_size=128)
quantization_config = TorchAoConfig(quant_type=quant_config)

# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    dtype="auto",
    device_map="cpu",
    quantization_config=quantization_config
)
# save the quantized model
output_dir = "llama-3.1-8b-torchao-int8"
quantized_model.save_pretrained(output_dir)

# reload the quantized model
reloaded_model = AutoModelForCausalLM.from_pretrained(
    output_dir,
    device_map="auto",
    dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(reloaded_model.device.type)

output = reloaded_model.generate(**input_ids, max_new_tokens=10)
print(tokenizer.decode(output[0], skip_special_tokens=True))

对于 int4，模型只能加载到其量化时所在的同一设备上，因为内存布局是特定于设备的。下面的示例演示了在 CPU 上量化并加载模型。

import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int4WeightOnlyConfig
from torchao.quantization.quantize_.workflows import Int4PackingFormat

quant_config = Int4WeightOnlyConfig(group_size=128, int4_packing_format=Int4PackingFormat.PLAIN_INT32)
quantization_config = TorchAoConfig(quant_type=quant_config)

# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    dtype="auto",
    device_map="cpu",
    quantization_config=quantization_config
)
# save the quantized model
output_dir = "llama-3.1-8b-torchao-int4-cpu"
quantized_model.save_pretrained(output_dir)

# reload the quantized model
reloaded_model = AutoModelForCausalLM.from_pretrained(
    output_dir,
    device_map="cpu",
    dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to(reloaded_model.device.type)

output = reloaded_model.generate(**input_ids, max_new_tokens=10)
print(tokenizer.decode(output[0], skip_special_tokens=True))

资源

为了更好地了解预期性能，请查看使用 CUDA 和 XPU 后端的各种模型的基准测试。您也可以运行以下代码自行对模型进行基准测试。

from torch._inductor.utils import do_bench_using_profiling
from typing import Callable

def benchmark_fn(func: Callable, *args, **kwargs) -> float:
    """Thin wrapper around do_bench_using_profiling"""
    no_args = lambda: func(*args, **kwargs)
    time = do_bench_using_profiling(no_args)
    return time * 1e3

MAX_NEW_TOKENS = 1000
print("int4wo-128 model:", benchmark_fn(quantized_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static"))

bf16_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", dtype=torch.bfloat16)
output = bf16_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static") # auto-compile
print("bf16 model:", benchmark_fn(bf16_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static"))

为了获得最佳性能，您可以通过调用 torchao.quantization.utils.recommended_inductor_config_setter() 来使用推荐设置。

有关更多示例和文档，请参阅其他可用的量化技术。

问题

如果您在使用 Transformers 集成时遇到任何问题，请在 Transformers 仓库中提交 issue。对于与 torchao 直接相关的问题，请在 torchao 仓库中提交 issue。

Update on GitHub