不知道是不是也有很多小伙伴像我一样,经常一装机就把最新的cuda安装上去了。等装到组件时,才发现cuda和其它组件各种不适配。这次把nvcc直接装到12.8了,然后发现flash attn没有合适的wheel,就只能自己编译了。下面记录下我的编译过程:
我的环境是ubuntu 24.04, ubuntu 22下也是一样的。
设置环境:
export CUDA_HOME=/usr/local/cuda-12.8 # 替换 <version> 为你的CUDA版本号
export PATH=CUDA_HOME/bin:PATH
export CPATH=CUDA_HOME/include:CPATH # 对于GCC, Clang等C/C++编译器
export LD_LIBRARY_PATH=CUDA_HOME/lib64:LD_LIBRARY_PATH
安装编译器:
pip install ninja
安装编译工具链:
apt-get update && sudo apt-get install -y nvidia-cuda-toolkit
pip install --upgrade pip setuptools wheel
下载 flash atten源码:
git clone https://github.com/Dao-AILab/flash-attention.git
准备编译:
cd flash-attention
python setup.py bdist_wheel
开始编译:
MAX_JOBS=1 python setup.py install
上面为的MAX_JOBS根据你的内存来设置,我的32G内存,默认3个JOB会OOM,被kill导致编译不过。所以就设置了1个并发编译