ao

环境要求

• 已安装 Python 与 pip 环境（版本要求见官方 compatibility table）
• PyTorch 生态依赖版本需与 torchao 匹配，参考 GitHub issue 2919 的兼容表
• 若使用 float8 / nvfp4 推理内核，建议额外安装 MSLK
• 若使用 mxfp8 MoE 训练内核，需额外安装 apache-tvm-ffi 与 nvidia-cutlass-dsl
• 如需预览最新特性，可改用 nightly 源安装

安装与启动步骤

1. 1安装稳定版

   在 Python 环境中用 pip 安装最新稳定版 TorchAO，这是最推荐的安装方式。

   pip install torchao

   复制

2. 2安装 MSLK 内核

   可选步骤。为 float8 和 nvfp4 推理提供加速内核，需 CUDA 13.0 对应的索引源。

   pip install mslk --index-url https://download.pytorch.org/whl/cu130

   复制

3. 3装 MoE 训练依赖

   可选步骤。使用 mxfp8 MoE 训练内核前需安装 tvm-ffi 与指定版本的 cutlass-dsl。

   pip install apache-tvm-ffi
   pip install nvidia-cutlass-dsl==4.5.2 nvidia-cutlass-dsl-libs-base==4.5.2 nvidia-cutlass-dsl-libs-cu13==4.5.2

   复制

4. 4按硬件换安装源

   根据本机 CUDA/XPU 环境选择对应索引源，纯 CPU 环境用 cpu 源。

   pip install torchao --index-url https://download.pytorch.org/whl/cu126
   pip install torchao --index-url https://download.pytorch.org/whl/cu129
   pip install torchao --index-url https://download.pytorch.org/whl/xpu
   pip install torchao --index-url https://download.pytorch.org/whl/cpu

   复制

5. 5安装 nightly 版

   想体验最新功能时，从 nightly 源安装预览版本，注意稳定性低于正式版。

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

   复制

6. 6开发者源码安装

   克隆仓库后在根目录源码安装，必须带 --no-build-isolation，否则编译会失败。

   git clone https://github.com/pytorch/ao.git
   cd ao
   USE_CUDA=1 pip install -e . --no-build-isolation

   复制

7. 7量化模型权重

   在代码中导入量化接口，用 Int4WeightOnlyConfig 把模型权重量化为 int4。

   import torch
   from torchao.quantization import Int4WeightOnlyConfig, quantize_
   quantize_(model, Int4WeightOnlyConfig(group_size=32, int4_packing_format="tile_packed_to_4d", int4_choose_qparams_algorithm="hqq"))

   复制

8. 8配置 float8 量化

   推理/服务侧可用 TorchAoConfig 指定 float8 动态激活加 float8 权重，粒度按行。

   quantization_config = TorchAoConfig(quant_type=Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()))

   复制

关键配置

如何确认成功

执行 Quick Start 中的 int4 量化代码片段，模型权重被成功量化为 int4 且无报错，即安装可用。

常见问题

注意事项

• MSLK、apache-tvm-ffi、cutlass-dsl 均为可选项，只在使用对应推理/MoE 训练内核时才需要安装。
• 安装源必须与硬件匹配，混用不同 CUDA 版本的 index-url 容易导致依赖冲突。
• nightly 版本为预览性质，生产环境建议使用 pip install torchao 的稳定版。
• 具体依赖版本以官方 compatibility table 为准，安装前先核对。