安装

本页介绍 Triton-distributed-ascend 的硬件与软件前置要求,以及在从源码构建之前如何准备 环境。

Triton-distributed-ascend 目前没有发布预编译的 wheel 包,因此安装即意味着从源码构建。 请先完成本页内容,然后参阅 从源码构建

硬件要求

Triton-distributed-ascend 面向以下昇腾产品:

  • Atlas A2

  • Atlas A3

  • Ascend 950

分布式 kernel 需要主机上至少有两张可见的 NPU,因为所有教程和测试都会为每张设备启动一个 进程。单设备机器足以构建本项目,但无法运行通信示例。

软件要求

组件

版本

说明

CANN

9.1.0 及以上

提供工具包、毕昇编译器以及 set_env.sh

Python

3.8 及以上

建议使用虚拟环境或 conda 环境

PyTorch

2.0 及以上

requirements.txt 中固定为 2.7.1

torch_npu

与所用 PyTorch 及 CANN 版本匹配

requirements.txt 中固定为 2.7.1.post8

CMake

3.19 及以上

C++/MLIR 构建所需

GLIBC

2.28 及以上

更低版本的发行版无法加载预编译的 LLVM

libstdc++

3.4.30 及以上

更低版本的发行版无法加载预编译的 LLVM

备注

PyTorch 与 torch_npu 的版本必须匹配。每个 torch_npu 版本都是针对某个特定的 PyTorch 次版本构建的,在其他版本上导入会失败。requirements.txt 中固定的组合 (torch 2.7.1 搭配 torch-npu 2.7.1.post8)是本项目测试验证过的组合。

1. 安装 CANN

CANN 提供昇腾工具包、用于构建 AscendNPU-IR 的毕昇编译器,以及后续每个步骤都需要 source 的环境脚本。请根据你的操作系统和安装方式参考官方指南:

请安装 toolkit 包,而不仅仅是 runtime,因为构建过程需要头文件和编译器可执行文件。

安装完成后,source 环境脚本。默认根目录为 /usr/local/Ascend

source /usr/local/Ascend/ascend-toolkit/set_env.sh

该脚本会设置 ASCEND_HOME_PATH,后续构建步骤依赖此变量。确认它已被正确赋值:

echo $ASCEND_HOME_PATH

如果该变量为空,说明脚本没有正确执行,AscendNPU-IR 与 shmem 的构建将会失败。每开一个新的 shell 都需要重新 source 该脚本,因此建议将其加入 shell 配置文件。

验证驱动能够识别你的 NPU:

npu-smi info

2. 安装 PyTorch 与 torch_npu

请为你的环境安装对应的 PyTorch 版本以及匹配的 torch_npu 插件。兼容性矩阵与安装方式请参阅 Ascend Extension for PyTorch 文档。

3. 安装 Python 依赖

在仓库根目录执行:

pip install -r requirements.txt

这会安装固定版本的 torch 与 torch-npu 组合、用于 Python 绑定的 pybind11,以及 CANN 自身所需的 Python 包(attrscythonnumpydecoratorsympycffipyyamlpathlib2psutilprotobufscipyrequestsabsl-py)。

备注

requirements.txt 从 CPU wheel 索引源拉取 torch。设备支持由 torch_npu 提供,而不是来自 特定设备的 torch 构建版本,因此这是预期行为。

验证环境

在继续之前,请检查基础软件栈能够正常导入并识别你的硬件:

python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version__); print(torch.npu.device_count())"

如果设备数为零,说明 torch_npu 无法访问驱动。请重新检查 npu-smi info,并确认已在当前 shell 中 source 了 set_env.sh

后续步骤