TorchAO: PyTorch-Native Quantization for Training and Inference

What TorchAO solves, and who it is for

Quantization in PyTorch has historically meant stitching together third-party kernels, custom autograd functions and framework-specific conversion scripts. TorchAO's premise is that the dtype, the quantization config and the serving path should all live in the same library as the model. The README describes it as "an easy to use quantization library for native PyTorch" that works out-of-the-box with torch.compile() and FSDP2 across most HuggingFace PyTorch models.

The audience is narrow but deep: engineers training or serving transformer models who are willing to trade portability for speed. The README's headline claims are all transformer-scale: pre-training Llama-3.1-70B with float8, quantizing Llama-3-8B to int4, and MXFP8 MoE training on Llama4 Scout and DeepSeekV3 671b. If you are quantizing a small CNN for a microcontroller, nothing here is aimed at you.

How the quantization workflows are organized

The library separates stable workflows from prototype ones, and the README points to a Workflows documentation page that maps hardware and dtypes to the supported path. That structure matters more than any single API. The same repository ships int4 weight-only inference, float8 training, MXFP8 training, quantization-aware training and semi-structured sparsity, and each has different kernel dependencies and different hardware requirements.

The core entry point shown in the README is quantize_, which mutates a model in place given a config object. The example config, Int4WeightOnlyConfig, exposes group_size, int4_packing_format and int4_choose_qparams_algorithm as separate knobs. That is a deliberate design choice: packing format and qparam algorithm are not hidden behind a single "int4" switch, because they interact with the kernel that will eventually run the matmul. The cost is that a user can produce a quantized model that is numerically valid but slow, if the packing format does not match what the target kernel expects.

Installing TorchAO and quantizing a model to int4

The stable release installs from PyPI. The README recommends the latest stable version, and notes that MSLK is an optional runtime dependency for accelerated kernels in some workflows.

pip install torchao
# optional - install MSLK for float8 and nvfp4 inference kernels
pip install mslk --index-url https://download.pytorch.org/whl/cu130

The README also documents nightly builds and per-CUDA index URLs, including cu126, cu129, XPU and CPU-only variants. After installing, the quick start quantizes model weights to int4 with a single call. The config below is copied from the README, including the tile_packed_to_4d packing format and the hqq qparam algorithm.

import torch
from torchao.quantization import Int4WeightOnlyConfig, quantize_
quantize_(model, Int4WeightOnlyConfig(group_size=32, int4_packing_format="tile_packed_to_4d", int4_choose_qparams_algorithm="hqq"))

quantize_ modifies the model in place, so the object you pass in is the object you run afterward. The example file examples/quantize_llama_4.py in the repository shows the same pattern applied to a Llama-family model rather than a toy module. If you are building from source instead, setup.py requires the --no-build-isolation flag, and pyproject.toml states that torch must already be installed in the environment because build-time dependencies are not declared.

Where TorchAO stops being the right tool

The README is explicit that some workflows are prototypes. The MXFP8 MoE training path lives under torchao/prototype/moe_training/ and needs apache-tvm-ffi plus pinned nvidia-cutlass-dsl packages. A prototype directory is a signal about API stability, not a guarantee of breakage, but it does mean the interface can move between releases.

Version alignment is the second constraint. The README directs readers to a compatibility table issue for version requirements for dependencies, and states that stable MSLK should be used with stable torchao while nightly MSLK goes with nightly torchao. Mixing stable and nightly across those two packages is the failure mode the documentation warns about. There is also a hardware boundary: the per-CUDA index URLs and the MSLK wheel target specific CUDA builds, and the CPU-only install exists but the README's performance claims are all GPU-scale. If your target is a non-NVIDIA accelerator not covered by the XPU or CPU indexes, the documentation does not describe a path for you.

TorchAO against bitsandbytes and GPTQ style pipelines

The practical alternative most teams compare against is a post-training quantization toolchain that produces a serialized quantized checkpoint, such as bitsandbytes or GPTQ-style pipelines used through HF Transformers. The difference is where quantization sits in the lifecycle. Those tools generally take a trained model and emit quantized weights for inference. TorchAO spans pre-training, fine-tuning and serving, which is why the README leads with float8 training speedups and QAT accuracy recovery rather than only inference.

That breadth changes the integration story. Because TorchAO is native PyTorch, it composes with torch.compile() and FSDP2 rather than sitting in front of them. The trade-off is ecosystem lock-in in the other direction: bitsandbytes and GPTQ checkpoints are widely consumed by inference servers that have never heard of TorchAO. TorchAO's own reach comes through integrations, and the README lists vLLM as a quantization backend, HF Transformers, GemLite and SGLang, Unsloth and Axolotl for QAT. If your serving stack is not on that list, you are on your own for the runtime side.

Maintenance cadence, licence and upgrade cost

The repository is not archived, and the last push was on 2026-09-15. Releases are frequent: v0.18.0 on 2026-08-03, v0.17.0 on 2026-03-30 and v0.16.0 on 2026-02-10, roughly a quarterly cadence with a shorter gap before the most recent one. That cadence is the upgrade cost. A library that ships float8, MXFP8, int4 and sparsity paths across CUDA, XPU and CPU builds has a large surface, and the README's own compatibility table exists because dependency versions matter.

The licence situation needs care. The repository metadata reports NOASSERTION, while the README's badge links to BSD 3-Clause in ./LICENSE. The setup.py header carries a Meta Platforms copyright notice and points to the LICENSE file in the source tree. Because the machine-readable licence field and the README badge disagree, check ./LICENSE directly before you rely on either. This is a factual discrepancy in the metadata, not legal advice.

Editorial conclusion

Adopt TorchAO if you are already on PyTorch and want quantization that stays inside the framework, especially for float8 or MXFP8 training and int4 serving on supported GPUs. Do not adopt it if you need a single config that works across every accelerator, or if you cannot pin torch, CUDA and optional kernel packages to matching versions. Verify first that your target dtype and hardware combination appears in the Workflows documentation, that your PyTorch version satisfies the compatibility table, and that your deployment path (vLLM, HF Transformers, or plain torch.compile) has a documented TorchAO backend.

Frequently asked questions

Official sources

1. Issues
2. Project website
3. pytorch/ao on GitHub
4. README
5. Releases