PyTorch原生训练至推理模型优化

• 使用float8训练，预训练Llama-3.1-70B速度提升1.5倍
• 通过QAT在Gemma3-4B上恢复67%的量化精度损失
• 将Llama-3-8B量化为int4，推理速度提升1.89倍，内存占用减少58%

📣 最新消息

• [10月25日] QAT现已集成到Unsloth中，适用于全模型和LoRA微调！使用此笔记本即可体验。
• [10月25日] MXFP8 MoE训练原型在Llama4 Scout的MoE层实现了约1.45倍加速，在DeepSeekV3 671b的MoE层则实现了约1.25倍加速——且数值与bfloat16相当！请查看文档以尝试。
• [9月25日] MXFP8训练在Crusoe B200集群上实现了1.28倍加速（https://pytorch.org/blog/accelerating-2k-scale-pre-training-up-to-1-28x-with-torchao-mxfp8-and-torchtitan-on-crusoe-b200-cluster/），且损失曲线与bfloat16几乎完全一致！
• [9月19日] TorchAO量化模型及量化配方现已上线Huggingface Hub！
• [6月25日] 我们的TorchAO论文已被ICML 2025的CodeML接收！

🌅 概述

TorchAO是一个易于使用的原生PyTorch量化库。TorchAO可与torch.compile()和FSDP2无缝配合，适用于大多数HuggingFace PyTorch模型。

有关不同硬件和数据类型下稳定及原型工作流的详细概述，请参阅工作流文档。

更多详情请查看我们的文档！

🚀 快速入门

首先，安装TorchAO。我们建议安装最新的稳定版本：

pip install torchao

将您的模型权重量化为int4！

import torch
from torchao.quantization import Int4WeightOnlyConfig, quantize_
if torch.cuda.is_available():
  # 在CUDA上量化
  quantize_(model, Int4WeightOnlyConfig(group_size=32, int4_packing_format="tile_packed_to_4d", int4_choose_qparams_algorithm="hqq"))
elif torch.xpu.is_available():
  # 在XPU上量化
  quantize_(model, Int4WeightOnlyConfig(group_size=32, int4_packing_format="plain_int32"))

更多详情请参阅我们的快速入门指南。

🛠 安装

要安装最新稳定版本：

pip install torchao

# 夜间版
pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu128

# 不同CUDA版本
pip install torchao --index-url https://download.pytorch.org/whl/cu126  # CUDA 12.6
pip install torchao --index-url https://download.pytorch.org/whl/cu129  # CUDA 12.9
pip install torchao --index-url https://download.pytorch.org/whl/xpu    # XPU
pip install torchao --index-url https://download.pytorch.org/whl/cpu    # 仅CPU

# 针对开发者
# 注意：需要使用`--no-build-isolation`标志。
USE_CUDA=1 pip install -e . --no-build-isolation
USE_XPU=1 pip install -e . --no-build-isolation
USE_CPP=0 pip install -e . --no-build-isolation

请参阅TorchAO兼容性表以了解依赖项的版本要求。

可选依赖

MSLK是TorchAO部分工作流中提供加速内核的可选运行时依赖。稳定的MSLK应与稳定的TorchAO搭配使用，而夜间版MSLK则应与夜间版TorchAO搭配使用。

# 稳定版
pip install mslk-cuda==1.0.0

# 夜间版
pip install --pre mslk --index-url https://download.pytorch.org/whl/nightly/cu128

🔎 推理

TorchAO 仅需少量代码改动即可带来显著的性能提升：

• Int4 权重量化：在 H100 上，Gemma3-12b-it 的推理速度提升 1.73 倍，内存占用减少 65%，且对精度影响较小。
• Float8 动态量化：在 H100 上，Gemma3-27b-it 的推理速度提升 1.5–1.6 倍；Flux.1-Dev* 和 CogVideoX-5b 的推理速度分别提升 1.54 倍和 1.27 倍，同时保持模型质量不变。
• Int8 激活量化与 Int4 权重量化：通过 ExecuTorch，在 iPhone 15 Pro 上，量化后的 Qwen3-4B 模型可达到 14.8 tokens/s 的吞吐量，内存占用仅为 3379 MB。
• Int4 + 2:4 稀疏化：在 Llama-3-8B 上，吞吐量提升 2.37 倍，内存占用减少 67.7%。

以下是我们的量化与部署推荐流程：

from transformers import TorchAoConfig, AutoModelForCausalLM
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, PerRow

# 创建量化配置
quantization_config = TorchAoConfig(quant_type=Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()))

# 加载并自动量化
quantized_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-32B",
    dtype="auto",
    device_map="auto",
    quantization_config=quantization_config
)

如果上述方法不适用，也可使用 快速入门指南 中的 quantize_ API 进行量化。

使用 vLLM 在单台 H100 机器上进行服务：

# 服务器端
VLLM_DISABLE_COMPILE_CACHE=1 vllm serve pytorch/Qwen3-32B-FP8 --tokenizer Qwen/Qwen3-32B -O3

# 客户端
curl http://localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{
  "model": "pytorch/Qwen3-32B-FP8",
  "messages": [
    {"role": "user", "content": "请简要介绍一下大型语言模型。"}
  ],
  "temperature": 0.6,
  "top_p": 0.95,
  "top_k": 20,
  "max_tokens": 32768
}'

对于扩散模型，可以使用 Hugging Face diffusers 进行量化：

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig
from torchao.quantization.granularity import PerGroup

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(granularity=PerGroup(128)))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    torch_dtype=torch.bfloat16,
    device_map="cuda"
)

我们还支持通过 ExecuTorch 将模型部署到边缘设备上，详情请参阅 量化与推理指南。此外，我们还在 Hugging Face 上发布了预量化模型。

🚅 训练

量化感知训练

后训练量化能够生成快速且紧凑的模型，但也可能导致精度下降。我们建议探索量化感知训练（QAT）以克服这一限制，尤其是在使用较低位宽数据类型（如 int4）时。与 TorchTune 合作，我们开发了一套 QAT 流程，相比传统 PTQ 显著提升了模型精度：与 PTQ 相比，Llama3 在 hellaswag 数据集上的准确率恢复了 96%，在 wikitext 数据集上的困惑度下降恢复了 68%。更多细节请参阅 QAT README 和 官方博客：

import torch
from torchao.quantization import quantize_, Int8DynamicActivationIntxWeightConfig, PerGroup
from torchao.quantization.qat import QATConfig

# 准备阶段
base_config = Int8DynamicActivationIntxWeightConfig(
    weight_dtype=torch.int4,
    weight_granularity=PerGroup(32),
)
quantize_(my_model, QATConfig(base_config, step="prepare"))

# 训练模型（未展示）

# 转换阶段
quantize_(my_model, QATConfig(base_config, step="convert"))

用户还可以结合 LoRA 和 QAT，利用此 微调配方，使训练速度相比纯 QAT 提升 1.89 倍。

量化训练

torchao.float8 实现了基于缩放浮点8位数据类型的训练流程，相关理论见 https://arxiv.org/abs/2209.05433。配合 torch.compile 使用时，当前结果表明，在最多 512 张 GPU、参数规模达 405B 的场景下，吞吐量最高可提升 1.5 倍（详情请参阅 PyTorch 博客）：

from torchao.float8 import convert_to_float8_training
convert_to_float8_training(m)

我们的 float8 训练已集成到 TorchTitan 的预训练流程，方便用户直接尝试。更多详情请参阅以下关于 float8 训练支持的博客文章：

• [使用 PyTorch Float8 行级量化加速大规模训练与收敛——基于 Crusoe 2K H200 集群]
• [利用 float8 和 FSDP2 超级加速训练]
• [使用 TorchTitan 在 Amazon SageMaker 上高效预训练 Llama 3 类似架构模型]
• [PyTorch 中的 Float8]

from torchao.sparsity.training import SemiSparseLinear, swap_linear_with_semi_sparse_linear
swap_linear_with_semi_sparse_linear(model, {"seq.0": SemiSparseLinear})

from torchao.optim import AdamW8bit, AdamW4bit, AdamWFp8
optim = AdamW8bit(model.parameters()) # 替换为 Adam4bit 和 AdamFp8 以使用 4 位和 FP8 版本

optim = CPUOffloadOptimizer(model.parameters(), torch.optim.AdamW, fused=True)
optim.load_state_dict(ckpt["optim"])

🔗 集成

TorchAO 已集成到一些领先的开源库中，包括：

• Unsloth 现已支持 QAT：阅读博客 和 指南。
• HuggingFace Transformers 拥有 内置推理后端 和 低比特优化器。
• HuggingFace diffusers 最佳实践，结合 torch.compile 和 TorchAO，位于独立仓库 diffusers-torchao 中。
• vLLM 用于 LLM 服务：使用说明、详细文档。
• 与 MSLK 集成，用于服务器 GPU 上的 SOTA 内核。
• 与 ExecuTorch 集成，用于边缘设备部署。
• Axolotl 支持 QAT 和 PTQ。
• TorchTitan 用于 float8 预训练。
• HuggingFace PEFT 使用 TorchAO 作为其 量化后端，用于 LoRA。
• TorchTune 提供我们的 NF4 QLoRA、QAT 以及 float8 量化微调 等配方。
• SGLang 用于 LLM 服务：使用说明。

🎥 视频

• GPU MODE IRL 主题演讲
• PyTorch 大会上关于低精度数据类型的演讲
• Mastering LLM's 课程中解决 OOM 问题的讲解
• CUDA MODE 高级量化讲座
• Chip Huyen 的 GPU 优化研讨会
• Cohere for AI 社区演讲

💬 引用

如果您觉得 torchao 库很有用，请在您的工作中按以下方式引用它。

@misc{or2025torchao,
  title={TorchAO: PyTorch 原生从训练到推理的模型优化},
  author={Andrew Or 和 Apurva Jain 和 Daniel Vega-Myhre 和 Jesse Cai 和 Charles David Hernandez 和 Zhenrui Zheng 和 Driss Guessous 和 Vasiliy Kuznetsov 和 Christian Puhrsch 和 Mark Saroufim 和 Supriya Rao 和 Thien Tran 和 Aleksandar Samardžić},
  year={2025},
  eprint={2507.16099},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2507.16099},
}

TorchAO 快速上手指南

TorchAO 是 PyTorch 原生的模型优化工具库，专注于从训练到推理的全流程量化与加速。它支持 int4、float8 等多种精度，能够显著降低显存占用并提升推理/训练速度，且无需修改模型架构即可与 torch.compile 和 HuggingFace 模型无缝集成。

环境准备

• 操作系统: Linux (推荐), Windows, macOS
• Python: 3.9 - 3.12
• PyTorch: 2.4+ (建议使用最新稳定版以獲得最佳兼容性)
• 硬件:
  • GPU: NVIDIA GPU (推荐 Ampere 架构及以上，如 A100/H100/RTX 30/40 系列) 或 Intel XPU。
  • CPU: 支持 ARM (Apple Silicon) 及 x86 架构的低比特算子。
• 前置依赖: 确保已安装对应 CUDA 版本的 PyTorch。

注意：TorchAO 强依赖 PyTorch 版本。若遇到兼容性问题，请参考 官方兼容性表。

安装步骤

1. 安装稳定版（推荐）

直接使用 pip 安装最新稳定版本：

pip install torchao

2. 指定 CUDA 版本安装

如果您的环境需要特定 CUDA 版本，可使用以下命令：

# CUDA 12.6
pip install torchao --index-url https://download.pytorch.org/whl/cu126

# CUDA 12.9
pip install torchao --index-url https://download.pytorch.org/whl/cu129

# CPU -only 版本
pip install torchao --index-url https://download.pytorch.org/whl/cpu

3. 开发者安装 (源码编译)

如需开发或体验最新特性（Nightly 版）：

# 安装 Nightly 版本
pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu128

# 或者从源码编译 (需设置 --no-build-isolation)
USE_CUDA=1 pip install -e . --no-build-isolation

可选依赖：MSLK 加速

MSLK 可为部分工作流提供加速内核：

# 稳定版 MSLK
pip install mslk-cuda==1.0.0

基本使用

TorchAO 提供了极简的 API 进行模型量化。以下是两种最常用的场景。

场景一：快速量化现有模型 (Int4 Weight-Only)

适用于推理加速，可大幅减少显存占用。

import torch
from torchao.quantization import Int4WeightOnlyConfig, quantize_

# 假设 model 已经加载并移动到设备
# model = ... 

if torch.cuda.is_available():
    # CUDA 设备量化配置
    config = Int4WeightOnlyConfig(
        group_size=32, 
        int4_packing_format="tile_packed_to_4d", 
        int4_choose_qparams_algorithm="hqq"
    )
elif torch.xpu.is_available():
    # XPU 设备量化配置
    config = Int4WeightOnlyConfig(
        group_size=32, 
        int4_packing_format="plain_int32"
    )
else:
    raise RuntimeError("No supported accelerator found")

# 执行原地量化
quantize_(model, config)

# 之后可直接用于推理
# output = model(input_ids)

场景二：加载即量化 (HuggingFace Transformers 集成)

通过 transformers 库直接加载并量化模型，无需手动修改模型代码。

from transformers import TorchAoConfig, AutoModelForCausalLM
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, PerRow

# 定义量化配置 (例如 Float8 动态量化)
quantization_config = TorchAoConfig(
    quant_type=Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
)

# 加载模型时自动应用量化
quantized_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-32B",       # 替换为目标模型
    dtype="auto",
    device_map="auto",
    quantization_config=quantization_config
)

# 模型已处于量化状态，可直接生成文本

场景三：扩散模型量化 (Diffusers)

针对 Stable Diffusion / Flux 等模型的量化示例：

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig
from torchao.quantization.granularity import PerGroup

# 配置 Transformer 部分的量化策略
pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(granularity=PerGroup(128)))}
)

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    torch_dtype=torch.bfloat16,
    device_map="cuda"
)