安装
本页介绍 Triton-distributed-ascend 的硬件与软件前置要求,以及在从源码构建之前如何准备 环境。
Triton-distributed-ascend 目前没有发布预编译的 wheel 包,因此安装即意味着从源码构建。 请先完成本页内容,然后参阅 从源码构建。
硬件要求
Triton-distributed-ascend 面向以下昇腾产品:
Atlas A2
Atlas A3
Ascend 950
分布式 kernel 需要主机上至少有两张可见的 NPU,因为所有教程和测试都会为每张设备启动一个 进程。单设备机器足以构建本项目,但无法运行通信示例。
软件要求
组件 |
版本 |
说明 |
|---|---|---|
CANN |
9.1.0 及以上 |
提供工具包、毕昇编译器以及 |
Python |
3.8 及以上 |
建议使用虚拟环境或 conda 环境 |
PyTorch |
2.0 及以上 |
|
torch_npu |
与所用 PyTorch 及 CANN 版本匹配 |
|
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 包(attrs、cython、numpy、decorator、sympy、cffi、
pyyaml、pathlib2、psutil、protobuf、scipy、requests、absl-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。