Windows环境下flash-attention安装

背景

机器学习里Transformer模型大量使用flash-attention这个自注意力模块，多次安装遇到一堆问题，希望能给大家一点帮助。结合了多篇前人经验，本文只提供我这次成功安装的一些步骤，也可能存在之前安装过部分内容文中有疏漏，还请根据提示多注意。

环境

• Windows 11 + NVDIA Geforce RTX3050 Laptop
• CUDA 12.3.2（好像至少要在12.1.6以上）
• Anaconda + Python
• PyTorch 12.1 + PyTorch内建CUDA 12.1

使用他人编译的版本

注意：如无特殊情况，请不要尝试自我编译，费时费力且存在大量依赖缺失、错误。下文自我编译的部分均请忽略，除非无可用的版本，也建议优先考虑其他自注意力模块，自我编译请做好大量失败和debug的心理准备，下文繁杂的过程最终也没有完全成功，仅供参考。使用他人编译的版本很快就成功了……

网络中大部分教程基于Linux系统，官网没有提供Windows的预编译版本。我们可以借助例如这位大佬编译好的版本进行安装（https://github.com/bdashore3/flash-attention/releases），选择匹配的版本，更新whl等包以防万一。

下载好以后cd过去，类似在线安装一样pip install就可以。

pip install <你下载的whl文件>

编译前的准备工作

（可参见FlashAttention提供的README.md）

• 安装PyTorch
• 安装packaging

 pip install packaging

• 安装ninja（加速编译）

pip install ninja

如果ninja存在问题可以考虑卸载后重装。

正式编译安装

1. 在已经安装Git的情况下，断开梯子，右键Git Gui Here -> Clone Existing Repository -> “Source Location”为"https://github.com/Dao-AILab/flash-attention"，然后自己选一个Target Directory并且注意如果是Browse选择的需要自己后面手打类似于“\flash-attention”（注意这个文件夹不能存在，他的父目录可以存在，但是最后一级必须不存在，要在点击Clone后Git自动创建） -> 确保勾选了“Recursively clone submodules too” -> 点击Clone
2. 在命令提示符中测试cl是否配置正确否则无法正常编译源码。

cl

如果返回的不是类似下面的结果需要配置环境变量。

用于 x64 的 Microsoft ® C/C++ 优化编译器 19.29.30151版
版权所有© Microsoft Corporation。保留所有权利。

用法： cl [ 选项… ] 文件名… [ /link 链接选项… ]

配置环境变量需要找到类似下面的路径（不同版本路径可能不同，在Microsoft Visual Studio里找一下）中是否存在cl.exe文件，然后复制

C:\Program Files (x86)\Microsoft Visual Studio\2019\BuildTools\VC\Tools\MSVC<版本号>\bin\Hostx64\x64

然后Windows搜索栏中搜索“编辑系统环境变量” -> 环境变量(N)… -> 选中“系统变量”中变量名为"Path” -> 编辑(I)… -> 新建(N) -> 粘贴刚刚找到的这个路径 -> 全部点“确定”
3. 去GitHub下载cutlass，解压以后配置CPATH环境变量，值为路径。
4. 解压下载的flash-attention源代码，然后命令提示符中cd到Git复制到的文件夹中，执行如下的命令

set MAX_JOBS=4
python setup.py install

如果不使用命令提示符而是用PowerShell就对应执行

$env:MAX_JOBS="4"
python setup.py install