在上一篇文章当中，我们详细介绍了 MQA 的动机、原理及其实现过程，并且也提到由于单一的 Key 和 Value 大大降低了模型的表达能力，所以在实际情况是 MQA 在大模型推理中的效果并不太理想，因此使用更为广泛的是分组查询注意力机制（ Grouped Query Attention, GQA）[1]。
那什么又是 GQA 呢？
简单来说它是介于 MQA 和 MHA 之间的一种解决方案，算是 MQA 和 MHA 之间的一种折衷处理，即将 Query 进行分组，然后每一组 Query 对应一个 Key 和 Value，如图1所示。
在图1右侧 GQA 当中，此时的 Queries 被分成了 3 组，因此对应就有 3 组 Keys-Values ，然后每一组内部所有的 Query 共享一个 Key 和 Value。 此时可以看出，MQA 其实就是 GQA 的一种特殊情况，也就是把所有的 Query 看成一组。
看到这里大家脑海有想到了什么吗？
是不是觉得 MQA 和 GQA 的做法有点类似于我们之前在深度学习基础中所介绍的两种算法？
想一想……
对，就是 层归一化和组归一化这两种算法，如图2所示。
在 LN 和 GN 中，LN 是将所有的 feature map 看成是一个整体进行归一化，而 GN 则是分组进行归一化，每一组内的 feature map 共享一组参数，可以看出 LN 也是 GN 的一种特殊情况，即此时每个通道为一个组。 关于这部分内容的详细介绍，可以参见「6.5 组归一化 GroupNorm：原理、公式与 PyTorch 实现」 内容。
介绍到这里你是不是以为 GQA 的内容已经完了？
不错，核心内容确实是完了，照着图1中 GQA 的结构就能够实现，不过 GQA 提出的动机却不仅仅只是这一个。
下面带着大家开始一步步走进 GQA 算法。
关键词：多头注意力、self-attention、MHA、MQA、GQA、LLaMA、Transformer
1. 动机#
尽管 MQA 在推理时的速度相比于 MHA 更快，但是从实际情况来看它的效果并不好；其次，从各个技术的出现时间上来看，MHA 在前而GQA 在后，因此还需要考虑的一个点便是如何将 MHA 方便地转换为 GQA 能够适用的情况。
下面，分别从这两个方面以 GQA 论文中的角度来为大家介绍 GQA 提出的动机。
1.1 MQA 既要快还要好#
尽管 MQA 在推理过程中能够极大地加快模型的推理速度，但是由于所有的注意力头都共享这一个 Key 和 Value ，使得模型的效果有了很大程度的降低，同时还影响了模型训练时的稳定性。因此，如果仅仅只是为了使用 MQA 来加快推理速度而再从头单独训练这样一个独立的模型显然是不可取的。
Multi-query attention (MQA), which only uses a single key-value head, drastically speeds up decoder inference. However, MQA can lead to quality degradation training instability, and moreover it may not be desirable to train a separate model just for faster inference.
什么意思呢，为什么说从头单独训练一个模型又是不可取的呢？
为了说明这个问题，我们先来看看这篇论文提出的背景及时间点。
GQA 是 2023年12月提出的，此时距离 GPT-4 提出已经过去了9个月，距离 1750 亿参数量的 GPT-3 （2020年5月）已经过去了快2年半。作为对比，MQA 提出的时间是 2019年11月，彼此模型的参数规模大约在数十亿这样的量级，例如 2019年2月发布的 GPT-2 参数规模约 15 亿，也就相当于15个BERT-base 左右的规模。
![图 3. 模型参数规模发展图 [2]](/images/featured/202508032228.png)
所以，从千亿规模参数的时代再回过头去看数十亿规模的模型，简直是不要太迷你，说从头开始训练就可以从头开始训练，一点不带犹豫。
然而，现在的大模型动不动就是上百亿千亿的规模，模型已经（采用 MHA ）训练好了，那我们总不能为了使用 MQA 再来重新训练一个模型，这显然是不可取的。
反观 MQA 时代，那时候的模型都还很小，所以可以随时从头训练。
所以，上面所谓的 “从头单独训练这样一个独立的模型” 作者想要表达的正真含义是：先前使用 MQA 的方式都不对，都是基于 MQA 从头开始训练模型的，这样虽然加快了模型的推理速度但是又耗费了资源（重新训练）同时模型效果又不好（共享一个 Key 和 Value）。
因此，在 GQA 这篇论文中，作者的第一个贡献就是提出了通过以增量训练（uptrain）的方式，将 MHA 中的参数复用到 MQA 中，这样既保证了模型推理时的速度，又使得模型的效果比重新从头训练一个 MQA 模型的效果更好。
This work contains two contributions for faster inference with large language models. First, we show that language model checkpoints with multi-head attention (MHA) can be uptrained to use MQA with a small fraction of original training compute. This presents a cost-effective method to obtain fast multi-query as well as high-quality MHA checkpoints.
1.2 比 MQA 更好的 GQA#
尽管通过增量训练的方式将预训练得到的 MHA 转换成 MQA 比单纯直接使用 MQA 从头训练得到的效果更好，但是这依旧没能改变一个事实—— MQA 中多头注意力共享一个 Key 和 Value，而这也极大地降低了模型的表达能力，因此，作者又提出了分组查询注意力机制（ Grouped Query Attention, GQA）这一方法。
最后，经过增量训练以后，GQA 既能在速度上与 MQA 类似，在效果上也与 MHA 几乎接近。
Second, we propose grouped-query attention (GQA), an interpolation between multi-head and multi-query attention with single key and value heads per subgroup of query heads. We show that uptrained GQA achieves quality close to multi-head attention while being almost as fast as multi-query attention.
如图1所示，GQA 的核心思想就是先将 Queries 进行分组，然后每一组内部共享同一个 Key 和 Value，可以看出 MQA 就是 GQA 的一种特殊情况。
总结起来这篇论文一共两点贡献：① 给出了如何通过增量训练将 MHA 中的模型参数迁移到 MQA 中；② 提出了比 MQA 更优的 GQA 算法，并同样也给出了将 MHA 复用到 GQA 的方法，并对 GQA 的性能进行了分析。
下面，开始分别介绍这两部分内容的原理。
2. MQA 增量训练#
所谓 MQA 增量训练指的是将已经完成训练模型中的 MHA 转换成适用于 MQA 的结构，然后再基于此时的权重参数进行追加训练。
那如何将 MHA 中的权重参数转换为 MQA 适用的结构呢？
2.1 MQA 增量训练原理#
对于如何将 MHA 中的权重参换为 MQA 适用的结构，作者给出的最简单的做法就是取平均值，也就类似于在卷积神经网络中的 Mean Pool 做法，如图4所示。
如图4所示，对于已经训练得到的 MHA 权重参数，我们分别取所有头对应的 $W^K$ 和 $W^V$ 的均值，作为 MQA 中的新的权重，这样便完成了模型权重参数的转换。
在完成模型参数的结构转换以后，只需要再使用之前同样的预训练方法，对基于 MQA 的模型进行训练即可。
Generating a multi-query model from a multi-head model takes place in two steps: first, converting the checkpoint, and second, additional pre-training to allow the model to adapt to its new structure.
并且，通常来说转换后的模型只需要训练少量的步数便可到达一个理想的效果。
The converted checkpoint is then pre-trained for a small proportion $\alpha$ of its original training steps on the same pre-training recipe.
从图5中可以看出，当 $\alpha=0.1$ 的时候（即再训练之前步数的 $10\%$），基于 GQA 增量训练后的模型便能达到 MHA 类似的效果，而 MQA 则是因为模型本身表达能力的欠缺。
当然，对于模型结构转换的方法，除了图4中的均值方法，还可以直接选择多头中的一组 Key-Value或者是直接随机初始化，然后再进行增量训练。不过根据作者实验来看，均值法是三者当中效果最好的。
The projection matrices for key and value heads are mean pooled into single projection matrices, which we find works better than selecting a single key and value head or randomly initializing new key and value heads from scratch.
如图6所示便是不同转换策略下模型的表现情况，其中 “First” 表示旋转第一个头对应的 Key 和 Value。从图中可以看成，均值法的效果是三者当中最好的。
2.2 MQA 增量训练实现#
在介绍完 MQA 增量训练的原理以后，我们再来看如何从零实现这部分代码，完整示例代码可参见 Code/C08_GQA/C01_mqa_uptrain.py 文件。
总体来看，如果我们拿到一个模型基于 MHA 训练得到的大模型，那么首先我们需要载入这些模型参数，然后解析并对每一层中所有的 Keys 和 Values 取均值得到对应层满足 MQA 结构的 Key 和 Value，最后再重新赋值给现在的基于 MQA 的网络结构。
2.2.1 定义网络结构#
为了模拟整个过程，我们首先需要定义一个基于 MHA 的网络结构，然后将训练结束的模型持久化到本地，示例代码如下所示：
1 class Decoder(nn.Module):
2 def __init__(self, num_layers=2, embed_dim=64, num_heads=8, attention_type='mha'):
3 super(Decoder, self).__init__()
4 if attention_type == 'mha':
5 attn = MultiHeadAttention
6 elif attention_type == 'mqa':
7 attn = MultiQueryAttention
8 print(f"## 正在训练基于 {attention_type}")
9 layers = [attn(embed_dim, num_heads) for _ in range(num_layers)]
10 self.decoder_layer = nn.ModuleList(layers)
11 self.classifier = nn.Linear(embed_dim, embed_dim, bias=False)
12 self.norm = nn.RMSNorm(embed_dim)
13
14 def forward(self, x, mask=None):
15 output = x
16 for layer in self.decoder_layer:
17 output = layer(output, output, output, mask)
18 output = self.norm(output)
19 logits = self.classifier(output)
20 return logits在上述代码中，第4~7行是根据参数选择不同的自注意力模块，这样后续我们只需要修改参数就能够复用整个代码。关于 MultiQueryAttention 模块的介绍可以参见文章「MQA 多查询注意力原理与实现：详解 KV Cache 显存优化与 MHA 的差异」。
进一步，我们基于上述代码完成一次前向传播过程，并保存模型参数，示例代码如下：
1 def train(config):
2 x = torch.randn([config.batch_size, config.seq_len, config.embed_dim])
3 mask = torch.full((1, 1, config.seq_len, config.seq_len), float("-inf"))
4 mask = torch.triu(mask, diagonal=1)
5 model = Decoder(config.num_layers, config.embed_dim,
6 config.num_heads, config.attention_type)
7 if os.path.exists(config.target_type):
8 print(f"## 存在模型文件 {config.target_type}，正在进行增量训练……")
9 model.load_state_dict(torch.load(config.target_type))
10 output = model(x, mask)
11 torch.save(model.state_dict(), f"{config.attention_type}.pt")
12 print(f"模型保存成功: {config.attention_type}.pt")
13 print(output.shape)在上述代码中，第7~9行是判断，如过当前已经存在转换后的模型参数，则进行增量训练。第10~11行便是模拟计算一次前向传播过程，然后再持久化保存到本地，文件名为 mha.pt。
2.2.2 转换模型参数#
在得到模型参数 mha.pt 以后，下一步我们需要将其转换成适合 MQA 的结构。
一般来说，在这之前我们可以先载入模型参数并打印出相关参数的信息，然后分析如何转换。
如下所示，便是 mha.pt 中每个参数的信息：
# para_names:[decoder_layer.0.wq.weight], shape:torch.Size([72, 72])
# para_names:[decoder_layer.0.wk.weight], shape:torch.Size([72, 72])
# para_names:[decoder_layer.0.wv.weight], shape:torch.Size([72, 72])
# para_names:[decoder_layer.0.wo.weight], shape:torch.Size([72, 72])
# para_names:[decoder_layer.1.wq.weight], shape:torch.Size([72, 72])
# para_names:[decoder_layer.1.wk.weight], shape:torch.Size([72, 72])
# para_names:[decoder_layer.1.wv.weight], shape:torch.Size([72, 72])
# para_names:[decoder_layer.1.wo.weight], shape:torch.Size([72, 72])
# para_names:[classifier.weight], shape:torch.Size([72, 72])
# para_names:[norm.weight], shape:torch.Size([72])可以看到，一共有2个 decoder_layer 层， 每一层中都有 wq、wk、wv 和 wo 这4个权重参数，并且最后还有一个分类层和归一化层。关于上面各个参数的命名规则介绍大家可以参见「第8讲 BERT SQuAD 阅读理解实战：预测与滑动窗口长文本处理」部分的内容。
这里尤其需要注意一点的是：对于 wq、wk 和 wv 这三个权重参数来说，尽管形状都是 [72,72] 也即 [embed_dim, embed_dim]，但是我们需要明白它的本质其实是 [embed_dim, num_head * head_dim]，并且由于 PyTorch 中 Linear 层权重参数的存储形式是列优先，所以上面 wq、wk 和 wv 的形状其实是 [num_head * head_dim, embed_dim]。
关于 Linear 中权重形状的介绍有兴趣可以参见文章「10090-Linear中权重的形状」。
此时我们知道，下一步我们只需要对每一层中 wk 和 wv 对应的权重参数进行转换即可，示例代码如下：
1 def convert_checkpoint(config):
2 __converted_key__ = ['wk', 'wv']
3 loaded_params = torch.load(config.source_type)
4 new_state_dict = deepcopy(loaded_params)
5 for layer_id in range(config.num_layers):
6 for key in __converted_key__:
7 para_name = f"decoder_layer.{layer_id}.{key}.weight"
8 para = loaded_params[para_name].transpose(0,1)
9 # 由 [num_head * head_dim, embed_dim] 转换成 [embed_dim, num_head * head_dim]
10 para = para.reshape(config.embed_dim, config.num_heads, -1)
11 # [embed_dim, num_heads, head_dim]
12 print(f"参数 {para_name} 转换成形状[embed_dim, num_heads, head_dim]: {para.shape}")
13 para_mean = torch.mean(para, dim=1)
14 print(f"参数 {para_name} 转换成形状[embed_dim, head_dim]: {para_mean.shape}")
15 new_state_dict[para_name] = para_mean.transpose(0, 1) # [head_dim, embed_dim]
16 print(f"参数 {para_name} 处理完毕\n")
17 torch.save(new_state_dict, f"{config.target_type}")在上述代码中，第3行是载入本地模型参数。第4行是深度拷贝一份载入后的模型参数，因为后续我们会在此基础上进行修改。第5行开始便是遍历每一层中的权重参数。第6行是对我们需要进行转换的参数进行处理。第7行是构造得到对应的参数名。第8行是把权重参数的形状转换成符合我们直觉的方式，即 [embed_dim, num_head * head_dim]。第9~13行是取到对应的参数，并以均值法进行转换。第13行则是得到转换后的参数。第15行是把权重参数的形状转换成符合 Linear 中规范的形式。第17行是将转换结束后的参数实例化保存到本地，文件名为 mqa.pt。
上述代码运行时将会输出类似如下信息：
1 参数 decoder_layer.0.wk.weight 转换成形状[embed_dim, num_heads, head_dim]: torch.Size([72, 6, 12])
2 参数 decoder_layer.0.wk.weight 转换成形状[embed_dim, head_dim]: torch.Size([72, 12])
3 参数 decoder_layer.0.wk.weight 处理完毕
4
5 参数 decoder_layer.0.wv.weight 转换成形状[embed_dim, num_heads, head_dim]: torch.Size([72, 6, 12])
6 参数 decoder_layer.0.wv.weight 转换成形状[embed_dim, head_dim]: torch.Size([72, 12])
7 参数 decoder_layer.0.wv.weight 处理完毕
8
9 参数 decoder_layer.1.wk.weight 转换成形状[embed_dim, num_heads, head_dim]: torch.Size([72, 6, 12])
10 参数 decoder_layer.1.wk.weight 转换成形状[embed_dim, head_dim]: torch.Size([72, 12])
11 参数 decoder_layer.1.wk.weight 处理完毕
12
13 参数 decoder_layer.1.wv.weight 转换成形状[embed_dim, num_heads, head_dim]: torch.Size([72, 6, 12])
14 参数 decoder_layer.1.wv.weight 转换成形状[embed_dim, head_dim]: torch.Size([72, 12])
15 参数 decoder_layer.1.wv.weight 处理完毕2.2.3 增量训练#
在完成参数结构的调整以后，我们可以再次基于上面的 train(config) 方法来完成基于 MQA 的模型的训练过程，示例代码如下：
1 if __name__ == '__main__':
2 config = Config()
3 # 第一步： 训练 mha 并保存模型
4 train(config)
5 # 第二步： 分析 mha 参数
6 analysis(config.source_type)
7 # 第三步： 将 mha 转换为 mqa 并保存
8 convert_checkpoint(config)
9 # 第四步： 基于 mqa 增量训练
10 config.__dict__['attention_type'] = 'mqa'
11 train(config)以上就是如何将基于 MHA 的模型参数转换为符合 MQA 结构的过程，以及如何继续进行增量训练。
3. GQA 增量训练#
在介绍完 MQA 的增量训练过程以后，我们再来看如何基于 GQA 来同样完成上述整个流程。不过在这之前先来给大家介绍一下 GQA 的具体原理，以及如何一步步实现 GQA 注意力结构。
3.1 GQA 原理#
假设多头输入 Query、Key 和 Value 分别为 $X_q\in\mathbb{R}^{T_t\times d}$ 、 $X_k\in\mathbb{R}^{T_s\times d}$ 和 $X_v\in\mathbb{R}^{T_s\times d}$ ，其中 $T_s$ 和 $T_t$ 分别表示源序列和目标序列的长度，$d$ 表示模型维度。同时，$h$ 表示多头个数，$g$ 表示分组数量，$s$ 表示每一组中多头的个数，即 $sg=h$。那么第 $i$ 个头对应的三个线性层映射权重分别为 $W^i_q\in\mathbb{R}^{d\times d_h}$ 、 $W^{\lfloor(i-1)/s\rfloor+1}_k\in\mathbb{R}^{d\times d_h}$ 和 $W^{\lfloor(i-1)/s\rfloor+1}_v\in\mathbb{R}^{d\times d_h}$ ，其中 $i=1,2,...,h$， $\lfloor x\rfloor$ 表示对 $x$ 向下取整（即整除）， 且通常 $d=h\times d_h$ 。
此时，第 $i$ 个头的注意力计算过程可以表示为
$$ \begin{aligned} Q_i &= X_qW^i_q\in\mathbb{R}^{T_t\times d_h}\\[2ex] \quad K_i &= X_kW^{\lfloor(i-1)/s\rfloor+1}_k\in\mathbb{R}^{T_s\times d_h}\\[2ex] \quad V_i &= X_vW^{\lfloor(i-1)/s\rfloor+1}_v\in\mathbb{R}^{T_s\times d_h}\\[2ex] \text{head}_i(Q_i,K_i,V_i)&=\text{softmax}\left(\frac{Q_iK_i^T}{\sqrt{d_h}}\right)V_i\in\mathbb{R}^{T_t\times d_h} \end{aligned}\tag{1} $$这里需要注意的是，因为同一组中的 Key 和 Value 是共享的，所以实际情况中每一组里的 $K_i$ 只需要计算一次即可，式(1)中只是一种通用的表达式。
进一步，多头注意力的计算过程为
$$ \begin{aligned} Y_{\text{MHA}}=\text{Concat}(\text{head}_1,...,\text{head}_h)W^O \end{aligned} \tag{2} $$其中 $W^O\in\mathbb{R}^{d\times d}$ 。
接着，根据式(1)我们可以得到计算 3 个线性变换的时间复杂度分别为 $O(T_tdd_hh)$ 、 $O(T_sdd_hg)$ 和 $O(T_sdd_hg)$ ，$Q_iK_i^T$ 对应的时间复杂度为 $O(T_td_hT_sh)$ ，计算注意力权重的时间复杂度为 $O(T_tT_sd_hh)$ 。因为模型超参数是固定的，GQA 的整体时间复杂度为 $O(T_tT_sd_hh)$。同时，由于 GQA 中每一组中所有头共享一个 Key 和 Value ，所以缓存 $K$ 和 $V$ 的空间复杂度均为 $O(T_sd_hg)$
此时，如果仅考虑推理时的场景那么 $T_t=T_s=n$，则 GQA 整体的时间复杂度为 $O(n^2d)$ ，空间复杂度为 $O(nd_hg)$ 。
3.2 GQA 实现过程#
3.2.1 GQA 实现策略#
根据式(1)和式(2)的计算过程，我们便可以完成整个 GQA 的实现过程。同时，由于每一组中的所有头共享一个 Key 和 Value ，所有在实际实现时我们可以先计算得到每个组中对应的 Key 和 Value ，然后再在组内进行复制，如图7所示。
在图7中，我们先计算得到每一组中共享的 Key 和 Value ，然后再在每组内部进行复制，这样便得到了每个头完整的三个状态，进一步便可以完整整个 GQA 多头注意力的计算过程。
因此，我们首先需要实现这么一个辅助函数来完成上面的复制过程，示例代码如下：
1 def repeat_kv(x, n_heads_group):
2 bs, seq_len, groups, head_dim = x.shape
3 # n_heads_group 每个 group 中头的个数
4 # groups 分组数量
5 if groups == 1:
6 return x
7 return (x[:, :, :, None, :] # [bs, seq_len, groups,1, head_dim ]
8 .expand(bs, seq_len, groups, n_heads_group, head_dim)
9 # [bs, seq_len, groups, n_heads_group, head_dim ]
10 .reshape(bs, seq_len, groups * n_heads_group, head_dim))在上述代码中，第5~6行表示如果只有一个组则直接返回，这是因为此时 Key 和 Value 的形状均为 [bs, seq_len, 1, head_dim] ，在计算时会自动出发 PyTorch 中的广播机制。如果第3个维度不是1则不会出发广播机制，需要手动实现。第7~10行便是整个复制过程，大家按维度标识即可理解。
3.2.2 从零实现 GQA#
进一步，我们便可以来实现 GQA 的整个计算过程，示例代码如下：
1 class GroupedQueryAttention(nn.Module):
2 def __init__(self, embed_dim=64, num_heads=8, num_groups=8):
3 super().__init__()
4 ......
5 self.num_groups = num_groups
6 self.num_heads_groups = num_heads // num_groups
7 self.wq = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=False)
8 self.wk = nn.Linear(self.embed_dim, self.num_groups * self.head_dim, bias=False)
9 self.wv = nn.Linear(self.embed_dim, self.num_groups * self.head_dim, bias=False)
10 self.wo = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
11
12 def forward(self, query, key, value, mask=None):
13 bsz, tgt_len, _ = query.shape
14 src_len = key.size(1)
15 xq = self.wq(query)
16 xk, xv = self.wk(key), self.wv(value)
17 xk = xk.view(bsz, src_len, self.num_groups, self.head_dim)
18 xv = xv.view(bsz, src_len, self.num_groups, self.head_dim)
19 xk = repeat_kv(xk, self.num_heads_groups) # [bsz, src_len, num_heads, head_dim]
20 xv = repeat_kv(xv, self.num_heads_groups)
21 queries = xq.view(bsz, tgt_len, self.num_heads, self.head_dim)
22 keys, values = xk.transpose(1, 2), xv.transpose(1, 2)
23 queries = queries.transpose(1, 2) # [bsz, num_heads, tgt_len, head_dim]
24 scores = torch.matmul(queries, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
25 if mask is not None: # [1, 1, tgt_len, src_len]
26 scores = scores + mask
27 scores = F.softmax(scores.float(), dim=-1).type_as(xq)
28 output = torch.matmul(scores, values) # [bsz, num_heads, tgt_len, head_dim]
29 output = output.transpose(1, 2).contiguous().view(bsz, tgt_len, -1)
30 return self.wo(output)在上述代码中，第5~6行是分别指定分组数量以及每一组中多头的数量。第8~9行是定义 Key 和 Value 对应的线性层映射权重。第15~16行则是分别计算三个线性变换，其中 xk 和 xv 的形状均为 [bsz, src_len, num_groups * head_dim]。第19~20行则是在组内对 Key 和 Value 进行复制，结束后的形状均为 [bsz, src_len, num_heads, head_dim] 。此时从形状上看便回到了正常的 MHA 计算过程，后续代码就不再赘述。
同时，如过我们将第22行代码中 keys 的部分值输出，即 print(keys[0][:4, :3, :5])，那将会得到类似如下结果：
tensor([[[-0.0911, 0.5632, 0.6265, 0.4045, -0.2997],
[ 0.1754, -0.2272, -0.1159, 0.4430, -0.5355],
[-1.0640, -0.2887, 0.2030, -0.1093, -1.0711]],
[[-0.0911, 0.5632, 0.6265, 0.4045, -0.2997],
[ 0.1754, -0.2272, -0.1159, 0.4430, -0.5355],
[-1.0640, -0.2887, 0.2030, -0.1093, -1.0711]],
[[ 1.3268, 0.0737, -0.3388, 0.1127, -0.1139],
[-0.4012, 0.1279, 1.1473, 0.5551, 0.5691],
[ 0.1228, 0.9680, -0.1508, -0.9461, 0.8109]],
[[ 1.3268, 0.0737, -0.3388, 0.1127, -0.1139],
[-0.4012, 0.1279, 1.1473, 0.5551, 0.5691],
[ 0.1228, 0.9680, -0.1508, -0.9461, 0.8109]]],
grad_fn=<SliceBackward0>)此时 keys 的形状为 [bsz, num_heads, src_len, head_dim] ，所以上面结果展示的便是前4个头的输出情况。根据 num_heads = 6 ，num_groups = 3 可知，每一组中有两个头，对应上面两组中每一组内部两个头的 Key 都是相等的。
到此，对于 GQA 的实现过程就介绍完了，以上完整示例代码可参见 Code/C08_GQA/C02_gqa.py 文件。
3.3 GQA 增量训练实现#
在完成 GQA 的实现以后，我们再来看如何基于 MHA 持久化的模型使用 GQA 结构进行增量训练。
3.3.1 转换模型参数#
如同上面介绍 MHA 的增量训练流程一样，首先仍旧需要实现一个辅助函数将 MHA 对应的权重参数转换为适用于 GQA 的结构。转换思路同 MQA 一样，采用 Mean Pool 即可，实现流程如图 8所示。
如图8所示便是对有6个头来说的多头注意力 MHA ，当分组数为3时，采用 Mean Pool 策略进行转换时的示意图。
When converting a multi-head checkpoint to a GQA checkpoint, we construct each group key and value head by mean-pooling all the original heads within that group.
当然，除了这样的均值法以外，我们也可以手动选择每一组 Query 中对应的 Key 和 Value 。
根据图8中的示意过程，我们便可以完成相应的代码实现过程，示例代码如下所示：
1 def convert_checkpoint(config):
2 __converted_key__ = ['wk', 'wv']
3 loaded_params = torch.load(config.source_type)
4 new_state_dict = deepcopy(loaded_params)
5 for layer_id in range(config.num_layers):
6 for key in __converted_key__:
7 para_name = f"decoder_layer.{layer_id}.{key}.weight"
8 para = loaded_params[para_name].transpose(0, 1)
9 # 由 [num_head * head_dim, embed_dim] 转换成 [embed_dim, num_head * head_dim]
10 para = para.reshape(config.embed_dim, config.num_heads, -1)
11 # [embed_dim, num_heads, head_dim]
12 print(f"{para_name} [embed_dim, num_heads, head_dim]: {para.shape}")
13 para = para.view(config.embed_dim, config.num_groups, config.num_heads_group, -1)
14 # [embed_dim, num_groups, num_heads_group, head_dim]
15 print(f" {para_name} [embed_dim,num_groups, num_heads_group, head_dim]: {para.shape}")
16 para_mean = torch.mean(para, dim=2).view(config.embed_dim, -1)
17 print(f" {para_name} [embed_dim,num_groups * head_dim]: {para_mean.shape}")
18 new_state_dict[para_name] = para_mean.transpose(0, 1) # [num_groups * head_dim, embed_dim]
19 print(f" {para_name} 处理完毕\n")
20 torch.save(new_state_dict, f"{config.target_type}")在上述代码中，第8行首先将权重参数的形状转换成符合我们直觉的方式，即 [embed_dim, num_head * head_dim]。第10行是将参数转换成 [embed_dim, num_heads, head_dim] 形状，然后再按组进行拆分，即第13行代码所示。第16行代码则是在 num_heads_group 找个维度上取均值。第18行是把权重参数的形状转换成符合 Linear 中规范的形式。第20行代码是将处理完成后的结果持久化到本地，文件名为 gqa.pt。
上述代码运行时将会输出类似如下信息：
1 decoder_layer.0.wk.weight [embed_dim, num_heads, head_dim]: torch.Size([72, 6, 12])
2 decoder_layer.0.wk.weight [embed_dim,num_groups, num_heads_group, head_dim]: torch.Size([72, 3, 2, 12])
3 decoder_layer.0.wk.weight [embed_dim,num_groups * head_dim]: torch.Size([72, 36])
4 decoder_layer.0.wk.weight 处理完毕
5
6 decoder_layer.0.wv.weight [embed_dim, num_heads, head_dim]: torch.Size([72, 6, 12])
7 decoder_layer.0.wv.weight [embed_dim,num_groups, num_heads_group, head_dim]: torch.Size([72, 3, 2, 12])
8 decoder_layer.0.wv.weight [embed_dim,num_groups * head_dim]: torch.Size([72, 36])
9 decoder_layer.0.wv.weight 处理完毕
10
11 decoder_layer.1.wk.weight [embed_dim, num_heads, head_dim]: torch.Size([72, 6, 12])
12 decoder_layer.1.wk.weight [embed_dim,num_groups, num_heads_group, head_dim]: torch.Size([72, 3, 2, 12])
13 decoder_layer.1.wk.weight [embed_dim,num_groups * head_dim]: torch.Size([72, 36])
14 decoder_layer.1.wk.weight 处理完毕
15
16 decoder_layer.1.wv.weight [embed_dim, num_heads, head_dim]: torch.Size([72, 6, 12])
17 decoder_layer.1.wv.weight [embed_dim,num_groups, num_heads_group, head_dim]: torch.Size([72, 3, 2, 12])
18 decoder_layer.1.wv.weight [embed_dim,num_groups * head_dim]: torch.Size([72, 36])
19 decoder_layer.1.wv.weight 处理完毕3.3.2 定义网络结构#
在得到转换后的参数结构以后，我们只需要对上面第2.2.1节中的 Decoder 做细微调整即可用于 GQA 的增量训练，示例代码如下：
1 class Decoder(nn.Module):
2 def __init__(self, num_layers=2, embed_dim=64, num_heads=8, attention_type='mha', num_groups=8):
3 super(Decoder, self).__init__()
4 if attention_type == 'mha':
5 attn = MultiHeadAttention
6 elif attention_type == 'gqa':
7 attn = GroupedQueryAttention
8 print(f"## 正在训练基于 {attention_type}")
9 layers = [attn(embed_dim, num_heads, num_groups=num_groups) for _ in range(num_layers)]
10 self.decoder_layer = nn.ModuleList(layers)
11 self.classifier = nn.Linear(embed_dim, embed_dim, bias=False)
12 self.norm = nn.RMSNorm(embed_dim)在上述代码中，第6~7行我们修改了 gqa 这一选项来使用基于 GQA 的网络结构。
3.3.3 增量训练#
最后，我们使用如下 train() 方法便可以模拟整个训练过程：
1 def train(config):
2 x = torch.randn([config.batch_size, config.seq_len, config.embed_dim])
3 mask = torch.full((1, 1, config.seq_len, config.seq_len), float("-inf"))
4 mask = torch.triu(mask, diagonal=1)
5 model = Decoder(config.num_layers, config.embed_dim, config.num_heads,
6 config.attention_type, config.num_groups)
7 if os.path.exists(config.target_type):
8 print(f"## 存在模型文件 {config.target_type}，正在进行增量训练……")
9 model.load_state_dict(torch.load(config.target_type))
10 output = model(x, mask)
11 torch.save(model.state_dict(), f"{config.attention_type}.pt")
12 print(f"模型保存成功: {config.attention_type}.pt")
13 print(output.shape)以上便是基于 GQA 结构的增量训练过程，完整示例代码可以参见 Code/C08_GQA/C03_gqa_uptrain.py 文件。
4. 复杂度分析#
在介绍完 GQA 的相关原理及实现以后，我们再来看看 GQA 的复杂度等情况，以及将其与之前的 MHA 和 MQA 进行一个对比。
4.1 复杂度分析对比#
根据上面第3.1节内容及上一篇文章的介绍，我们分析得到如果仅考虑推理时的场景那么 $T_t=T_s=n$，则 MHA 整体的时间复杂度为 $O(n^2d)$ ，空间复杂度为 $O(nd)$ ；MQA 整体的时间复杂度为 $O(n^2d)$ ，空间复杂度为 $O(nd_h)$ ； GQA 整体的时间复杂度为 $O(n^2d)$ ，空间复杂度为 $O(nd_hg)$ 。
根据以上结果可以看出，MHA、MQA 和 GQA 在时间复杂度上并没有降低，这因为它们的 Query 数量没有变，所以仍旧需要计算 $h$ 次 $Q_iK^T$ ，所以在时间开销的量级上并没有减少。
但是，从空间复杂度上来看，MHA 为 $O(nd)$ 、 MQA 为 $O(nd_h)$ 而 GQA 为 $O(nd_hg)$ ，所以 MQA 相较于 MHA 会减少 $h$ 倍存储空间，而 GQA 则将取决于分组数量 $g$ 的影响。因此 MQA 对显存节省非常显著，而 GQA 则是两者之间的一个平衡。
4.2 参数缓存对比#
接着，我们继续来看看 MHA 、 MQA 和 GQA 的参数量和 KV 缓存大小的差异。
在 MHA 中因为每个注意力头有独立的 Query、Key、Value 投影矩阵，其参数量为 $(3dd_h)h=3d^2$ 。在 MQA 中，所有的头共享一个 Key 和 Value ，所以分别对应只有一个投影矩阵，其参数量为 $hdd_h+2dd_h=d^2+2dd_h$ 。在 GQA 中，每一组头重共享一个 Key 和 Value ，所以分别对应只有 $g$ 个投影矩阵，其参数量为 $hdd_h+2gdd_h=d^2+2gdd_h$ 。
可以看出，MQA 参数量是最小的，而 MHA 是大的，而 GQA 则居于两者之间。
同时，由于在 MHA 中 KV Cache 随 $h$ 数线性增长，所以长序列推理时极容易出现显存爆炸的情况；而在 MQA 中只是 MHA 的 $1/h$ 倍的显存；在 GQA 中则是 MHA 的 $g/h$ 倍的显存。
现在假设模型在推理时 $h=32$，$n=2048$，$d_h=128$，参数精度为 float16 （即每个参数占用 2 个字节）。由此可以得到此时如果是MHA 其对应 KV 缓存占用显存大小为 $(nd_hh2)2=2048\times128\times32\times2\times2=33554432$ 字节（32MB） ； MQA 对应 KV 缓存大小为 $(nd_h2)2=2048\times128\times1\times2\times2=1048576$ 字节（1MB）；GQA 对应的 KV 缓存大小为 $(nd_hg2)2=2048\times128\times g\times2\times2=1048576g$ 字节（g MB）。
根据上述信息，我们可以得到如下对比结果。
| 时间复杂度 | 空间复杂度 | 参数量 | KV 缓存 | |
|---|---|---|---|---|
| MHA | $O(n^2d)$ | $O(nd)$ | $3d^2$ | $hx$ |
| GQA | $O(n^2d)$ | $O(nd_hg)$ | $d^2+2gdd_h$ | $gx$ |
| MQA | $O(n^2d)$ | $O(nd_h)$ | $d^2+2dd_h$ | $x$ |
5. 总结#
虽然在 MQA 中所有头都共享一个 Key 和 Value 极大减少缓存和带宽，但是效果却并不理想。但是作为 MQA 的泛化， GQA 能同时兼顾推理效率和模型质量，所以更加适合大模型在推理环境中使用。同时，在 GQA 中我们可以通过调整分组数来平衡模型的推理速度和之旅，因此也更加灵活。
引用#
[1] Ainslie J, Lee-Thorp J, De Jong M, et al. Gqa: Training generalized multi-query transformer models from multi-head checkpoints[J]. arXiv preprint arXiv:2305.13245, 2023.
[2] https://shaheryaryousaf.medium.com/the-evolution-of-language-models-ee8d696ec352
[3] Hugo Touvron, Louis Martin, et al. Llama 2: Open Foundation and Fine-Tuned Chat Models[J]. arXiv preprint arXiv:2307.09288, 2023.