TransnormerLLM 中 FlashLinearAttention 的纯pytorch实现

Github 仓库:https://github.com/One-sixth/flash-linear-attention-pytorch

flash-linear-attention-pytorch

纯 Pytorch 实现 TransnormerLLM 中快速线性注意力算子。

用于学习目的。

如果你希望用于训练模型,你可能要修改为 CUDA 或 Triton 的实现,不然会很慢。

注意

这个算子有精度问题,误差较大,是正常的。

这是因为注意力矩阵没有激活函数,导致注意力矩阵的值很大。

在使用 float16 类型时需要特别小心。

这是一个简单的缓解方法:限制 q 和 k 的值,从而减少float16溢出的可能性。

python 复制代码
q = q / q.norm(-1, keepdim=True)
k = k / k.norm(-1, keepdim=True)
o = linear_attention(q, k, v, m)

使用方法

python 复制代码
import torch
from flash_linear_attention_ops import flash_linear_attention, normal_linear_attention


batch_size = 16
seq_len = 1024
dim = 64
n_head = 12
device = 'cuda'
dtype = torch.float32


Q = torch.randn(batch_size, n_head, seq_len, dim, requires_grad=True, dtype=dtype, device=device)
K = torch.randn(batch_size, n_head, seq_len, dim, requires_grad=True, dtype=dtype, device=device)
V = torch.randn(batch_size, n_head, seq_len, dim, requires_grad=True, dtype=dtype, device=device)
M = torch.randint(0, 2, (1, 1, seq_len, seq_len), device=device, dtype=dtype)

O_flash = flash_linear_attention(Q, K, V, M)
O_normal = normal_linear_attention(Q, K, V, M)

print('O_flash.shape', O_flash.shape)
print('O_normal.shape', O_normal.shape)

print('O diff', (O_flash - O_normal).abs().max().item())

参考引用

https://github.com/OpenNLPLab/TransnormerLLM

https://github.com/shreyansh26/FlashAttention-PyTorch

相关推荐
马尚来6 分钟前
移动端自动化测试Appium,从入门到项目实战Python版
python
天才测试猿19 分钟前
WebUI自动化测试:POM设计模式全解析
自动化测试·软件测试·python·selenium·测试工具·设计模式·测试用例
MonkeyKing_sunyuhua30 分钟前
python线程间怎么通信
android·网络·python
跳跳的向阳花43 分钟前
01、大模型部署方案与Dify的使用
python
西柚小萌新1 小时前
【Python从入门到精通】--Pycharm增加内存
开发语言·python·pycharm
西柚小萌新1 小时前
【深入浅出PyTorch】--7.1.PyTorch可视化1
人工智能·pytorch·python
我是华为OD~HR~栗栗呀1 小时前
华为OD-23届考研-Java面经
java·c++·后端·python·华为od·华为·面试
PKNLP1 小时前
Transformer模型
人工智能·深度学习·transformer
渡我白衣1 小时前
深度学习进阶(一)——从 LeNet 到 Transformer:卷积的荣光与注意力的崛起
人工智能·深度学习·transformer
许泽宇的技术分享1 小时前
百刀打造ChatGPT:nanochat极简LLM全栈实现深度解析
chatgpt·transformer·大语言模型·nanochat