TorchAO: 基于 PyTorch 的训练到服务模型优化-计算机科学-大语言模型-原生模型优化.pdf

TorchAO:基于PyTorch的训练到服务模型优化

11111

AndrewOrApurvaJainDanielVega-MyhreJesseCaiCharlesDavidHernandez

11111

ZhenruiZhangDrissGuessousVasiliyKuznetsovChristianPuhrschMarkSaroufim

123

SupriyaRaoThienTranAleksandarSamardžić

Abstract1.介绍

大型语言模型（LLMs）一直处于内容创作、文本

我们介绍了TorchAO，一个利用量化和稀

摘要、聊天机器人和代码生成等多种用途的前沿。

疏性提供从训练到服务的端到端工作流程

然而，这些能力通常需要大量的基础设施支持，

的PyTorch原生模型优化框架。TorchAO

本支持多种流行的模型优化技术，包括FP8如顶级模型Qwen3（235B参数）(Qwen3,2025)、

译量化训练、量化感知训练（QAT）、后训练DeepSeek-v3（671B）(Guoetal.,2025)、Llama3.1

中量化（PTQ）以及2:4稀疏性，并利用一（405B）(Grattafiorietal.,2024)和Llama4Behe-

1种新颖的张量子类抽象表示广泛使用的与moth（2T）(MetaAI,2025a)所示。

v后端无关的低精度数据类型，包括INT4、这些模型的计算成本和内存占用在大规模语言模型

9

9INT8、FP8、MXFP4、MXFP6和MXFP8。管道的每个步骤中，从训练到微调再到服务，都带

0

6来了显著的挑战。例如，训练Llama3.1需要16K

1TorchAO在模型优化流水线的每个步骤中

.H100GPU上的3084万GPU小时(Grattafiorietal.,

7都与更广泛的生态系统紧密集成，从预训

02024)，甚至以原始BF16精度提供该模型也需要至

5练（TorchTitan(Liangetal.,2024)）到少800GB的聚合内存来容纳模型，超过了配备8个

2微调（TorchTune(torchtune,2025),Ax-

:H100GPU的单一服务器的内存限制。即使在较小

volotl(Axolotl,2025)）再到服务（Hugging-

i的1-8B参数规模下，减小这些模型的大小对