大模型中常用的注意力机制GQA详解以及Pytorch代码实现

本文涉及的产品
实时数仓Hologres,5000CU*H 100GB 3个月
智能开放搜索 OpenSearch行业算法版,1GB 20LCU 1个月
实时计算 Flink 版,5000CU*H 3个月
简介: GQA是一种结合MQA和MHA优点的注意力机制,旨在保持MQA的速度并提供MHA的精度。它将查询头分成组,每组共享键和值。通过Pytorch和einops库,可以简洁实现这一概念。GQA在保持高效性的同时接近MHA的性能,是高负载系统优化的有力工具。相关论文和非官方Pytorch实现可进一步探究。

分组查询注意力 (Grouped Query Attention) 是一种在大型语言模型中的多查询注意力 (MQA) 和多头注意力 (MHA) 之间进行插值的方法,它的目标是在保持 MQA 速度的同时实现 MHA 的质量。

这篇文章中,我们将解释GQA的思想以及如何将其转化为代码。

GQA是在论文 GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints paper.中提出,这是一个相当简单和干净的想法,并且建立在多头注意力之上。

GQA

标准多头注意层(MHA)由H个查询头、键头和值头组成。每个头都有D个维度。Pytorch的代码如下:

 fromtorch.nn.functionalimportscaled_dot_product_attention

 # shapes: (batch_size, seq_len, num_heads, head_dim)
 query=torch.randn(1, 256, 8, 64)
 key=torch.randn(1, 256, 8, 64)
 value=torch.randn(1, 256, 8, 64)

 output=scaled_dot_product_attention(query, key, value)
 print(output.shape) # torch.Size([1, 256, 8, 64])
AI 代码解读

对于每个查询头,都有一个对应的键。这个过程如下图所示:

而GQA将查询头分成G组,每组共享一个键和值。可以表示为:

使用可视化的表示就能非常清楚的了解GQA的工作原理,就像我们上面说的那样,GQA是一个相当简单和干净的想法

Pytorch代码实现

让我们编写代码将这种将查询头划分为G组,每个组共享一个键和值。我们可以使用einops库有效地执行对张量的复杂操作。

首先,定义查询、键和值。然后设置注意力头的数量,数量是随意的,但是要保证num_heads_for_query % num_heads_for_key = 0,也就是说要能够整除。我们的定义如下:

 importtorch

 # shapes: (batch_size, seq_len, num_heads, head_dim)
 query=torch.randn(1, 256, 8, 64)
 key=torch.randn(1, 256, 2, 64)
 value=torch.randn(1, 256, 2, 64)

 num_head_groups=query.shape[2] //key.shape[2]
 print(num_head_groups) # each group is of size 4 since there are 2 kv_heads
AI 代码解读

为了提高效率,交换seq_len和num_heads维度,einops可以像下面这样简单地完成:

 fromeinopsimportrearrange

 query=rearrange(query, "b n h d -> b h n d")
 key=rearrange(key, "b s h d -> b h s d")
 value=rearrange(value, "b s h d -> b h s d")
AI 代码解读

然后就是需要在查询矩阵中引入”分组“的概念。

 fromeinopsimportrearrange
 query=rearrange(query, "b (h g) n d -> b g h n d", g=num_head_groups)
 print(query.shape) # torch.Size([1, 4, 2, 256, 64])
AI 代码解读

上面的代码我们将二维重塑为二维:对于我们定义的张量,原始维度8(查询的头数)现在被分成两组(以匹配键和值中的头数),每组大小为4。

最后最难的部分是计算注意力的分数。但其实它可以在一行中通过insum操作完成的

 fromeinopsimporteinsum, rearrange
 # g stands for the number of groups
 # h stands for the hidden dim
 # n and s are equal and stands for sequence length

 scores=einsum(query, key, "b g h n d, b h s d -> b h n s")
 print(scores.shape) # torch.Size([1, 2, 256, 256])
AI 代码解读

scores张量和上面的value张量的形状是一样的。我们看看到底是怎么操作的

einsum帮我们做了两件事:

1、一个查询和键的矩阵乘法。在我们的例子中,这些张量的形状是(1,4,2,256,64)和(1,2,256,64),所以沿着最后两个维度的矩阵乘法得到(1,4,2,256,256)。

2、对第二个维度(维度g)上的元素求和——如果在指定的输出形状中省略了维度,einsum将自动完成这项工作,这样的求和是用来匹配键和值中的头的数量。

最后是注意分数与值的标准乘法:

 importtorch.nn.functionalasF

 scale=query.size(-1) **0.5
 attention=F.softmax(similarity/scale, dim=-1)

 # here we do just a standard matrix multiplication
 out=einsum(attention, value, "b h n s, b h s d -> b h n d")

 # finally, just reshape back to the (batch_size, seq_len, num_kv_heads, hidden_dim)
 out=rearrange(out, "b h n d -> b n h d")
 print(out.shape) # torch.Size([1, 256, 2, 64])
AI 代码解读

这样最简单的GQA实现就完成了,只需要不到16行python代码:

最后再简单提一句MQA:多查询注意(MQA)是另一种简化MHA的流行方法。所有查询将共享相同的键和值。原理图如下:

可以看到,MQA和MHA都可以从GQA推导出来。具有单个键和值的GQA相当于MQA,而具有与头数量相等的组的GQA相当于MHA。

GQA的好处是什么?

GQA是最佳性能(MQA)和最佳模型质量(MHA)之间的一个很好的权衡。

下图显示,使用GQA,可以获得与MHA几乎相同的模型质量,同时将处理时间提高3倍,达到MQA的性能。这对于高负载系统来说可能是必不可少的。

在pytorch中没有GQA的官方实现。所以我找到了一个比较好的非官方实现,有兴趣的可以试试:

https://avoid.overfit.cn/post/58ee0d8f5ed14414bc856080ab748047

作者:Max Shap

目录
打赏
0
4
4
0
547
分享
相关文章
加速LLM大模型推理,KV缓存技术详解与PyTorch实现
大型语言模型(LLM)的推理效率是AI领域的重要挑战。本文聚焦KV缓存技术,通过存储复用注意力机制中的Key和Value张量,减少冗余计算,显著提升推理效率。文章从理论到实践,详细解析KV缓存原理、实现与性能优势,并提供PyTorch代码示例。实验表明,该技术在长序列生成中可将推理时间降低近60%,为大模型优化提供了有效方案。
59 15
加速LLM大模型推理,KV缓存技术详解与PyTorch实现
9个主流GAN损失函数的数学原理和Pytorch代码实现:从经典模型到现代变体
生成对抗网络(GAN)的训练效果高度依赖于损失函数的选择。本文介绍了经典GAN损失函数理论,并用PyTorch实现多种变体,包括原始GAN、LS-GAN、WGAN及WGAN-GP等。通过分析其原理与优劣,如LS-GAN提升训练稳定性、WGAN-GP改善图像质量,展示了不同场景下损失函数的设计思路。代码实现覆盖生成器与判别器的核心逻辑,为实际应用提供了重要参考。未来可探索组合优化与自适应设计以提升性能。
108 7
9个主流GAN损失函数的数学原理和Pytorch代码实现:从经典模型到现代变体
从零开始用Pytorch实现LLaMA 4的混合专家(MoE)模型
近期发布的LLaMA 4模型引入混合专家(MoE)架构,以提升效率与性能。尽管社区对其实际表现存在讨论,但MoE作为重要设计范式再次受到关注。本文通过Pytorch从零实现简化版LLaMA 4 MoE模型,涵盖数据准备、分词、模型构建(含词元嵌入、RoPE、RMSNorm、多头注意力及MoE层)到训练与文本生成全流程。关键点包括MoE层实现(路由器、专家与共享专家)、RoPE处理位置信息及RMSNorm归一化。虽规模小于实际LLaMA 4,但清晰展示MoE核心机制:动态路由与稀疏激活专家,在控制计算成本的同时提升性能。完整代码见链接,基于FareedKhan-dev的Github代码修改而成。
50 9
从零开始用Pytorch实现LLaMA 4的混合专家(MoE)模型
比扩散策略更高效的生成模型:流匹配的理论基础与Pytorch代码实现
扩散模型和流匹配是生成高分辨率数据(如图像和机器人轨迹)的先进技术。扩散模型通过逐步去噪生成数据,其代表应用Stable Diffusion已扩展至机器人学领域形成“扩散策略”。流匹配作为更通用的方法,通过学习时间依赖的速度场将噪声转化为目标分布,适用于图像生成和机器人轨迹生成,且通常以较少资源实现更快生成。 本文深入解析流匹配在图像生成中的应用,核心思想是将图像视为随机变量的实现,并通过速度场将源分布转换为目标分布。文中提供了一维模型训练实例,展示了如何用神经网络学习速度场,以及使用最大均值差异(MMD)改进训练效果。与扩散模型相比,流匹配结构简单,资源需求低,适合多模态分布生成。
73 13
比扩散策略更高效的生成模型:流匹配的理论基础与Pytorch代码实现
从零实现基于扩散模型的文本到视频生成系统:技术详解与Pytorch代码实现
本文介绍了一种基于扩散模型的文本到视频生成系统,详细展示了模型架构、训练流程及生成效果。通过3D U-Net结构和多头注意力机制,模型能够根据文本提示生成高质量视频。
63 1
从零实现基于扩散模型的文本到视频生成系统:技术详解与Pytorch代码实现
基于昇腾用PyTorch实现传统CTR模型WideDeep网络
本文介绍了如何在昇腾平台上使用PyTorch实现经典的WideDeep网络模型,以处理推荐系统中的点击率(CTR)预测问题。
295 66
用PyTorch从零构建 DeepSeek R1:模型架构和分步训练详解
本文详细介绍了DeepSeek R1模型的构建过程,涵盖从基础模型选型到多阶段训练流程,再到关键技术如强化学习、拒绝采样和知识蒸馏的应用。
372 3
用PyTorch从零构建 DeepSeek R1:模型架构和分步训练详解
近端策略优化(PPO)算法的理论基础与PyTorch代码详解
近端策略优化(PPO)是深度强化学习中高效的策略优化方法,广泛应用于大语言模型的RLHF训练。PPO通过引入策略更新约束机制,平衡了更新幅度,提升了训练稳定性。其核心思想是在优势演员-评论家方法的基础上,采用裁剪和非裁剪项组成的替代目标函数,限制策略比率在[1-ϵ, 1+ϵ]区间内,防止过大的策略更新。本文详细探讨了PPO的基本原理、损失函数设计及PyTorch实现流程,提供了完整的代码示例。
842 10
近端策略优化(PPO)算法的理论基础与PyTorch代码详解
Pytorch学习笔记(九):Pytorch模型的FLOPs、模型参数量等信息输出(torchstat、thop、ptflops、torchsummary)
本文介绍了如何使用torchstat、thop、ptflops和torchsummary等工具来计算Pytorch模型的FLOPs、模型参数量等信息。
927 2
优化技巧与策略:提高 PyTorch 模型训练效率
【8月更文第29天】在深度学习领域中,PyTorch 是一个非常流行的框架,被广泛应用于各种机器学习任务中。然而,随着模型复杂度的增加以及数据集规模的增长,如何有效地训练这些模型成为了一个重要的问题。本文将介绍一系列优化技巧和策略,帮助提高 PyTorch 模型训练的效率。
753 0

热门文章

最新文章

AI助理

你好,我是AI助理

可以解答问题、推荐解决方案等