pytorch 使用 xformers 库加速多头注意力计算和大幅节省显存

news/2024/4/16 15:41:29/文章来源:https://blog.csdn.net/ONE_SIX_MIX/article/details/129062323

效果概览：
好处：使用 google PALM 架构的小模型做生成任务，改为 xformers 实现后，加速比为 2倍，显存消耗为原来的 1/3 ，非常给力。
缺点：相比pytorch的原生实现，误差略大。。。

xformers 官方github仓库：https://github.com/facebookresearch/xformers
xformers 官方文档：https://facebookresearch.github.io/xformers/
https://facebookresearch.github.io/xformers/components/ops.html#module-xformers.ops

前两周 xformers 官方提供了 pypi 和 whl 轮包
windows 和 linux 均可用，最低版本要求 pytorch 1.13.1 版本

pip 安装 xformers

pip install -U xformers

如果需要用于编码器或需要位置偏置，则需要安装 0.17 以上版本
当前（2023/2/26） v0.17 为预发行版，需要使用 --pre 来安装

pip install --pre -U xformers

使用方法

import torch
from xformers.ops import memory_efficient_attention, LowerTriangularMaskdevice='cuda'
batch = 4
n_head = 8
head_dim = 16
seq_len = 128q = torch.rand(batch, seq_len, n_head, head_dim).to(device)
k = torch.rand(batch, seq_len, n_head, head_dim).to(device)
v = torch.rand(batch, seq_len, n_head, head_dim).to(device)# 使用 causal 掩码
o = memory_efficient_attention(q, k, v, LowerTriangularMask())# 不使用编码
o = memory_efficient_attention(q, k, v)# 使用自定义的 attn_bias，要求 xformers 版本 大于等于 0.17
## 这里的 from_len，to_len 分别代表Decoder的序列长度，Encoder的序列长度
from_len = seq_len
to_len = seq_len
attn_bias = torch.rand(batch, n_head, from_len, to_len).to(device)
o = memory_efficient_attention(q, k, v, attn_bias)

memory_efficient_attention 的大概的等效pytorch实现
来自 https://facebookresearch.github.io/xformers/components/ops.html#module-xformers.ops


def memory_efficient_attention_pytorch(query, key, value, attn_bias=None, p=0., scale=None):# q [batch, seq_len, n_head, head_dim]# k [batch, seq_len, n_head, head_dim]# v [batch, seq_len, n_head, head_dim]# attn_bias [batch, n_head, seq_len, seq_len]if scale is None:scale = 1 / query.shape[-1] ** 0.5query = query * scaleattn = query @ key.transpose(-2, -1)if attn_bias is not None:attn = attn + attn_biasattn = attn.softmax(-1)attn = F.dropout(attn, p)return attn @ value