arXiv:2507.16099cs.LG2025-07被引 30

TorchAO统一优化训练到部署流程,支持多种量化与稀疏技术。

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

  • 基于PyTorch原生设计,通过张量子类抽象支持多种低精度数据类型
  • 实现FP8训练、量化感知训练等技术,支持Llama 3.2 1B/3B等模型量化发布
  • 打通从训练到服务全链路,适配HuggingFace、vLLM等主流框架

我们提出TorchAO,一个基于PyTorch的模型优化框架,通过量化与稀疏化技术,实现从训练到部署的端到端优化流程。TorchAO支持多种主流优化方法,包括FP8量化训练、量化感知训练(QAT)、训练后量化(PTQ)和2:4稀疏性,并利用创新的张量子类抽象,表示多种通用、后端无关的低精度数据类型,如INT4、INT8、FP8、MXFP4、MXFP6和MXFP8。该框架在模型优化全流程中深度集成,涵盖预训练(TorchTitan)、微调(TorchTune、Axolotl)到部署(HuggingFace、vLLM、SGLang、ExecuTorch),将原本碎片化的生态整合为统一工作流。TorchAO已用于发布量化版Llama 3.2 1B/3B及LlamaGuard3-8B模型,项目开源地址为https://github.com/pytorch/ao/。

原文摘要 · Abstract (English)

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/.

模型优化量化PyTorch推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。