PyTorch-Native Training-to-Serving Model Optimization

• Pre-train Llama-3.1-70B 1.5x faster with float8 training
• Recover 77% of quantized perplexity degradation on Llama-3.2-3B with QAT
• Quantize Llama-3-8B to int4 for 1.89x faster inference with 58% less memory

Latest News

• [Jun 25] Our TorchAO paper was accepted to CodeML @ ICML 2025!
• [May 25] QAT is now integrated into Axolotl for fine-tuning (docs)!
• [Apr 25] Float8 rowwise training yielded 1.34-1.43x training speedup at 2k H100 GPU scale
• [Apr 25] TorchAO is added as a quantization backend to vLLM (docs)!
• [Mar 25] Our 2:4 Sparsity paper was accepted to SLLM @ ICLR 2025!
• [Jan 25] Our integration with GemLite and SGLang yielded 1.1-2x faster inference with int4 and float8 quantization across different batch sizes and tensor parallel sizes
• [Jan 25] We added 1-8 bit ARM CPU kernels for linear and embedding ops

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: * Float8 training and inference for speedups without compromising accuracy * MX training and inference, provides MX tensor formats based on native PyTorch MX dtypes (prototype) * Quantization-Aware Training (QAT) for mitigating quantization degradation * Post-Training Quantization (PTQ) for int4, int8, fp6 etc, with matching kernels targeting a variety of backends including CUDA, ARM CPU, and XNNPACK * Sparsity, includes different techniques such as 2:4 sparsity and block sparsity

Check out our docs for more details!

From the team that brought you the fast series: * 9.5x inference speedups for Image segmentation models with sam-fast * 10x inference speedups for Language models with gpt-fast * 3x inference speedup for Diffusion models with sd-fast (new: flux-fast) * 2.7x inference speedup for FAIRs Seamless M4T-v2 model with seamlessv2-fast

Quick Start

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

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!

Installation

To install the latest stable version: pip install torchao

Please see the torchao compability table for version requirements for dependencies.

Integrations

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

• HuggingFace transformers with a builtin inference backend and low bit optimizers
• HuggingFace diffusers best practices with torch.compile and TorchAO in a standalone repo diffusers-torchao
• HuggingFace PEFT for LoRA using TorchAO as their quantization backend
• Mobius HQQ backend leveraged our int4 kernels to get 195 tok/s on a 4090
• TorchTune for our NF4 QLoRA, QAT, and float8 quantized fine-tuning recipes
• TorchTitan for float8 pre-training
• VLLM for LLM serving: usage, detailed docs
• SGLang for LLM serving: usage and the major PR.
• Axolotl for QAT and PTQ

Inference

TorchAO delivers substantial performance gains with minimal code changes:

• Int4 weight-only: 1.89x throughput with 58.1% less memory on Llama-3-8B
• Float8 dynamic quantization: 1.54x and 1.27x speedup on Flux.1-Dev* and CogVideoX-5b respectively on H100 with preserved quality
• Int4 + 2:4 Sparsity: 2.37x throughput with 67.7% memory reduction on Llama-3-8B

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

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

Option 2: HuggingFace Integration

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

Create quantization configuration

quantizationconfig = TorchAoConfig(quanttype=Int4WeightOnlyConfig(groupsize=128, usehqq=True))

Load and automatically quantize

quantizedmodel = AutoModelForCausalLM.frompretrained( "microsoft/Phi-4-mini-instruct", torchdtype="auto", devicemap="auto", quantizationconfig=quantizationconfig ) ```

Deploy quantized models in vLLM with one command:

shell 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:

```python from torchao.quantization import quantize_, Int8DynamicActivationInt4WeightConfig from torchao.quantization.qat import QATConfig

prepare

baseconfig = Int8DynamicActivationInt4WeightConfig(groupsize=32) quantize(mymodel, QATConfig(base_config, step="prepare"))

train model (not shown)

convert

quantize(mymodel, QATConfig(base_config, step="convert")) ```

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

python 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: * Accelerating Large Scale Training and Convergence with PyTorch Float8 Rowwise on Crusoe 2K H200s * Supercharging Training using float8 and FSDP2 * Efficient Pre-training of Llama 3-like model architectures using torchtitan on Amazon SageMaker * Float8 in PyTorch

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:

python 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

python 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:

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

Videos

• Keynote talk at GPU MODE IRL
• Low precision dtypes at PyTorch conference
• Slaying OOMs at the Mastering LLM's course
• Advanced Quantization at CUDA MODE
• Chip Huyen's GPU Optimization Workshop
• Cohere for AI community talk

Citation

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

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