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

相关推荐
蓝纹绿茶20 小时前
bash:**:pip:***python: 错误的解释器: 没有那个文件或目录
开发语言·python·pip
START_GAME20 小时前
深度学习Diffusers:用 DiffusionPipeline 实现图像生成
开发语言·python·深度学习
Deamon Tree21 小时前
后端开发常用Linux命令
linux·运维·python
卡卡恩1 天前
使用uv创建系统全局python执行环境
python
查士丁尼·绵1 天前
笔试-座位调整
python
飞翔的佩奇1 天前
【完整源码+数据集+部署教程】【运动的&足球】足球场地区域图像分割系统源码&数据集全套:改进yolo11-RFAConv
前端·python·yolo·计算机视觉·数据集·yolo11·足球场地区域图像分割系统
MYX_3091 天前
第四章 多层感知机
开发语言·python
盼哥PyAI实验室1 天前
《Python爬虫 + 飞书自动化上传》全流程详细讲解
爬虫·python·飞书
StarPrayers.1 天前
(补)CNN 模型搭建与训练:PyTorch 实战 CIFAR10 任务的应用
人工智能·pytorch·cnn
时空无限1 天前
conda 管理 python 版本和虚拟环境
python·conda