windows下安装大模型加速包flash-attn

liuian 2025-01-06 14:01 75 浏览

Flash-attn（Flash Attention），作为一种用于神经网络模型的注意力机制，在处理序列数据时有显著的优势，可以提高计算效率并减少内存使用，特别是Transformer模型。它通过优化数据布局和计算流程，降低了注意力计算的内存访问开销，在本地部署大模型可以减轻GPU的开销，加快推理速度。项目地址：https://github.com/Dao-AILab/flash-attention

然而，官方发布的版本只支持linux，没有windows包，如果要在windows下安装，要么使用docker,要么自己编译，而自己编译耗时很长，而且也会遇到各种问题。这里我使用的是Brian Dashore编译好的whl. 下载地址：https://github.com/bdashore3/flash-attention/releases

要注意选择对应的版本，首先查看自己的cuda版本，在CMD 里输入：nvcc --version

我的本地环境如下：

系统 windows10 ,python版本11,cuda版本12.1

下载对应的版本为：flash_attn-2.5.9.post1+cu122torch2.3.1cxx11abiFALSE-cp311-cp311-win_amd64.whl

下载到本地后执行，pip安装：

pip install d:\downloads\flash_attn-2.5.9.post1+cu122torch2.3.1cxx11abiFALSE-cp311-cp311-win_amd64.whl

注意：安装完成后，如果你以前已经安装过PyTorch，需要重新安装一遍，否则运行会报错。

打开网址https://pytorch.org/get-started/locally/ 选择对应的版本

根据pytorch.org给出的提示命令，在本地执行安装：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

安装完成，写一个简单的脚本检验一下。

test.py

import flash_attn
print(flash_attn.__version__)

运行结果如果正确输出版本号，则表示安装成功。

pip安装指定版本的包

上一篇：从全流程的角度来了解python包的使用，你也许会有不一样的认识
下一篇：Python包管理入门:轻松掌握pip命令与虚拟环境

windows下安装大模型加速包flash-attn

相关推荐

飞牛OS入门安装遇到问题，如何解决?

系统C盘清理:微信PC端文件清理，扩大C盘可用空间步骤

如何在 iPhone 和 Android 上恢复已删除的抖音消息

Boost高性能并发无锁队列指南:boost::lockfree::queue

大模型手册: 保姆级用CherryStudio知识库

西门子博途中如何输入读取和编辑date and time变量

用什么工具在Win中查看8G大的log文件?

我的Excel打开后是一堆乱码，如何解决?

如何在 Windows 10 或 11 上通过命令行安装 Node.js 和 NPM

威联通NAS安装阿里云盘WebDAV服务并添加到Infuse