GitHub - pytorch/ao: PyTorch native quantization and sparsity for training and inference (original) (raw)

TorchAO

PyTorch-Native Training-to-Serving Model Optimization

📣 Latest News

🌅 Overview

TorchAO is a PyTorch-native model optimization framework leveraging quantization and sparsity to provide an end-to-end, training-to-serving workflow for AI models. TorchAO works out-of-the-box with torch.compile() and FSDP2 across most HuggingFace PyTorch models. Key features include:

Check out our docs for more details!

From the team that brought you the fast series:

🚀 Quick Start

First, install TorchAO. We recommend installing the latest stable version:

Other installation options 

# Nightly
pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu126

# Different CUDA versions
pip install torchao --index-url https://download.pytorch.org/whl/cu126  # CUDA 12.6
pip install torchao --index-url https://download.pytorch.org/whl/cpu    # CPU only

# For developers
USE_CUDA=1 python setup.py develop

Quantize your model weights to int4!

from torchao.quantization import Int4WeightOnlyConfig, quantize_
quantize_(model, Int4WeightOnlyConfig(group_size=32))

Compared to a torch.compiled bf16 baseline, your quantized model should be significantly smaller and faster on a single A100 GPU:

int4 model size: 1.25 MB
bfloat16 model size: 4.00 MB
compression ratio: 3.2

bf16 mean time: 30.393 ms
int4 mean time: 4.410 ms
speedup: 6.9x

For the full model setup and benchmark details, check out our quick start guide. Alternatively, try quantizing your favorite model using our HuggingFace space!

🔗 Integrations

TorchAO is integrated into some of the leading open-source libraries including:

🔎 Inference

TorchAO delivers substantial performance gains with minimal code changes:

Quantize any model with nn.Linear layers in just one line (Option 1), or load the quantized model directly from HuggingFace using our integration with HuggingFace transformers (Option 2):

Option 1: Direct TorchAO API

from torchao.quantization.quant_api import quantize_, Int4WeightOnlyConfig quantize_(model, Int4WeightOnlyConfig(group_size=128, use_hqq=True))

Option 2: HuggingFace Integration

from transformers import TorchAoConfig, AutoModelForCausalLM from torchao.quantization.quant_api import Int4WeightOnlyConfig

Create quantization configuration

quantization_config = TorchAoConfig(quant_type=Int4WeightOnlyConfig(group_size=128, use_hqq=True))

Load and automatically quantize

quantized_model = AutoModelForCausalLM.from_pretrained( "microsoft/Phi-4-mini-instruct", torch_dtype="auto", device_map="auto", quantization_config=quantization_config )

Deploy quantized models in vLLM with one command:

vllm serve pytorch/Phi-4-mini-instruct-int4wo-hqq --tokenizer microsoft/Phi-4-mini-instruct -O3

With this quantization flow, we achieve 67% VRAM reduction and 12-20% speedup on A100 GPUs while maintaining model quality. For more detail, see this step-by-step quantization guide. We also release some pre-quantized models here.

🚅 Training

Quantization-Aware Training

Post-training quantization can result in a fast and compact model, but may also lead to accuracy degradation. We recommend exploring Quantization-Aware Training (QAT) to overcome this limitation, especially for lower bit-width dtypes such as int4. In collaboration with TorchTune, we've developed a QAT recipe that demonstrates significant accuracy improvements over traditional PTQ, recovering 96% of the accuracy degradation on hellaswag and 68% of the perplexity degradation on wikitext for Llama3 compared to post-training quantization (PTQ). For more details, please refer to the QAT README and the original blog:

from torchao.quantization import quantize_ from torchao.quantization.qat import FakeQuantizeConfig, IntXQuantizationAwareTrainingConfig activation_config = FakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False) weight_config = FakeQuantizeConfig(torch.int4, group_size=32) qat_config = IntXQuantizationAwareTrainingConfig(activation_config, weight_config), quantize_(my_model, qat_config)

Users can also combine LoRA + QAT to speed up training by 1.89x compared to vanilla QAT using this fine-tuning recipe.

Float8

torchao.float8 implements training recipes with the scaled float8 dtypes, as laid out in https://arxiv.org/abs/2209.05433. With torch.compile on, current results show throughput speedups of up to 1.5x on up to 512 GPU / 405B parameter count scale (details):

from torchao.float8 import convert_to_float8_training convert_to_float8_training(m)

Our float8 training is integrated into TorchTitan's pre-training flows so users can easily try it out. For more details, check out these blog posts about our float8 training support:

Sparse Training

We've added support for semi-structured 2:4 sparsity with 6% end-to-end speedups on ViT-L. Full blog here. The code change is a 1 liner with the full example available here:

from torchao.sparsity.training import SemiSparseLinear, swap_linear_with_semi_sparse_linear swap_linear_with_semi_sparse_linear(model, {"seq.0": SemiSparseLinear})

Memory-efficient optimizers

Optimizers like ADAM can consume substantial GPU memory - 2x as much as the model parameters themselves. TorchAO provides two approaches to reduce this overhead:

1. Quantized optimizers: Reduce optimizer state memory by 2-4x by quantizing to lower precision

from torchao.optim import AdamW8bit, AdamW4bit, AdamWFp8 optim = AdamW8bit(model.parameters()) # replace with Adam4bit and AdamFp8 for the 4 / fp8 versions

Our quantized optimizers are implemented in just a few hundred lines of PyTorch code and compiled for efficiency. While slightly slower than specialized kernels, they offer an excellent balance of memory savings and performance. See detailed benchmarks here.

2. CPU offloading: Move optimizer state and gradients to CPU memory

For maximum memory savings, we support single GPU CPU offloading that efficiently moves both gradients and optimizer state to CPU memory. This approach can reduce your VRAM requirements by 60% with minimal impact on training speed:

optim = CPUOffloadOptimizer(model.parameters(), torch.optim.AdamW, fused=True) optim.load_state_dict(ckpt["optim"])

🎥 Videos

💬 Citation

If you find the torchao library useful, please cite it in your work as below.

@software{torchao, title={TorchAO: PyTorch-Native Training-to-Serving Model Optimization}, author={torchao}, url={https://github.com/pytorch/torchao}, license={BSD-3-Clause}, month={oct}, year={2024} }