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

本文涉及的产品
智能开放搜索 OpenSearch行业算法版,1GB 20LCU 1个月
实时计算 Flink 版,5000CU*H 3个月
实时数仓Hologres,5000CU*H 100GB 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
546
分享
相关文章
加速LLM大模型推理,KV缓存技术详解与PyTorch实现
大型语言模型(LLM)的推理效率是AI领域的重要挑战。本文聚焦KV缓存技术,通过存储复用注意力机制中的Key和Value张量,减少冗余计算,显著提升推理效率。文章从理论到实践,详细解析KV缓存原理、实现与性能优势,并提供PyTorch代码示例。实验表明,该技术在长序列生成中可将推理时间降低近60%,为大模型优化提供了有效方案。
112 15
加速LLM大模型推理,KV缓存技术详解与PyTorch实现
一键部署 Qwen3! 0 代码,2 种方式全新体验
Qwen3 正式发布并开源 8 款混合推理模型,包括两款 MoE 模型(Qwen3-235B-A22B 和 Qwen3-30B-A3B)及六个 Dense 模型。这些模型支持 119 种语言,在代码、数学等测试中表现优异,并提供思考与非思考两种模式。依托阿里云函数计算 FC 算力,FunctionAI 平台支持模型服务和应用模板部署,适用于多种场景。用户可通过 Serverless 架构快速构建高弹性、智能化应用,显著降低开发成本,提升效率。试用链接及详细文档已提供,欢迎体验。
9个主流GAN损失函数的数学原理和Pytorch代码实现:从经典模型到现代变体
生成对抗网络(GAN)的训练效果高度依赖于损失函数的选择。本文介绍了经典GAN损失函数理论,并用PyTorch实现多种变体,包括原始GAN、LS-GAN、WGAN及WGAN-GP等。通过分析其原理与优劣,如LS-GAN提升训练稳定性、WGAN-GP改善图像质量,展示了不同场景下损失函数的设计思路。代码实现覆盖生成器与判别器的核心逻辑,为实际应用提供了重要参考。未来可探索组合优化与自适应设计以提升性能。
141 7
9个主流GAN损失函数的数学原理和Pytorch代码实现:从经典模型到现代变体
阶跃星辰推出新一代多模态推理模型 Step-R1-V-Mini:视觉+数学+代码全能王
Step-R1-V-Mini是阶跃星辰推出的多模态推理模型,采用PPO强化学习策略,在视觉推理、数学求解和代码生成等任务中表现优异,支持图文输入与文字输出。
78 0
阶跃星辰推出新一代多模态推理模型 Step-R1-V-Mini:视觉+数学+代码全能王
小米7B参数推理大模型首次开源!Xiaomi MiMo:数学代码双杀,超越32B巨头
小米开源的MiMo推理大模型通过联动预训练与强化学习算法,在7B参数规模下实现数学推理与代码生成能力的突破性提升,技术报告显示其性能超越部分32B级模型。
361 74
小米7B参数推理大模型首次开源!Xiaomi MiMo:数学代码双杀,超越32B巨头
从零开始用Pytorch实现LLaMA 4的混合专家(MoE)模型
近期发布的LLaMA 4模型引入混合专家(MoE)架构,以提升效率与性能。尽管社区对其实际表现存在讨论,但MoE作为重要设计范式再次受到关注。本文通过Pytorch从零实现简化版LLaMA 4 MoE模型,涵盖数据准备、分词、模型构建(含词元嵌入、RoPE、RMSNorm、多头注意力及MoE层)到训练与文本生成全流程。关键点包括MoE层实现(路由器、专家与共享专家)、RoPE处理位置信息及RMSNorm归一化。虽规模小于实际LLaMA 4,但清晰展示MoE核心机制:动态路由与稀疏激活专家,在控制计算成本的同时提升性能。完整代码见链接,基于FareedKhan-dev的Github代码修改而成。
61 9
从零开始用Pytorch实现LLaMA 4的混合专家(MoE)模型
比扩散策略更高效的生成模型:流匹配的理论基础与Pytorch代码实现
扩散模型和流匹配是生成高分辨率数据(如图像和机器人轨迹)的先进技术。扩散模型通过逐步去噪生成数据,其代表应用Stable Diffusion已扩展至机器人学领域形成“扩散策略”。流匹配作为更通用的方法,通过学习时间依赖的速度场将噪声转化为目标分布,适用于图像生成和机器人轨迹生成,且通常以较少资源实现更快生成。 本文深入解析流匹配在图像生成中的应用,核心思想是将图像视为随机变量的实现,并通过速度场将源分布转换为目标分布。文中提供了一维模型训练实例,展示了如何用神经网络学习速度场,以及使用最大均值差异(MMD)改进训练效果。与扩散模型相比,流匹配结构简单,资源需求低,适合多模态分布生成。
85 13
比扩散策略更高效的生成模型:流匹配的理论基础与Pytorch代码实现
从零实现基于扩散模型的文本到视频生成系统:技术详解与Pytorch代码实现
本文介绍了一种基于扩散模型的文本到视频生成系统,详细展示了模型架构、训练流程及生成效果。通过3D U-Net结构和多头注意力机制,模型能够根据文本提示生成高质量视频。
73 1
从零实现基于扩散模型的文本到视频生成系统:技术详解与Pytorch代码实现
GenPRM:思维链+代码验证,通过生成式推理的过程奖励让大模型推理准确率显著提升
本文提出GenPRM,一种生成式过程奖励模型,通过显式Chain-of-Thought推理与代码验证提升大型语言模型性能。针对传统PRMs的局限,GenPRM结合相对进展估计和监督微调,优化推理评估精度。实验表明,GenPRM在ProcessBench及数学任务中显著优于现有方法,且可通过测试时扩展进一步增强性能。然而,该方法在计算开销和跨领域应用上仍存在局限性。
43 0
GenPRM:思维链+代码验证,通过生成式推理的过程奖励让大模型推理准确率显著提升
我定制的通义灵码 Project Rules,用 AI 写出“更懂我”的代码
本文分享了一名全栈开发同学使用通义灵码做代码生成、接口注释、测试代码补全等工作,效率明显提升的体会。
AI助理

你好,我是AI助理

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