torchao

torchao é uma biblioteca de otimização de arquitetura do PyTorch com suporte a tipos de dados de alto desempenho personalizados, quantização e esparsidade. Ela é composável com recursos nativos do PyTorch, como torch.compile, para inferência e treinamento ainda mais rápidos.

Veja a tabela abaixo para mais recursos do torchao.

Recurso	Descrição
Quantization Aware Training (QAT)	Treine modelos quantizados com perda mínima de precisão (veja QAT README)
Float8 Training	Treinamento de alto rendimento com formatos float8 (veja a documentação de torchtitan e Accelerate)
Suporte a Esparsidade	Esparsidade semi-estruturada (2:4) para inferência mais rápida (veja o post no blog Acelerando o Treinamento de Redes Neurais com Esparsidade Semi-Estruturada (2:4))
Quantização do Otimizador	Reduza a memória do estado do otimizador com as variantes de 4 e 8 bits do Adam
Quantização de Cache KV	Permite inferência de contexto longo com menor uso de memória (veja Quantização de KV Cache)
Suporte a Kernels Personalizados	use suas próprias ops compatíveis com `torch.compile`
FSDP2	Componível com FSDP2 para treinamento

Consulte o README.md do torchao para mais detalhes sobre a biblioteca.

O torchao suporta as técnicas de quantização abaixo.

Quantização dinâmica Float8 A16W8
Quantização A16W8 Float8 WeightOnly
Quantização Dinâmica Int8 A8W8
Quantização apenas de pesos Int8 A16W8
Quantização apenas de pesos Int4 A16W4
Quantização apenas de pesos Int4 A16W4 + Esparsidade 2:4
Autoquantização

O torchao também suporta configuração em nível de módulo, especificando um dicionário com o nome totalmente qualificado do módulo e sua configuração de quantização correspondente. Isso permite pular a quantização de determinadas camadas e usar configurações de quantização diferentes para diferentes módulos.

Consulte a tabela abaixo para verificar se seu hardware é compatível.

Componente	Compatibilidade
Versões do CUDA	✅ cu118, cu126, cu128
Versões de XPU	✅ pytorch2.8
CPU	✅ altere `device_map="cpu"` (veja os exemplos abaixo)

Instale o torchao a partir do PyPi ou do índice do PyTorch com os seguintes comandos.

PyPi

Índice do PyTorch

torchao >= 0.15.0 is required. The string-based API (e.g., TorchAoConfig("int4_weight_only")) has been removed — use AOBaseConfig objects instead (see examples below).

Exemplos de quantização

O TorchAO oferece uma variedade de configurações de quantização. Cada configuração pode ser personalizada ainda mais com parâmetros como group_size, scheme e layout para otimizar para arquiteturas de hardware e de modelo específicas.

Para uma lista completa das configurações disponíveis, consulte a documentação da API de quantização.

Você pode escolher manualmente os tipos e configurações de quantização ou selecionar automaticamente os tipos de quantização.

Crie um TorchAoConfig e especifique o tipo de quantização e o group_size dos pesos a serem quantizados (para int8 weight only e int4 weight only). Defina o cache_implementation como "static" para torch.compile automaticamente o método forward.

Mostraremos exemplos de métodos de quantização recomendados com base nos hardwares, por exemplo, GPU A100, GPU H100, CPU.

O torchao compila o modelo automaticamente durante a primeira inferência se definirmos cache_implementation="static". O modelo é recompilado sempre que o tamanho do lote ou max_new_tokens é modificado. Passe disable_compile=True em generate() para quantizar sem compilação.

GPU H100

float8-dynamic-and-weight-only

int4-weight-only

GPU A100

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

Quantização por Módulo

1. Pular a quantização para determinadas camadas

Com FqnToConfig, podemos especificar uma configuração padrão para todas as camadas, ignorando a quantização para determinadas camadas.

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. Quantizar diferentes camadas com diferentes configurações de quantização (sem regex)

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. Quantizar diferentes camadas com diferentes configurações de quantização (com regex)

Também podemos usar regex para especificar a configuração de todos os módulos que possuem module_fqn que corresponde à regex. Todas as regex devem começar com re:. Por exemplo, re:layers\..*\.gate_proj corresponderá a todas as camadas como layers.0.gate_proj. Veja a documentação aqui.

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)

Serialização

Salvar o modelo quantizado com save_pretrained (no formato safetensors) é suportado apenas para torchao >= v0.15. Para qualquer versão anterior, é possível apenas salvar manualmente como checkpoints inseguros .bin com torch.save.

save-locally

push-to-huggingface-hub

Carregando modelos quantizados

Carregar um modelo quantizado depende do esquema de quantização. Para esquemas de quantização, como int8 e float8, você pode quantizar o modelo em qualquer dispositivo e também carregá-lo em qualquer dispositivo. O exemplo abaixo demonstra a quantização de um modelo na CPU e, em seguida, o carregamento em CUDA ou 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))

Para int4, o modelo só pode ser carregado no mesmo dispositivo em que foi quantizado, pois o layout é específico para o dispositivo. O exemplo abaixo demonstra a quantização e o carregamento de um modelo na 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))

Recursos

Para ter uma melhor noção do desempenho esperado, consulte os benchmarks de vários modelos com backends CUDA e XPU. Você também pode executar o código abaixo para fazer o benchmark de um modelo por conta própria.

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"))

Para obter o melhor desempenho, você pode usar as configurações recomendadas chamando torchao.quantization.utils.recommended_inductor_config_setter()

Consulte Outras Técnicas de Quantização Disponíveis para mais exemplos e documentação.

Problemas

Se você encontrar algum problema com a integração do Transformers, abra uma issue no repositório Transformers. Para problemas diretamente relacionados ao torchao, abra uma issue no repositório torchao.

Update on GitHub