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

相关推荐
闲人编程几秒前
图像去雾算法:从物理模型到深度学习实现
图像处理·人工智能·python·深度学习·算法·计算机视觉·去雾
咔咔学姐kk5 分钟前
大模型微调技术宝典:Transformer架构,从小白到专家
人工智能·深度学习·学习·算法·transformer
Kyln.Wu1 小时前
【python实用小脚本-211】[硬件互联] 桌面壁纸×Python梦幻联动|用10行代码实现“开机盲盒”自动化改造实录(建议收藏)
开发语言·python·自动化
Ms_Big2 小时前
ppliteseg改rknn,部署在嵌入式板,加速模型
人工智能·python·深度学习
折翼的恶魔3 小时前
数据分析:合并
python·数据分析·pandas
百锦再3 小时前
在 CentOS 系统上实现定时执行 Python 邮件发送任务
java·linux·开发语言·人工智能·python·centos·pygame
I'm a winner3 小时前
第五章:Python 数据结构:列表、元组与字典(二)
数据结构·python
番薯大佬3 小时前
Python学习-day8 元组tuple
java·python·学习
小文数模4 小时前
2025高教社国赛数学建模C题参考论文(含模型和代码)
python·数学建模·matlab
鸡哥爱技术5 小时前
Django入门笔记
笔记·python·django