torchao

torchao est une bibliothèque d’optimisation d’architecture PyTorch prenant en charge des types de données haute performance personnalisés, la quantification et la sparsité. Elle est composable avec les fonctionnalités natives de PyTorch, telles que torch.compile, pour une inférence et un entraînement encore plus rapides.

Consultez le tableau ci-dessous pour connaître les fonctionnalités supplémentaires de torchao.

Fonctionnalité	Description
Entraînement tenant compte de la quantification (QAT)	Entraîner des modèles quantifiés avec une perte de précision minimale (voir README QAT)
Entraînement Float8	Entraînement à haut débit avec des formats float8 (voir la documentation de torchtitan et de Accelerate)
Prise en charge de la sparsité	Sparsity semi-structurée (2:4) pour une inférence plus rapide (voir l'article de blog Accelerating Neural Network Training with Semi-Structured (2:4) Sparsity)
Quantification de l'optimiseur	Réduire la mémoire de l'état de l'optimiseur avec les variantes 4 et 8 bits d'Adam
Quantification de la mémoire cache KV	Permet une inférence sur de longs contextes avec une consommation mémoire réduite (voir Quantification du cache KV)
Prise en charge des noyaux personnalisés	utilisez vos propres opérateurs compatibles avec `torch.compile`
FSDP2	Composable avec FSDP2 pour l'entraînement

Consultez le README.md de torchao pour plus de détails sur la bibliothèque.

torchao prend en charge les techniques de quantification suivantes.

Quantification dynamique Float8 A16W8
Quantification A16W8 Float8 WeightOnly
Quantification dynamique Int8 A8W8
Quantification A16W8 Int8 des poids uniquement
Quantification A16W4 Int4 des poids uniquement
Quantification A16W4 Int4 des poids uniquement + Sparsité 2:4
Autoquantization

torchao prend également en charge la configuration au niveau du module en spécifiant un dictionnaire associant le nom entièrement qualifié du module à sa configuration de quantification correspondante. Cela permet de sauter la quantification de certaines couches et d’utiliser des configurations de quantification différentes pour différents modules.

Consultez le tableau ci-dessous pour vérifier si votre matériel est compatible.

Composant	Compatibilité
Versions de CUDA	✅ cu118, cu126, cu128
Versions XPU	✅ pytorch2.8
CPU	✅ modifiez `device_map="cpu"` (voir les exemples ci-dessous)

Installez torchao depuis PyPi ou l'index PyTorch à l'aide des commandes suivantes.

PyPi

Index de PyTorch

torchao >= 0.15.0 est requis. L'API basée sur des chaînes de caractères (par exemple, TorchAoConfig("int4_weight_only")) a été supprimée — utilisez à la place des objets AOBaseConfig (voir les exemples ci-dessous).

Exemples de quantification

TorchAO propose une variété de configurations de quantification. Chaque configuration peut être personnalisée à l'aide de paramètres tels que group_size, scheme et layout afin d'optimiser le fonctionnement sur des matériels et des architectures de modèles spécifiques.

Pour obtenir la liste complète des configurations disponibles, consultez la documentation de l’API de quantification.

Vous pouvez choisir manuellement les types et les paramètres de quantisation, ou sélectionner automatiquement les types de quantisation.

Créez un TorchAoConfig et spécifiez le type de quantification ainsi que la group_size des poids à quantiser (pour int8 weight only et int4 weight only). Définissez cache_implementation sur "static" afin de torch.compile automatiquement la méthode forward.

Nous présenterons des exemples de méthodes de quantisation recommandées en fonction du matériel, par exemple GPU A100, GPU H100, CPU.

torchao compile automatiquement le modèle lors de la première inférence si nous définissons cache_implementation="static". Le modèle est recompilé à chaque modification de la taille du lot ou de max_new_tokens. Passez disable_compile=True dans generate() pour quantiser sans compilation.

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

Quantification par module

1. Ignorer la quantification pour certaines couches

Avec FqnToConfig, nous pouvons spécifier une configuration par défaut pour toutes les couches tout en ignorant la quantification pour certaines couches.

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. Quantifier différentes couches avec différentes configurations de quantification (sans expression régulière)

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. Quantifier différentes couches avec différentes configurations de quantification (avec des expressions régulières)

Nous pouvons également utiliser des expressions régulières pour spécifier la configuration de tous les modules dont le module_fqn correspond à l'expression régulière. Toutes les expressions régulières doivent commencer par re:. Par exemple, re:layers\..*\.gate_proj correspondra à toutes les couches comme layers.0.gate_proj. Consultez la documentation.

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)

Sérialisation

La sauvegarde du modèle quantifié avec save_pretrained (au format safetensors) n'est prise en charge que pour torchao >= v0.15. Pour toute version antérieure, il n'est possible de sauvegarder manuellement que des points de contrôle non sécurisés .bin à l'aide de torch.save.

save-locally

push-to-huggingface-hub

Chargement de modèles quantifiés

Le chargement d'un modèle quantifié dépend du schéma de quantisation. Pour des schémas de quantisation tels que int8 et float8, vous pouvez quantiser le modèle sur n'importe quel appareil et le charger également sur n'importe quel appareil. L'exemple ci-dessous illustre la quantisation d'un modèle sur le CPU, puis son chargement sur 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))

Pour int4, le modèle ne peut être chargé que sur le même appareil sur lequel il a été quantifié, car la disposition est spécifique à l’appareil. L’exemple ci-dessous illustre la quantification et le chargement d’un modèle sur le 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))

Ressources

Pour avoir une meilleure idée des performances attendues, consultez les benchmarks de divers modèles avec les backends CUDA et XPU. Vous pouvez également exécuter le code ci-dessous pour effectuer vous-même le benchmark d'un modèle.

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

Pour obtenir les meilleures performances, vous pouvez utiliser les paramètres recommandés en appelant torchao.quantization.utils.recommended_inductor_config_setter()

Reportez-vous à Autres techniques de quantification disponibles pour plus d’exemples et de documentation.

Problèmes

Si vous rencontrez des problèmes avec l'intégration Transformers, veuillez ouvrir un ticket sur le dépôt Transformers. Pour les problèmes directement liés à torchao, veuillez ouvrir un ticket sur le dépôt torchao.

Update on GitHub