TorchAO
PyTorch-Native Training-to-Serving Model Optimization
Pre-train Llama-3.1-70B 1.5x faster with float8 training
Recover 67% of quantized accuracy degradation on Gemma3-4B with QAT
Quantize Llama-3-8B to int4 for 1.89x faster inference with 58% less memory
📣 Latest News
[Oct 25] QAT is now integrated into
for both full and LoRA fine-tuning! Try it out using
.
[Oct 25] MXFP8 MoE training prototype achieved ~1.45x speedup for MoE layer in Llama4 Scout, and ~1.25x speedup for MoE layer in DeepSeekV3 671b - with comparable numerics to bfloat16! Check out the
to try it out.
[Sept 25] MXFP8 training achieved
1.28x speedup on Crusoe B200 cluster
with virtually identical loss curve to bfloat16!
[Sept 19]
TorchAO Quantized Model and Quantization Recipes Now Available on Huggingface Hub
!
[Jun 25] Our
was accepted to CodeML @ ICML 2025!
Older news[May 25] QAT is now integrated into
for fine-tuning (
)!
[Apr 25] Float8 rowwise training yielded
at 2k H100 GPU scale
[Apr 25] TorchAO is added as a
(
)!
[Mar 25] Our
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
for linear and embedding ops
[Nov 24] We achieved
1.43-1.51x faster pre-training
on Llama-3.1-70B and 405B using float8 training
[Oct 24] TorchAO is added as a quantization backend to HF Transformers!
[Sep 24] We officially launched TorchAO. Check out our blog
!
[Jul 24] QAT
recovered up to 96% accuracy degradation
from quantization on Llama-3-8B
[Jun 24] Semi-structured 2:4 sparsity
achieved 1.1x inference speedup and 1.3x training speedup
on the SAM and ViT models respectively
[Jun 24] Block sparsity
achieved 1.46x training speeedup
on the ViT model with <2% drop in accuracy
🌅 Overview
TorchAO is an easy to use quantization library for native PyTorch. TorchAO works out-of-the-box with torch.compile() and FSDP2 across most HuggingFace PyTorch models.
For a detailed overview of stable and prototype workflows for different hardware and dtypes, see the
.
Check out our
for more details!
🚀 Quick Start
First, install TorchAO. We recommend installing the latest stable version:
pip install torchao # optional - install MSLK for float8 and nvfp4 inference kernels pip install mslk --index-url https://download.pytorch.org/whl/cu130 # optional - install apache-tvm-ffi and cutedsl for mxfp8 MoE training kernels pip install apache-tvm-ffi pip install nvidia-cutlass-dsl==4.5.2 nvidia-cutlass-dsl-libs-base==4.5.2 nvidia-cutlass-dsl-libs-cu13==4.5.2Quantize your model weights to int4!
importtorchfromtorchao.quantizationimportInt4WeightOnlyConfig, quantize_quantize_(model, Int4WeightOnlyConfig(group_size=32, int4_packing_format="tile_packed_to_4d", int4_choose_qparams_algorithm="hqq"))See our
for more details.
🛠 Installation
To install the latest stable version:
pip install torchaoOther installation options# Nightly pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu128 # Different CUDA versions pip install torchao --index-url https://download.pytorch.org/whl/cu126 # CUDA 12.6 pip install torchao --index-url https://download.pytorch.org/whl/cu129 # CUDA 12.9 pip install torchao --index-url https://download.pytorch.org/whl/xpu # XPU pip install torchao --index-url https://download.pytorch.org/whl/cpu # CPU only # For developers # Note: the `--no-build-isolation` flag is required. USE_CUDA=1 pip install -e . --no-build-isolation USE_XPU=1 pip install -e . --no-build-isolation USE_CPP=0 pip install -e . --no-build-isolation Please see the
for version requirements for dependencies.
Optional Dependencies
is an optional runtime dependency that provides accelerated kernels for some of the workflows in torchao. Stable MSLK should be used with stable torchao, and nightly MSLK with nightly torchao.
# Stable pip install mslk --index-url https://download.pytorch.org/whl/cu130 # Nightly pip install --pre mslk --index-url https://download.pytorch.org/whl/nightly/cu130apache-tvm-ffi and nvidia-cutlass-dsl are used for MoE mxfp8 training kernels.
pip install apache-tvm-ffi pip install nvidia-cutlass-dsl==4.5.2 nvidia-cutlass-dsl-libs-base==4.5.2 nvidia-cutlass-dsl-libs-cu13==4.5.2TorchAO delivers substantial performance gains with minimal code changes:
Int4 weight-only:
1.73x speedup with 65% less memory
for Gemma3-12b-it on H100 with slight impact on accuracy
Float8 dynamic quantization:
1.5-1.6x speedup on gemma-3-27b-it
and
1.54x and 1.27x speedup on Flux.1-Dev* and CogVideoX-5b respectively
on H100 with preserved quality
Int8 activation quantization and int4 weight quantization: Quantized Qwen3-4B running with 14.8 tokens/s with 3379 MB memory usage on iPhone 15 Pro through
Following is our recommended flow for quantization and deployment:
fromtransformersimportTorchAoConfig, AutoModelForCausalLMfromtorchao.quantizationimportFloat8DynamicActivationFloat8WeightConfig, PerRow# Create quantization configurationquantization_config=TorchAoConfig(quant_type=Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())) # Load and automatically quantizequantized_model=AutoModelForCausalLM.from_pretrained( "Qwen/Qwen3-32B", dtype="auto", device_map="auto", quantization_config=quantization_config )Alternative quantization API to use when the above doesn't work is quantize_ API in
.
Serving with vllm on 1xH100 machine:
# Server VLLM_DISABLE_COMPILE_CACHE=1 vllm serve pytorch/Qwen3-32B-FP8 --tokenizer Qwen/Qwen3-32B -O3# Client curl http://localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{ "model": "pytorch/Qwen3-32B-FP8", "messages": [ {"role": "user", "content": "Give me a short introduction to large language models."} ], "temperature": 0.6, "top_p": 0.95, "top_k": 20, "max_tokens": 32768}'For diffusion models, you can quantize using Hugging Face diffusers
importtorchfromdiffusersimportDiffusionPipeline, PipelineQuantizationConfig, TorchAoConfigfromtorchao.quantizationimportInt8WeightOnlyConfigfromtorchao.quantization.granularityimportPerGrouppipeline_quant_config=PipelineQuantizationConfig( quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(granularity=PerGroup(128)))} ) pipeline=DiffusionPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16, device_map="cuda" )We also support deployment to edge devices through ExecuTorch, for more detail, see
quantization and serving guide
. We also release pre-quantized models
.
🚅 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
, 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
and the
:
importtorchfromtorchao.quantizationimportquantize_, Int8DynamicActivationIntxWeightConfig, PerGroupfromtorchao.quantization.qatimportQATConfig# preparebase_config=Int8DynamicActivationIntxWeightConfig( weight_dtype=torch.int4, weight_granularity=PerGroup(32), ) quantize_(my_model, QATConfig(base_config, step="prepare")) # train model (not shown)# convertquantize_(my_model, QATConfig(base_config, step="convert"))Users can also combine LoRA + QAT to speed up training by
compared to vanilla QAT using this
.
Quantized training
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 (
):
fromtorchao.float8importconvert_to_float8_trainingconvert_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
Other features (memory efficient optimizers)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
fromtorchao.optimimportAdamW8bit, AdamW4bit, AdamWFp8optim=AdamW8bit(model.parameters()) # replace with Adam4bit and AdamFp8 for the 4 / fp8 versionsOur 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
.
2. CPU offloading: Move optimizer state and gradients to CPU memory
For maximum memory savings, we support
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:
optim=CPUOffloadOptimizer(model.parameters(), torch.optim.AdamW, fused=True) optim.load_state_dict(ckpt["optim"])🔗 Integrations
TorchAO is integrated into some of the leading open-source libraries including:
Unsloth now supports QAT:
and
.
HuggingFace transformers with a
and
HuggingFace
best practices with torch.compile and TorchAO in a standalone repo
vLLM for LLM serving:
,
Integration with
for SOTA kernels on server GPUs
Integration with
for edge device deployment
Axolotl for
and
TorchTitan for
HuggingFace PEFT for LoRA using TorchAO as their
TorchTune for our NF4
,
, and
recipes
SGLang for LLM serving:
🎥 Videos
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
💬 Citation
If you find the torchao library useful, please cite it in your work as below.
@misc{or2025torchao, title={TorchAO: PyTorch-Native Training-to-Serving Model Optimization}, author={Andrew Or and Apurva Jain and Daniel Vega-Myhre and Jesse Cai and Charles David Hernandez and Zhenrui Zheng and Driss Guessous and Vasiliy Kuznetsov and Christian Puhrsch and Mark Saroufim and Supriya Rao and Thien Tran and Aleksandar Samardžić}, year={2025}, eprint={2507.16099}, archivePrefix={arXiv}, primaryClass={cs.LG}, url={https://arxiv.org/abs/2507.16099}, }