西安搭建网站,第三方电子商务交易平台,深圳市手机网站建设品牌,基础网站建设代码FlashAttention安装教程
FlashAttention 是一种高效且内存优化的注意力机制实现#xff0c;旨在提升大规模深度学习模型的训练和推理效率。 高效计算#xff1a;通过优化 IO 操作#xff0c;减少内存访问开销#xff0c;提升计算效率。 内存优化#xff1a;降低内存占用…FlashAttention安装教程
FlashAttention 是一种高效且内存优化的注意力机制实现旨在提升大规模深度学习模型的训练和推理效率。 高效计算通过优化 IO 操作减少内存访问开销提升计算效率。 内存优化降低内存占用使得在大规模模型上运行更加可行。 精确注意力保持注意力机制的精确性不引入近似误差。 FlashAttention-2 是 FlashAttention 的升级版本优化了并行计算策略充分利用硬件资源。改进了工作负载分配进一步提升计算效率。 FlashAttention-3FlashAttention-3 是专为 Hopper GPU如 H100优化的版本目前处于 Beta 测试阶段。 常见问题: 安装成功后实际模型代码运行时报错未安装核心原因就是cxx11abiFALSE这个参数表示该包在构建时不启用 C11 ABI。 必须开启不使用才行。否则报错如下 ImportError: This modeling file requires the following packages that were not found in your environment: flash_attn. 最佳安装步骤(方法1)
安装依赖 基础环境cuda12.1、nvcc.安装python示例3.10。安装PyTorch示例orchtorch2.3.0; torchvision0.18.0ninja Python 包 获取releases对应的whl包 - 地址https://github.com/Dao-AILab/flash-attention/releases - 按照系统环境选whl 3. 我的环境对应的包是flash_attn-2.7.2.post1cu12torch2.3cxx11abiTRUE-cp310-cp310-linux_x86_64.whl解释如下 flash_attn: 包的名称表示这个 Wheel 文件是 flash_attn 包的安装文件。2.7.2.post1: 包的版本号遵循 PEP 440 版本规范。 2.7.2: 主版本号表示这是 flash_attn 的第 2.7.2 版本。post1: 表示这是一个“后发布版本”post-release通常用于修复发布后的某些问题。 cu12torch2.3cxx11abiFALSE: 构建标签表示该 Wheel 文件是在特定环境下构建的。 cu12: 表示该包是针对 CUDA 12 构建的。torch2.3: 表示该包是针对 PyTorch 2.3 构建的。cxx11abiFALSE: 表示该包在构建时不启用 C11 ABIApplication Binary Interface。如果安装包后不识别就要选为False的版本。 cp310: Python 版本的标签表示该包是为 Python 3.10 构建的。 cp310: 是 cpython 3.10 的缩写表示该包适用于 CPython 解释器的 3.10 版本。 linux_x86_64: 平台标签表示该包是为 Linux 操作系统和 x86_64 架构即 64 位 Intel/AMD 处理器构建的。.whl: 文件扩展名表示这是一个 Python Wheel 文件。Wheel 是 Python 的一种二进制分发格式用于快速安装包。
如何安装
可以使用 pip 安装这个 Wheel 文件
pip install flash_attn-2.7.2.post1cu12torch2.3cxx11abiTRUE-cp310-cp310-linux_x86_64.whl --no-build-isolation 常规安装步骤方法二 安装依赖 CUDA 工具包或 ROCm 工具包PyTorch 1.12 及以上版本packaging 和 ninja Python 包 pip install packaging ninja安装 FlashAttention # 后面--no-build-isolation参数是为了pip 会直接在当前环境中构建包使用当前环境中已安装的依赖项。
# 如果当前环境缺少构建所需的依赖项构建过程可能会失败。
pip install flash-attn --no-build-isolation或从源码编译 # 下载源码后进行编译
cd flash-attention
python setup.py install运行测试 export PYTHONPATH$PWD
pytest -q -s test_flash_attn.py补充说明 4.1 上面运行时建议设置参数MAX_JOBS限制最大进程数不然系统容易崩。本人在docker下安装直接干重启了所以建议如下方式运行 MAX_JOBS4 pip install flash-attn --no-build-isolation4.2 如果运行时会出现警告且推理速度依旧很慢需要继续从源码安装rotary和layer_normcd到源码的那两个文件夹执行 python setup.py install进行安装如果命令报错弃用可能要用easy_install命令。
接口使用
import flash_attn_interface
flash_attn_interface.flash_attn_func()硬件支持
NVIDIA CUDA 支持
支持 GPUAmpere、Ada 或 Hopper 架构 GPU如 A100、RTX 3090、RTX 4090、H100。数据类型FP16 和 BF16。头维度支持所有头维度最大至 256。
AMD ROCm 支持
支持 GPUMI200 或 MI300 系列 GPU。数据类型FP16 和 BF16。后端支持 Composable Kernel (CK) 和 Triton 后端。
性能优化
Triton 后端
Triton 后端的 FlashAttention-2 实现仍在开发中目前支持以下特性
前向和反向传播支持因果掩码、变长序列、任意 Q 和 KV 序列长度、任意头大小。多查询和分组查询注意力目前仅支持前向传播反向传播支持正在开发中。
性能改进
并行编译使用 ninja 工具进行并行编译显著减少编译时间。内存管理通过设置 MAX_JOBS 环境变量限制并行编译任务数量避免内存耗尽。
结论
FlashAttention 系列通过优化计算和内存使用显著提升了注意力机制的效率。无论是研究人员还是工程师都可以通过本文提供的安装和使用指南快速上手并应用于实际项目中。随着 FlashAttention-3 的推出针对 Hopper GPU 的优化将进一步推动大规模深度学习模型的发展。
参考链接
FlashAttention 源码
