Repository files navigation

• Pre-train Llama-3.1-70B1.5x faster with float8 training
• Recover67% of quantized accuracy degradation on Gemma3-4B with QAT
• Quantize Llama-3-8B to int4 for1.89x faster inference with58% less memory

• [Oct 25] QAT is now integrated intoUnsloth for both full and LoRA fine-tuning! Try it out usingthis notebook.
• [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 thedocs to try it out.
• [Sept 25] MXFP8 training achieved1.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] OurTorchAO paper was accepted to CodeML @ ICML 2025!

TorchAO is an easy to use quantization library for native PyTorch. TorchAO works out-of-the-box withtorch.compile() andFSDP2 across most HuggingFace PyTorch models.

🟢 = stable, 🟡 = prototype, 🟠 = planned, ⚪ = not supported

🟢 = stable, 🟡 = prototype, 🟠 = planned, ⚪ = not supported

• Quantization-Aware Training (QAT) README.md
• Post-Training Quantization (PTQ) README.md
• Sparsity README.md, includes different techniques such as 2:4 sparsity and block sparsity
• the prototype folder for other prototype features

Check out ourdocs for more details!

First, install TorchAO. We recommend installing the latest stable version:

pip install torchao

Quantize your model weights to int4!

fromtorchao.quantizationimportInt4WeightOnlyConfig,quantize_quantize_(model,Int4WeightOnlyConfig(group_size=32,int4_packing_format="tile_packed_to_4d",int4_choose_qparams_algorithm="hqq"))

See ourquick start guide for more details.

To install the latest stable version:

pip install torchao

# Nightlypip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu128# Different CUDA versionspip install torchao --index-url https://download.pytorch.org/whl/cu126  # CUDA 12.6pip install torchao --index-url https://download.pytorch.org/whl/cu129  # CUDA 12.9pip 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-isolationUSE_CPP=0 pip install -e . --no-build-isolation

Please see thetorchao compability table for version requirements for dependencies.

TorchAO 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 and1.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 throughExecuTorch
• Int4 + 2:4 Sparsity:2.37x throughput with 67.7% memory reduction on Llama-3-8B

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 isquantize_ API inquick start guide.

Serving with vllm on 1xH100 machine:

# ServerVLLM_DISABLE_COMPILE_CACHE=1 vllm serve pytorch/Qwen3-32B-FP8 --tokenizer Qwen/Qwen3-32B -O3

# Clientcurl 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}'

We also support deployment to edge devices through ExecuTorch, for more detail, seequantization and serving guide. We also release pre-quantized modelshere.

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 withTorchTune, we've developed a QAT recipe that demonstrates significant accuracy improvements over traditional PTQ, recovering96% 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 theQAT README and theoriginal blog:

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 by1.89x compared to vanilla QAT using thisfine-tuning recipe.

torchao.float8 implements training recipes with the scaled float8 dtypes, as laid out inhttps://arxiv.org/abs/2209.05433. Withtorch.compile on, current results show throughput speedups of up to1.5x on up to 512 GPU / 405B parameter count scale (details):

fromtorchao.float8importconvert_to_float8_trainingconvert_to_float8_training(m)

Our float8 training is integrated intoTorchTitan'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
• Float8 in PyTorch

fromtorchao.sparsity.trainingimportSemiSparseLinear,swap_linear_with_semi_sparse_linearswap_linear_with_semi_sparse_linear(model, {"seq.0":SemiSparseLinear})

fromtorchao.optimimportAdamW8bit,AdamW4bit,AdamWFp8optim=AdamW8bit(model.parameters())# replace with Adam4bit and AdamFp8 for the 4 / fp8 versions

optim=CPUOffloadOptimizer(model.parameters(),torch.optim.AdamW,fused=True)optim.load_state_dict(ckpt["optim"])

TorchAO is integrated into some of the leading open-source libraries including:

• Unsloth now supports QAT:Read blog andguide.
• HuggingFace transformers with abuiltin inference backend andlow bit optimizers
• HuggingFace diffusers best practices withtorch.compile and TorchAO in a standalone repodiffusers-torchao
• vLLM for LLM serving:usage,detailed docs
• Integration withFBGEMM for SOTA kernels on server GPUs
• Integration withExecuTorch for edge device deployment
• Axolotl forQAT andPTQ
• TorchTitan forfloat8 pre-training
• HuggingFace PEFT for LoRA using TorchAO as theirquantization backend
• TorchTune for our NF4QLoRA,QAT, andfloat8 quantized fine-tuning recipes
• SGLang for LLM serving:usage

• Keynote talk at GPU MODE IRL
• 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
• Cohere for AI community talk

If you find the torchao library useful, please cite it in your work as below.

@software{torchao,title={TorchAO: PyTorch-Native Training-to-Serving Model Optimization},author={torchao},url={https://github.com/pytorch/ao},license={BSD-3-Clause},month={oct},year={2024}}