FlashAttention(flash-attn)安装
Flash Attention是一种注意力算法,更有效地缩放基于transformer的模型,从而实现更快的训练和推理。由于很多llm模型运行的时候都需要安装flash_attn,比如Llama3,趟了不少坑,最后建议按照已有环境中Python、PyTorch和CUDA的版本精确下载特定的whl文件安装是最佳方式。
FlashAttention的论文:FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
FlashAttention的github地址,建议安装之前读一下官方文档。
官方给出的安装依赖:
- CUDA 11.6 and above.
- PyTorch 1.12 and above.
- Make sure that ninja is installed and that it works correctly (e.g. ninja --version then echo $? should return exit code 0). If not (sometimes ninja --version then echo $? returns a nonzero exit code), uninstall then reinstall ninja (pip uninstall -y ninja && pip install ninja). Without ninja, compiling can take a very long time (2h) since it does not use multiple CPU cores. With ninja compiling takes 3-5 minutes on a 64-core machine.
1. 最终的成功版本
按照自己的Pytorch版本、CUDA版本和Python版本来选择对应版本的flash-attn
,否则会默认下载最新版。