TorchAO：面向训练至部署的PyTorch原生模型优化框架
机器学习
2025-07-23 v1
摘要
我们提出TorchAO，一款基于量化和稀疏化的PyTorch原生模型优化框架，为AI模型提供端到端的训练到部署工作流程。TorchAO支持多种流行的模型优化技术，包括FP8量化训练、量化感知训练（QAT）、后训练量化（PTQ）和2:4稀疏化，利用一种新颖的张量子类抽象来表示各种广泛使用的、与后端无关的低精度数据类型，包括INT4、INT8、FP8、MXFP4、MXFP6和MXFP8。TorchAO在模型优化管道的每个环节都与更广泛的生态系统紧密集成，涵盖从预训练（TorchTitan）到微调（TorchTune、Axolotl）到部署（HuggingFace、vLLM、SGLang、ExecuTorch），将原本碎片化的空间连接为统一的工作流程。TorchAO已使最近推出的量化Llama 3.2 1B/3B和LlamaGuard3-8B模型成为可能，该项目开源于https://github.com/pytorch/ao/。
引用
@article{arxiv.2507.16099,
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ć},
journal= {arXiv preprint arXiv:2507.16099},
year = {2025}
}
备注
5 pages, 3 figures, published in CODEML@ICML25