Favicon of TorchAO

TorchAO

A PyTorch quantization library that reduces model memory use for training and inference, with vLLM serving and ExecuTorch on-device deployment.

Screenshot of TorchAO website

TorchAO is a PyTorch library for developers who want to train and run models on their own hardware with less memory. It supports language models and image or video generation, and works with torch.compile and FSDP2 across most Hugging Face PyTorch models.

  • Quantized inference: Reduce precision with int4 weights, float8 quantization, or combined int8 activations and int4 weights, with workflows for Gemma, Qwen3, FLUX.1-dev and CogVideoX.
  • Accuracy recovery: Quantization-aware training helps reduce accuracy loss from lower precision and can work alongside LoRA fine-tuning through integrations such as TorchTune, Axolotl and Unsloth.
  • Float8 training: Train models at lower precision using recipes integrated into TorchTitan's pre-training workflows.
  • Training memory: Quantized optimizers reduce the memory used by optimizer state; single-GPU CPU offloading moves that state and gradients into CPU memory.
  • Serving and device deployment: Use vLLM or SGLang for self-hosted LLM serving, or ExecuTorch for edge deployment, including quantized Qwen3 on iPhone.

Hugging Face Transformers includes a TorchAO inference backend, while Diffusers supports its use with diffusion models. Hardware support depends on the workflow; inference examples include NVIDIA H100 GPUs and iPhone 15 Pro. MSLK is an optional runtime dependency that supplies accelerated kernels for some workflows. Post-training quantization can reduce accuracy, particularly at lower precision such as int4.

Similar to TorchAO