TorchAO

TorchAO is a PyTorch-native library for custom data types and optimizations, hosted under the pytorch GitHub organization. It quantizes and sparsifies weights, gradients, optimizers, and activations for both training and inference. The latest GitHub release listed is v0.18.0, published 2026-08-03.

What It Is

TorchAO is a Python library that adds quantization, quantization-aware training, float8 training, and sparsity to native PyTorch models. The README states it works out of the box with torch.compile() and FSDP2 across most Hugging Face PyTorch models. Users apply it by calling APIs such as quantize_ with a config object (for example an int4 weight-only config), or convert_to_float8_training for float8 training.

Workflows: Training to Serving

The documentation organizes the library into quantized training, quantization-aware training (QAT), and quantized inference. QAT uses a prepare step, a training step, and a convert step to reduce the accuracy loss from low-bit quantization. The project also provides memory-efficient optimizers (8-bit, 4-bit, and fp8 AdamW variants) and CPU offloading of optimizer state. The docs also cover PT2E quantization, static quantization, serialization, and writing custom quantized tensor subclasses.

Integrations and Deployment

The README lists integrations with Hugging Face Transformers, Diffusers, and PEFT, vLLM, SGLang, ExecuTorch for edge devices, TorchTitan, TorchTune, Axolotl, Unsloth, and MSLK kernels. Installation is via pip, with index options for nightly builds, CUDA 12.6 and 12.9, XPU, and CPU-only. The project reports benchmark figures such as 1.5x faster float8 pre-training of Llama-3.1-70B and 1.89x faster int4 inference for Llama-3-8B with 58% less memory.