【python】flash-attn安装

时间:2024-10-30 09:08:12

在这里插入图片描述
这个命令:

确保使用正确的 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