TorchAO: PyTorch-Native Training-to-Serving Model Optimization

2025/07/21 by Andrew Or, Or, Andrew, Apurva Jain +21 · 5 citations
Computer Science · #Advanced Neural Network Applications #FOS: Computer and information sciences #Human Pose and Action Recognition #Machine Learning (cs.LG) #Machine Learning and Data Classification

paper · pdf · doi:10.48550/arxiv.2507.16099

openalex publication_date 2025/07/21 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

We present TorchAO, a PyTorch-native model optimization framework leveraging quantization and sparsity to provide an end-to-end, training-to-serving workflow for AI models. TorchAO supports a variety of popular model optimization techniques, including FP8 quantized training, quantization-aware training (QAT), post-training quantization (PTQ), and 2:4 sparsity, and leverages a novel tensor subclass abstraction to represent a variety of widely-used, backend agnostic low precision data types, including INT4, INT8, FP8, MXFP4, MXFP6, and MXFP8. TorchAO integrates closely with the broader ecosystem at each step of the model optimization pipeline, from pre-training (TorchTitan) to fine-tuning (TorchTune, Axolotl) to serving (HuggingFace, vLLM, SGLang, ExecuTorch), connecting an otherwise fragmented space in a single, unified workflow. TorchAO has enabled recent launches of the quantized Llama 3.2 1B/3B and LlamaGuard3-8B models and is open-source at https://github.com/pytorch/ao/.

Citations

• Accelerating Transformer Inference and Training with 2:4 Activation Sparsity
• ParetoQ: Improving Scaling Laws in Extremely Low-bit LLM Quantization
• TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training
• The Llama 3 Herd of Models
• SpinQuant: LLM quantization with learned rotations
• GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection
• SGLang: Efficient Execution of Structured Language Model Programs
• Llama Guard: LLM-based Input-Output Safeguard for Human-AI Conversations
• Microscaling Data Formats for Deep Learning
• Efficient Memory Management for Large Language Model Serving with PagedAttention
• AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration
• Direct Preference Optimization: Your Language Model is Secretly a Reward Model
• QLoRA: Efficient Finetuning of Quantized LLMs
• SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models
• LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale
• HuggingFace's Transformers: State-of-the-art Natural Language Processing
• Stochastic Sign Descent Methods: New Algorithms and Better Theory

Cited by

• Scalable Neural Decoders for Practical Fault-Tolerant Quantum Computation
• UniQL: Unified Quantization and Low-rank Compression for Adaptive Edge LLMs
• The Unseen Frontier: Pushing the Limits of LLM Sparsity with Surrogate-Free ADMM
• Scaling LLM Test-Time Compute with Mobile NPU on Smartphones
• AttestLLM: Efficient Attestation Framework for Billion-scale On-device LLMs

Related

• TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training
• Torchbearer: A Model Fitting Library for PyTorch
• TorchBench: Benchmarking PyTorch with High API Surface Coverage
• PyLO: Towards Accessible Learned Optimizers in PyTorch
• COAT: Compressing Optimizer states and Activation for Memory-Efficient FP8 Training