pith. sign in

arxiv: 2507.16099 · v1 · pith:ZB3JCTAXnew · submitted 2025-07-21 · 💻 cs.LG

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

Andrew Or , Apurva Jain , Daniel Vega-Myhre , Jesse Cai , Charles David Hernandez , Zhenrui Zheng , Driss Guessous , Vasiliy Kuznetsov
show 5 more authors
Christian Puhrsch Mark Saroufim Supriya Rao Thien Tran Aleksandar Samard\v{z}i\'c
This is my paper
classification 💻 cs.LG
keywords torchaomodeloptimizationincludingmodelspytorch-nativequantizationquantized
0
0 comments X
read the original 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/.

This paper has not been read by Pith yet.

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Forward citations

Cited by 6 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. Zero-Shot Quantization via Weight-Space Arithmetic

    cs.CV 2026-04 unverdicted novelty 8.0

    A quantization vector derived from a donor model via weight-space arithmetic can be added to a receiver model to improve post-PTQ Top-1 accuracy by up to 60 points in 3-bit settings without receiver-side QAT or data.

  2. LoKA: Low-precision Kernel Applications for Recommendation Models At Scale

    cs.LG 2026-05 unverdicted novelty 6.0

    LoKA enables practical FP8 use in numerically sensitive large recommendation models via profiling, model adaptations, and runtime kernel orchestration.

  3. LoKA: Low-precision Kernel Applications for Recommendation Models At Scale

    cs.LG 2026-05 unverdicted novelty 6.0

    LoKA enables practical FP8 use in numerically sensitive large recommendation models via online profiling of activations, reusable model modifications for stability, and dynamic kernel dispatching.

  4. StoSignSGD: Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models

    cs.LG 2026-04 unverdicted novelty 6.0

    StoSignSGD resolves SignSGD divergence on non-smooth objectives via structural stochasticity, matching optimal convex rates and improving non-convex bounds while delivering 1.44-2.14x speedups in FP8 LLM pretraining.

  5. torchtune: PyTorch native post-training library

    cs.LG 2026-05 unverdicted novelty 5.0

    torchtune is a modular PyTorch library for LLM post-training that delivers competitive performance and memory efficiency while supporting rapid research iteration through hackable components.

  6. CAT-Translate: Building Compact Open-Source Models for Japanese-English Translation

    cs.CL 2026-06 unverdicted novelty 3.0

    Compact 0.8B-7B models for bidirectional Japanese-English translation outperform large multilingual models on real-world domain benchmarks.