torchao

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

使用以下命令安装 torchao。

# Updating 🤗 Transformers to the latest version, as the example script below uses the new auto compilation
pip install --upgrade torch torchao transformers

torchao 支持多种针对不同数据类型（int4、float8、仅权重等）的量化类型。从 0.10.0 版本开始，torchao 通过 AOBaseConfig API 提供增强的灵活性，允许进行更自定义的量化配置，并完全访问 torchao 库中提供的技术。

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

手动

自动

创建一个 TorchAoConfig 并指定要量化的权重的量化类型和 group_size。将 cache_implementation 设置为 "static" 以自动 torch.compile forward 方法。

通过将 device_map 更改为 "cpu" 并将 layout 更改为 Int4CPULayout()，在 CPU 上运行量化模型。这仅在 torchao 0.8.0+ 中可用。

在 torchao 0.10.0+ 中，您可以使用更灵活的 AOBaseConfig 方法来代替字符串标识符

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

# Using AOBaseConfig instance (torchao >= 0.10.0)
quant_config = Int4WeightOnlyConfig(group_size=128)
quantization_config = TorchAoConfig(quant_type=quant_config)

# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype="auto",
    device_map="auto",
    quantization_config=quantization_config
)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

# auto-compile the quantized model with `cache_implementation="static"` to get speed up
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))

可用的量化方案

TorchAO 提供了多种量化配置

Int4WeightOnlyConfig
Int8WeightOnlyConfig
Int8DynamicActivationInt8WeightConfig
Float8WeightOnlyConfig

每个配置都可以使用诸如 group_size、scheme 和 layout 等参数进一步自定义，以针对特定硬件和模型架构进行优化。

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

⚠️ 弃用警告

从 0.10.0 版本开始，用于量化配置的基于字符串的 API（例如，TorchAoConfig("int4_weight_only", group_size=128)）已被弃用，并将在未来版本中删除。

请改用新的基于 AOBaseConfig 的方法
# Old way (deprecated)
quantization_config = TorchAoConfig("int4_weight_only", group_size=128)

# New way (recommended)
from torchao.quantization import Int4WeightOnlyConfig
quant_config = Int4WeightOnlyConfig(group_size=128)
quantization_config = TorchAoConfig(quant_type=quant_config)
新的 API 提供了更高的灵活性、更好的类型安全性和对 torchao 中提供的全套功能的访问。

迁移指南

以下是如何从常见的字符串标识符迁移到其 AOBaseConfig 等效项的方法

旧的字符串 API 新的 AOBaseConfig API

"int4_weight_only" Int4WeightOnlyConfig()

"int8_weight_only" Int8WeightOnlyConfig()

"int8_dynamic_activation_int8_weight" Int8DynamicActivationInt8WeightConfig()

所有配置对象都接受用于自定义的参数（例如，group_size、scheme、layout）。

旧的字符串 API	新的 `AOBaseConfig` API
`"int4_weight_only"`	`Int4WeightOnlyConfig()`
`"int8_weight_only"`	`Int8WeightOnlyConfig()`
`"int8_dynamic_activation_int8_weight"`	`Int8DynamicActivationInt8WeightConfig()`

以下是 torchao < 0.9.0 的 API

import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer

quantization_config = TorchAoConfig("int4_weight_only", group_size=128)
quantized_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype="auto",
    device_map="auto",
    quantization_config=quantization_config
)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

# auto-compile the quantized model with `cache_implementation="static"` to get speed up
output = quantized_model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))

运行以下代码以基准测试量化模型的性能。

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", torch_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() 来使用推荐的设置。

序列化

torchao 实现了 torch.Tensor 子类，以在支持新的量化 torch.Tensor 格式时实现最大灵活性。Safetensors 序列化和反序列化不适用于 torchao。

为了避免任意用户代码执行，torchao 在 torch.load 中设置 weights_only=True，以确保仅加载张量。任何已知的用户函数都可以使用 add_safe_globals 列入白名单。

# don't serialize model with Safetensors
output_dir = "llama3-8b-int4wo-128"
quantized_model.save_pretrained("llama3-8b-int4wo-128", safe_serialization=False)

资源

为了更好地了解预期性能，请查看各种模型在 CUDA 和 XPU 后端上的基准测试。

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

< > 在 GitHub 上更新