【python】flash-attn安装

x66ccff2024-10-30 6:05

这个命令：

确保使用正确的 CUDA 12.6 工具链

设置必要的 CUDA 环境变量

包含了常见的 GPU 架构支持

利用你的128核心进行并行编译

复制代码

# 清理之前的安装
proxychains4 pip uninstall -y flash-attn

# 获取 CUDA 路径
CUDA_PATH=$(dirname $(dirname $(which nvcc)))

# 使用 proxychains4 安装
CUDA_HOME=$CUDA_PATH \
TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0" \
MAKEFLAGS="-j128" \
CMAKE_BUILD_PARALLEL_LEVEL=128 \
CMAKE_GENERATOR="Ninja" \
CFLAGS="-march=native -O3" \
CXXFLAGS="-march=native -O3" \
proxychains4 pip install flash-attn --no-build-isolation

# 如果上面失败，可以尝试这个备选命令（从源码安装）
CUDA_HOME=$CUDA_PATH \
TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0" \
MAKEFLAGS="-j128" \
CMAKE_BUILD_PARALLEL_LEVEL=128 \
CMAKE_GENERATOR="Ninja" \
CFLAGS="-march=native -O3" \
CXXFLAGS="-march=native -O3" \
proxychains4 pip install git+https://github.com/Dao-AILab/flash-attention.git