torchao

// Overview

torchao is a quantization library built for native PyTorch. It lets you shrink model weights to low-bit data types such as int4 and float8, which reduces memory use and speeds up both training and inference. It works out of the box with torch.compile() and FSDP2 across most HuggingFace PyTorch models.

It is aimed at ML engineers and researchers who train or serve large models and want to lower hardware costs without rewriting their model code. Applying quantization usually takes only a few extra lines, using the quantize_ function with a config object.

Within the inference and serving space, torchao covers the quantization and compression step. It is integrated as a quantization backend in tools like vLLM, SGLang, and HuggingFace Transformers, and also supports quantization-aware training (QAT) to recover accuracy lost during quantization.

// What it does

• Quantize model weights to int4 with the quantize_ API and an Int4WeightOnlyConfig, for faster inference and lower memory use
• float8 training support for large models such as Llama-3.1-70B
• Quantization-aware training (QAT) to recover accuracy lost when quantizing
• Works with torch.compile() and FSDP2 across most HuggingFace PyTorch models
• Used as a quantization backend in vLLM, SGLang, and HuggingFace Transformers
• Supports CUDA, XPU, and CPU builds, plus low-bit ARM CPU kernels for linear and embedding ops

// Getting started

Install torchao from pip, then quantize an existing PyTorch model's weights in place.

pip install torchao

import torch
from torchao.quantization import Int4WeightOnlyConfig, quantize_
if torch.cuda.is_available():
  # quantize on CUDA
  quantize_(model, Int4WeightOnlyConfig(group_size=32, int4_packing_format="tile_packed_to_4d", int4_choose_qparams_algorithm="hqq"))
elif torch.xpu.is_available():
  # quantize on XPU
  quantize_(model, Int4WeightOnlyConfig(group_size=32, int4_packing_format="plain_int32"))

Commands and code are distilled from the project's own documentation — always check the official repo for the latest.

// When to use it

• Reduce the memory footprint and speed up inference of a large language model by quantizing its weights to int4
• Pre-train or fine-tune large models faster using float8 training
• Recover accuracy lost to quantization with quantization-aware training before serving
• Serve quantized models through vLLM, SGLang, or HuggingFace Transformers using torchao as the quantization backend

// How torchao compares

torchao alongside other open-source quantization & compression tools AI/TLDR tracks, ranked by GitHub stars.