Flash-Attention代码调用尝试

原创

已于 2023-11-30 10:00:16 修改 · 9.5k 阅读

标签

#深度学习 #人工智能

于 2023-10-10 16:48:36 首次发布

本文介绍了如何通过源码方式在PyTorch中应用Flash-Attention，包括原理、环境配置、模型ChatGLM2-6b的调用方法和优化后的性能比较，展示了FlashAttention在内存占用和速度上的优势。

Flash-Attention代码调用尝试

本文主要介绍通过如何通过源码方式使用flash-attention，以实现更自由的调用。

1.介绍

Flash-attention原理：

论文：
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
Paper: https://arxiv.org/abs/2205.14135
IEEE Spectrum article about our submission to the MLPerf 2.0 benchmark using FlashAttention. FlashAttention

FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Tri Dao

Paper: https://tridao.me/publications/flash2/flash2.pdf

源码：
https://github.com/Dao-AILab/flash-attention

FlashAttention-2 硬件支持

Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100).
Turing GPUs 只能使用FlashAttention 1.x.

Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).
All head dimensions up to 256. Head dim > 192 backward requires A100/A800 or H100/H800.

pytorch2.0版本已内置了flash-attention1

2.环境配置

CUDA 11.6 and above
PyTorch 1.12 and above

以及LLM基本运行环境
我的环境是：
transformers 4.33.1
torch 2.0.1+cu118
torchaudio 2.0.2+cu118
torchvision 0.15.2+cu118
accelerate 0.22.0
sentencepiece 0.1.99

install flash-attention
在线安装：
pip install flash-attn --no-build-isolation

源码编译安装：
python setup.py install

3.代码实现

模型chatglm2-6b

3.1 模型调用

模型加载时使用modeling_chatglm.py 而非transformers的AutoModel加载，因为要对modeling中的AttenCore进行修改

from transformers import AutoTokenizer
from modeling_chatglm import ChatGLMModel, ChatGLMForConditionalGeneration

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = ChatGLMForConditionalGeneration.from_pretrained(model_path, trust_remote_code=True).cuda

最低0.47元/天解锁文章