flash-attn windows环境安装踩坑日记

本文介绍了如何在Windows环境中安装FlashAttention开源包,由于官方提供的是Linux版本,故需编译源码。作者分享了解决编译问题的方法,包括选择合适的PyTorch和CUDA版本,以及下载预编译的Windowswheel文件。
摘要由CSDN通过智能技术生成

        flash-attn python包是对Flash attention的开源实现。Flash attention是一种具有IO-Awareness的快速和内存效率的精确注意力,具体信息参考:FlashAttention: 具有IO-Awareness的快速和内存效率的精确注意力 - 知乎
论文:https://arxiv.org/abs/2205.14135

        因为官方发布的wheel是linux版的,所以在windows环境下安装时需要编译源码。编译源码过程非常慢,耗时很长且易于出错。尝试若干次,遇到各种疑难杂症后放弃。网上找到了有大牛编译好的windows版,下载地址:https://github.com/bdashore3/flash-attention/releases。这里有四个候选项:

        这里注意,需要安装PyTorch2.1.x和cuda 12.x。具体选择哪个下载,需要先运行pip debug --verbose,根据输出里面的Compatible tags来选择兼容的wheel文件。

        下载相应的wheel文件并安装:pip install "flash_attn-2.4.1+cu121torch2.1cxx11abiFALSE-cp39-cp39-win_amd64.whl"