在autoDL平台上从源码安装openAi Triton

介绍

• 这是一个如何在autoDL平台上下载并安装openAi Triton的简单步骤。
• Triton由openAi开发，它是一种为了深度学习和高性能计算设计的一种编译器和编程语言。

配置

autoDL 镜像：Miniconda conda3 Python 3.10(ubuntu22.04) Cuda 11.8
GPU：RTX 3090(24GB) * 1
CPU：14 vCPU Intel® Xeon® Platinum 8362 CPU @ 2.80GHz
内存：45GB

参考文档

• OpenAI/Triton MLIR 第零章: 源码编译
• triton-lang/triton

参考官方文档和网上一些大牛的文章，发现有很多方法可以安装Triton，快速安装、从源码编译（gcc，python)、源码编译（llvm，make)、源码编译（llvm，python)。因为后续希望研究Triton的源码，且本人只成功使用llvm和pip -e 命令成功安装了Triton。这里只介绍最后一种方法。

步骤

1.登录autoDL平台

注册并登录服务器的方法可参考AutoDL官网

2.学术加速

用于github资源下载的加速
source /etc/network_turbo

3.克隆 triton

cd ~
git clone --recursive https://github.com/openai/triton.git

后续的命令是为了克隆triton的子项目triton_shared，但是貌似triton和一些之项目之间的连接已经失效了。

git submodule update --init --recursive
git submodule status

这里替换成下述命令

cd ~/triton/third_party
export TRITON_PLUGIN_DIRS=$(pwd)/triton_shared
git clone --recurse-submodules https://github.com/microsoft/triton-shared.git triton_shared

这里会clone triton_shared成功，且再次clone一遍triton失败。（这里没搞清楚为什么还要clone一遍triton源码，不过不影响后续使用）

4. 建立venv环境，并安装支持包

cd ~/triton
python -m venv .venv --prompt triton
source /root/triton/.venv/bin/activate

为了后续开机自动启动环境可vim ~/.bashrc
在文件中加入source /root/triton/.venv/bin/activate
source ~/.bashrc

安装一些支持包，为了保证正确运行，最好都安装。

sudo apt update
sudo apt install git gcc clang clangd lldb tree tldr
pip install --upgrade pip
pip install --upgrade setuptools  
pip install ninja cmake wheel; # build-time dependencies
pip install scipy numpy torch pytest lit pandas matplotlib

5. 构建 LLVM

查看当前triton需要的LLVM版本
cat ~/triton/cmake/llvm-hash.txt
我这里是
e4790ce2096c805b7598c353f9e7fa513701f843

cd ~
git clone https://github.com/llvm/llvm-project
cd llvm-project  # your clone of LLVM.
git checkout e4790ce2096c805b7598c353f9e7fa513701f843
mkdir build
cd build
cmake -G Ninja -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=ON ../llvm -DLLVM_ENABLE_PROJECTS="mlir;llvm" 
-DLLVM_TARGETS_TO_BUILD="host;NVPTX;AMDGPU"）
或者
cmake -G Ninja -DCMAKE_BUILD_TYPE=Release ../llvm -DLLVM_ENABLE_PROJECTS="mlir;llvm"    -DLLVM_BUILD_EXAMPLES=ON   -DLLVM_PARALLEL_COMPILE_JOBS=7    -DLLVM_PARALLEL_LINK_JOBS=7    -DLLVM_TARGETS_TO_BUILD="X86;NVPTX;AMDGPU"    -DMLIR_ENABLE_CUDA_RUNNER=ON    -DCMAKE_BUILD_TYPE=Release    -DLLVM_ENABLE_ASSERTIONS=ON    -DCMAKE_C_COMPILER=clang    -DCMAKE_CXX_COMPILER=clang++    -DLLVM_ENABLE_RTTI=ON    -DLLVM_INSTALL_UTILS=ON    -DMLIR_INCLUDE_INTEGRATION_TESTS=ON  

这里要注意DLLVM_TARGETS_TO_BUILD中不能有RISCV，如果选择后面编译会报错，貌似是现在RISCV的支持还不完善。

ninja -j8

6. 构建 triton

export LLVM_BUILD_DIR=/root/llvm-project/build
cd ~/triton
LLVM_INCLUDE_DIRS=$LLVM_BUILD_DIR/include \
LLVM_LIBRARY_DIR=$LLVM_BUILD_DIR/lib \
LLVM_SYSPATH=$LLVM_BUILD_DIR \
pip install -e python -i  https://pypi.tuna.tsinghua.edu.cn/simple

最后一行命令可以多试几次，失败原因大都是网络不好。

7. 测试

• 在python中import triton如果没有错代表triton安装成功。
• 查找triton_shared地址sudo find / -name triton-shared-opt
  找到地址后输入下述命令（替换成你的地址），没有错代表triton_shared安装成功。
  /root/triton/python/build/cmake.linux-x86_64-cpython-3.12/third_party/triton_shared/tools/triton-shared-opt --triton-tolinalg --help

后面可以使用triton愉快的玩耍了。