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

相关推荐
databook7 小时前
Manim实现脉冲闪烁特效
后端·python·动效
程序设计实验室7 小时前
2025年了,在 Django 之外,Python Web 框架还能怎么选?
python
倔强青铜三9 小时前
苦练Python第46天:文件写入与上下文管理器
人工智能·python·面试
用户25191624271112 小时前
Python之语言特点
python
刘立军13 小时前
使用pyHugeGraph查询HugeGraph图数据
python·graphql
数据智能老司机16 小时前
精通 Python 设计模式——创建型设计模式
python·设计模式·架构
数据智能老司机17 小时前
精通 Python 设计模式——SOLID 原则
python·设计模式·架构
c8i18 小时前
django中的FBV 和 CBV
python·django
c8i19 小时前
python中的闭包和装饰器
python
这里有鱼汤1 天前
小白必看:QMT里的miniQMT入门教程
后端·python