CMU 10423 Generative AI：HW1（编程部分：在GPT-2模型中实现RoPE、GQA）

完整代码和PDF笔记：https://github.com/YM2025/CMU_10423_2024S

1 概述

在“Programming: RoPE and GQA”部分，主要任务是通过结合RoPE（旋转位置嵌入）和GQA（Grouped Query Attention，分组查询注意力）这两种机制，改进现有的GPT模型，并观察这些改进对模型性能的影响。以下是对RoPE和GQA的介绍：

Rotary Positional Embeddings (RoPE)

RoPE是一种相对位置嵌入方法，用来取代传统的绝对位置嵌入。在传统Transformer中，位置信息通过将位置嵌入直接加到输入的词向量中进行传播。而RoPE直接在每一层注意力计算中引入相对位置信息，旋转每个查询和键向量的一部分来嵌入这些信息。

RoPE的数学表达为：通过旋转矩阵对查询向量和键向量进行旋转，计算出新的旋转后的向量，并基于这些旋转向量计算注意力得分。与传统的绝对位置编码不同，RoPE使得位置信息能够在不同层之间保持一致，并且具有较好的相对位置信息表达能力。

你将在此部分作业中实现RoPE，并通过与传统minGPT模型的对比，分析其在训练过程中的表现。你需要在mingpt/model.py文件中的RotaryPositionalEmbeddings和CausalSelfAttention类中进行修改。

Grouped Query Attention (GQA)

GQA是一种用于优化注意力机制的技术。它通过将查询头分组，每个组共享一个键和值头，减少了注意力计算中的资源消耗。GQA的设计提供了一种在计算效率和模型质量之间取得平衡的方案。

GQA的基本思想是，将查询头按组进行分组，每组查询头使用相同的键和值头，从而减少了键和值的计算量。与标准的多头注意力机制不同，GQA在保持较好的注意力表现的同时，能够大幅度减少计算开销和内存占用。

你将在mingpt/model.py文件中实现GQA。你需要在GroupedQueryAttention类中实现这一机制，并进行一系列实验，验证其性能和效率改进。

实验任务

1. RoPE实验：实现RoPE后，你需要绘制RoPE实现和传统minGPT在600次迭代中的训练损失对比图，并进一步分析RoPE在更长序列长度（128和256）下的训练表现。
2. GQA实验：实现GQA后，你需要测量GQA在不同键头数量（如1、2、3、6）下的计算时间和内存消耗，并与标准的多头注意力机制进行对比。此外，还需要观察GQA和多头注意力机制在600次迭代中的训练损失表现。

通过完成这些实验，你将对RoPE和GQA在提升模型性能和效率方面的效果有深入理解。

2 项目文件

该项目的起始代码文件包含在名为“hw1”的目录中，以下是这些文件及其功能的概述：

1. requirements.txt

• 功能：列出了此项目所需的Python依赖包。
• 内容：该文件列出了两种主要的依赖：torch 和 einops，这两个包分别用于深度学习框架和高效的张量操作。

2. input.txt

• 功能：存储了用于训练模型的数据集。
• 内容：这是完整的莎士比亚作品集，大小约为1.1MB。

3. chargpt.py

• 功能：项目的主入口，用于训练Transformer模型。
• 内容：通过命令python chargpt.py启动训练，可以通过添加标志来调整Transformer的配置。此文件使用提供的RoPE和GQA配置来调整模型。

4. mingpt/

这是项目中的核心代码库，包含以下文件：

a. model.py

• 功能：定义了GPT模型的构建，包括Transformer的各个层和注意力机制。

• 内容：这是作业中需要修改的文件，包含基础的Transformer实现。你需要在此文件中实现以下两个类：
  i. RotaryPositionalEmbeddings：用于实现RoPE旋转位置嵌入。
  ii. GroupedQueryAttention：用于实现分组查询注意力机制（GQA）。

  你还需要修改CausalSelfAttention类以整合RoPE。

  b. trainer.py

  • 功能：定义了训练循环和模型优化逻辑。

• 内容：这个文件包含训练GPT模型的逻辑，负责训练步骤、损失计算和参数更新。

  c. utils.py

  • 功能：提供一些辅助功能，如保存日志和配置。

• 内容：这个文件包含帮助函数，用于模型的日志记录、保存模型和配置的管理。

参数配置（Flags）

• 用途：通过向chargpt.py添加不同的标志，你可以修改训练时的各种参数配置。
• 示例：你可以通过命令行标志设置序列长度、查询头和键/值头的数量等。例如：
  • --data.block_size=128：设置模型序列长度为128。
  • --model.rope=True：启用RoPE嵌入。
  • --model.n_query_head=6：设置查询头的数量为6。
  • --model.n_kv_head=3：设置键/值头的数量为3。

通过这些起始代码文件，你可以实现并训练一个基于RoPE和GQA的GPT模型，同时测试其在小型数据集上的性能。

3 本项目的模型剖析

3.1 源数据、数据集的制作

源数据就一个input.txt，里面是莎士比亚《科利奥兰纳斯》戏剧文字内容。内容长这样：

网上找了翻译大概看了下，如下：

txt文本概述：

• 数据量：总计 40000 行文本。
• 字符种类：65种字符，包括大小写字母、标点符号、数字、空格、换行符。
• 总字符数：文本中包含 1,115,394 个字符（包括空格和换行符）。

数据集制作流程：

1. 样本构建：

• 首先，将 txt 文件中的所有文本行首尾相接，形成一个连续的字符序列，总长度为 1,115,394 个字符。
• 在训练过程中，每次从该连续字符序列中，在随机索引位置提取 128 个（这是个超参数）连续的字符，作为一个训练样本。

2. 标签生成：

• 标签与样本对应，但标签内容相对样本内容往后挪了一位。因为要训练模型基于前面的字符预测下一个字符。

3.2 模型结构

模型主体采用了GPT-2的基本结构，包含6层解码器。本项目有两种位置编码方式：

1. 第一种方法是先生成一个 [0, 1, ..., 127] 的序列，然后通过 nn.Embedding() 进行编码，其结果再与样本的nn.Embedding()编码结果进行融合，相当于把128个字符的索引位置编码进去了。
2. 第二种方法就是在attention里面应用RoPE位置编码。

模型架构完整参数如下：

3.3 优化器、学习器、损失函数、超参数等

优化器：adamW

学习率调度器：没用

损失函数：交叉熵

学习率：3e-4

batchsize：64

4 Rotary Position Embeddings (RoPE)

4.1 讲义原文

在本部分中，你将实现旋转位置嵌入（Rotary Position Embeddings，RoPE）（Su et al., 2021）。

背景

在标准的Transformer语言模型中，绝对位置嵌入（Absolute Position Embeddings）被添加到词嵌入的第一层中。随后的层从底层向上传播位置信息。

传统的注意力机制定义如下：

$$q_j=W_q^T{x}_j,\forall j$$

$$k_j=W_k^T{x}_j,\forall j$$

$$s_{t,j}=\frac{k_j^T{q}_t}{\sqrt{d_k}},\forall j,t$$

$$a_t=\text{softmax}(s_t),\forall t$$

其中 $d_k$ 是查询/键/值向量的大小。

RoPE机制

旋转位置嵌入（Rotary Position Embeddings，RoPE）（Su et al., 2021）直接将位置信息融入到每一层的注意力计算中。如果下一层注意力机制的输入是 $X=[x_1,...,x_N{]}^T$，那么我们引入两个函数 $f_q(x_j,j)$ 和 $f_k(x_j,j)$，分别计算位置感知的查询和键。然后，注意力得分计算如下：

$$q_j=W_q^T{x}_j,\forall j$$

$$k_j=W_k^T{x}_j,\forall j$$

$$\widetilde{q_j}=R_{\Theta ,j}{q}_j,\widetilde{k_j}=R_{\Theta ,j}{k}_j$$

$$s_{t,j}=\frac{{\widetilde{k_j}}^T\widetilde{q_t}}{\sqrt{d_k}},\forall j,t$$

$$a_t=\text{softmax}(s_t),\forall t$$

其中 $d=\frac{d_k}{2}$， $W_k,W_q\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$。对于某个固定的绝对位置 $m$，旋转矩阵 $R_{\Theta ,m}\in {\mathbb{R}}^{d_k\times d_k}$ 被定义为：

$$R_{\Theta ,m}=\left(\begin{array}{ccccccc}\cos (m{\theta }_1)& -\sin (m{\theta }_1)& 0& 0& \dots & 0& 0\\ \sin (m{\theta }_1)& \cos (m{\theta }_1)& 0& 0& \dots & 0& 0\\ 0& 0& \cos (m{\theta }_2)& -\sin (m{\theta }_2)& \dots & 0& 0\\ 0& 0& \sin (m{\theta }_2)& \cos (m{\theta }_2)& \dots & 0& 0\\ ⋮& ⋮& ⋮& ⋮& \ddots & ⋮& ⋮\\ 0& 0& 0& 0& \dots & \cos (m{\theta }_{d_k/2})& -\sin (m{\theta }_{d_k/2})\\ 0& 0& 0& 0& \dots & \sin (m{\theta }_{d_k/2})& \cos (m{\theta }_{d_k/2})\end{array}\right)$$

参数 ${\theta }_i$ 在训练之前被固定，定义如下：

Θ = { θ i = 1000 0 − 2 i − 1 d , i ∈ [ 1 , 2 , . . . , d / 2 ] } \Theta = \{ \theta_i = 10000^{-\frac{2i-1}{d}}, \quad i \in [1, 2, ..., d/2] \} Θ={θi​=10000−d2i−1​,i∈[1,2,...,d/2]}

由于 $R_{\Theta ,m}$ 中的稀疏结构，我们可以更高效地计算 $R_{\Theta ,m}$ 与任意向量 $y$ 的矩阵-向量乘法：

$$R_{\Theta ,m}y=\left(\begin{array}{c}{y}_1\\ y_2\\ y_3\\ y_4\\ ⋮\\ y_{d-1}\\ y_d\end{array}\right)\otimes \left(\begin{array}{c}\cos (m{\theta }_1)\\ \cos (m{\theta }_1)\\ \cos (m{\theta }_2)\\ \cos (m{\theta }_2)\\ ⋮\\ \cos (m{\theta }_{d/2})\\ \cos (m{\theta }_{d/2})\end{array}\right)+\left(\begin{array}{c}-y_2\\ y_1\\ -y_4\\ y_3\\ ⋮\\ -y_d\\ y_{d-1}\end{array}\right)\otimes \left(\begin{array}{c}\sin (m{\theta }_1)\\ \sin (m{\theta }_1)\\ \sin (m{\theta }_2)\\ \sin (m{\theta }_2)\\ ⋮\\ \sin (m{\theta }_{d/2})\\ \sin (m{\theta }_{d/2})\end{array}\right)$$

在PyTorch中高效地实现这一过程仍然需要一些技巧。如果我们有一个嵌入矩阵 $Y=[y_1,...,y_N{]}^T\in {\mathbb{R}}^{N\times d_k}$（在实际操作中，这个 $Y$ 可能是查询 $Q$ 或键 $K$），那么我们想要构建一个新的矩阵 $\tilde{Y}=g(Y,\Theta )$，其中 ${\tilde{Y}}_{m,\cdot }=R_{\Theta ,m}{y}_m$。为了简化，我们可以将向量的索引重新排列，使得可以分别处理向量的前半部分和后半部分。简记为 $d=d_k$：

$$\tilde{Y}=g(Y,\Theta )=\left[\begin{array}{cccccc}{Y}_{1,1}& \dots & Y_{1,d/2}& Y_{1,d/2+1}& \dots & Y_{1,d}\\ ⋮& \dots & ⋮& ⋮& \dots & ⋮\\ Y_{N,1}& \dots & Y_{N,d/2}& Y_{N,d/2+1}& \dots & Y_{N,d}\end{array}\right]\otimes \left[\begin{array}{ccc}\cos (1{\theta }_1)& \dots & \cos (1{\theta }_{d/2})\\ ⋮& \dots & ⋮\\ \cos (N{\theta }_1)& \dots & \cos (N{\theta }_{d/2})\end{array}\right]+\left[\begin{array}{cccccc}-Y_{1,d/2+1}& \dots & -Y_{1,d}& Y_{1,1}& \dots & Y_{1,d/2}\\ ⋮& \dots & ⋮& ⋮& \dots & ⋮\\ -Y_{N,d/2+1}& \dots & -Y_{N,d}& Y_{N,1}& \dots & Y_{N,d/2}\end{array}\right]\otimes \left[\begin{array}{ccc}\sin (1{\theta }_1)& \dots & \sin (1{\theta }_{d/2})\\ ⋮& \dots & ⋮\\ \sin (N{\theta }_1)& \dots & \sin (N{\theta }_{d/2})\end{array}\right]$$

或者更简洁地表示为：

$$C=\left[\begin{array}{ccc}1{\theta }_1& \dots & 1{\theta }_{d/2}\\ ⋮& \dots & ⋮\\ N{\theta }_1& \dots & N{\theta }_{d/2}\end{array}\right]$$

$$\tilde{Y}=g(Y,\Theta )=\left[\begin{array}{cc}{Y}_{\cdot ,1:d/2}& Y_{\cdot ,d/2+1:d}\end{array}\right]\otimes \cos (C)+\left[\begin{array}{cc}-Y_{\cdot ,d/2+1:d}& Y_{\cdot ,1:d/2}\end{array}\right]\otimes \sin (C)$$

现在，我们可以高效地计算RoPE嵌入：

$$Q=X{W}_q,K=X{W}_k$$

$$\tilde{Q}=g(Q,\Theta ),\tilde{K}=g(K,\Theta )$$

$$S=\frac{\tilde{Q}\cdot {\tilde{K}}^T}{\sqrt{d_k}}$$

$$A=\text{softmax}(S)$$

你不需要理解论文中的所有数学细节，但可以阅读以便对RoPE形成直觉。

实现：

你将在 minGPT 中实现RoPE。为此，你需要在 mingpt/model.py 文件中修改 RotaryPositionalEmbeddings 和 CausalSelfAttention 类。

4.2 RoPE的代码实现

class RotaryPositionalEmbeddings(nn.Module):
    """
    实现RoPE（旋转位置嵌入）。该模块对输入的查询和键应用旋转位置嵌入，主要通过构建旋转角度的 cos 和 sin 矩阵来完成。
    """
    
    def __init__(self, d: int, base: int = 10_000):
        """
        初始化RoPE模块。
        
        参数：
        d (int): 输入的嵌入维度。
        base (int): 用于生成旋转角度的基数，默认为10_000。
        """
        super().__init__()
        self.d = d  # 嵌入维度（例如，d_model 或 d_query）
        self.base = base  # 基数，控制旋转角度的幅度
        self.cosine_mat = None  # 用于缓存 cos 矩阵
        self.sine_mat = None    # 用于缓存 sin 矩阵

    def _build_cache(self, x: torch.Tensor):
        """
        计算和缓存 cos 和 sin 矩阵，这些矩阵用于对输入进行旋转嵌入。
        此步骤根据输入的维度（序列长度和嵌入维度）动态生成。

        参数：
        x (torch.Tensor): 输入张量，包含批次大小、头数、序列长度和嵌入维度。
        """
        device = x.device  # 获取输入的设备（例如 CPU 或 GPU）
        self.N = x.shape[-2]  # 获取序列长度 N

        # 生成 theta 作为旋转角度，它基于嵌入维度的前一半
        # positions 是嵌入维度的一半（即 d // 2）个位置，从0到 (d//2 - 1)
        positions = torch.arange(0, self.d // 2, device=device)
        # theta 是旋转频率的参数，随着位置增大而衰减
        theta = torch.pow(self.base, -2 ** (positions) / self.d)

        # 生成矩阵 c_matrix，其尺寸为 (N, d//2)，用于储存每个序列位置的旋转频率
        # arange_N 是一个从 1 到 N 的向量，它表示每个位置（从1开始）
        arange_N = torch.arange(1, self.N + 1, device=device).unsqueeze(1)
        # c_matrix 是每个位置乘以 theta，用于生成 cos 和 sin 值
        c_matrix = arange_N * theta.unsqueeze(0)

        # 将 c_matrix 复制两份，形成 (N, d) 维度的矩阵，其中前半部分和后半部分相同
        # 生成对应的 cos 和 sin 矩阵，并扩展维度为 (1, 1, N, d)，以便广播操作
        self.cosine_mat = torch.cos(torch.cat([c_matrix, c_matrix], dim=-1)).unsqueeze(0).unsqueeze(0)
        self.sine_mat = torch.sin(torch.cat([c_matrix, c_matrix], dim=-1)).unsqueeze(0).unsqueeze(0)

    def forward(self, x: torch.Tensor):
        """
        前向传播：对输入张量 x 进行旋转位置嵌入。

        参数：
        x (torch.Tensor): 输入张量，形状为 (batch_size, n_heads, seq_len, d)。
        
        返回：
        torch.Tensor: 应用旋转位置嵌入后的张量。
        """
        # 检查是否需要重新生成 cos 和 sin 矩阵（基于输入的序列长度和嵌入维度）
        if (self.cosine_mat is None or x.shape[-2] != self.cosine_mat.shape[-2] or x.shape[-1] != self.cosine_mat.shape[-1]):
            self._build_cache(x)  # 如果缓存的矩阵尺寸不匹配，则重新构建

        # 获取嵌入维度的一半，便于将张量分成两部分
        half_d = self.d // 2
        
        # 计算 x 的 cos 部分：每个位置的值乘以相应的 cos 值
        x_cos = x * self.cosine_mat
        
        # 计算 x 的 sin 部分：
        # 将 x 的后半部分移到前面，前半部分移到后面，并乘以相应的 sin 值
        x_sin = torch.cat([-x[:, :, :, half_d:], x[:, :, :, :half_d]], dim=3) * self.sine_mat
        
        # 返回两个部分的和，作为旋转嵌入后的输出
        return x_cos + x_sin

4.3 RoPE问题

1 （4 分）

绘制你实现的 RoPE 和原始 minGPT 模型在序列长度为 128 时，训练 600 次迭代的训练损失图。

2 （4 分）

绘制你实现的 RoPE 和原始 minGPT 模型在总计 800 次训练迭代中的训练损失图：前 600 次迭代使用序列长度 128，接下来的 200 次迭代使用序列长度 256。

3 （2 分）

提供一个来自你训练了 600 次迭代（序列长度为 128）的 RoPE 模型的样本。样本的生成应基于你最喜欢的莎士比亚剧本中的第一行作为条件。

4 （2 分）

提供一个来自你训练了 600 次迭代（序列长度为 128），然后又训练了 200 次迭代（序列长度为 256）的 RoPE 模型的样本。样本的生成应基于你最喜欢的莎士比亚剧本中的第一行作为条件。

解答：

以下分别是训练600个step（序列长度=128）和800个step（序列长度=256）情况：

从下图可看出2点：

• RoPE位置嵌入远好于绝对位置嵌入。
• 使用绝对位置嵌入，当序列长度增大时，同样step情况下loss会变大。而如果使用RoPE位置嵌入，当序列长度增大时，训练时的loss会更小。

训练600个step后，切换序列长度的训练结果可以直接参考CMU学生的文件：

文本生成情况展示：

我同时测试了把训练集《科利奥兰纳斯》的第一句话作为输入，以及《罗密欧与朱丽叶》第一句话作为输入，评估模型生成文本的效果。

以下2张图是《罗密欧与朱丽叶》原文：

600step+RoPE编码的模型的预测结果如下：

GPT4的评估

根据你提供的第一句“Gregory, o’my word, we’ll not carry coals.”以及后续生成的文本，我可以从以下几个方面评估大模型的文本生成效果：

1. 语言风格的匹配性（权重：40%）：

• 这句话来自莎士比亚的作品，因此具有古英语的语法和风格。大模型生成的文本尝试模仿这种古英语的形式，比如使用了词汇如“thou”，“hath”，“shall”等。

• 然而，整体生成的句子语法显得相当混乱，句子结构也不太连贯，难以真正展现莎士比亚那种优雅的风格。语言的流畅性和一致性有待提高。

• 语义连贯性（权重：30%）：

  • 文本中的句子大多是独立的，缺乏上下文的连贯性。例如“Shall him by disher; and the may best thy sea;”这类句子含义模糊，且难以理解。语义上，文本没有明确传递某个主题或逻辑进展。
  • 生成的内容尽管模仿了古英语的形式，但没有形成一个有意义的段落，语义比较混乱。

• 创造力（权重：20%）：

  • 尽管生成的文本不连贯，但模型展现了一定的创造力，生成了许多词汇组合和结构。比如“Tit his suppetion”和“deather in toudious”，尽管没有明显的意义，但能看出模型在尝试创造新词和表达。
  • 这种创造力在文本生成中是很重要的，但它需要在连贯的框架下展现出来。

• 模型对特定任务的适应性（权重：10%）：

  • 如果这是一个模仿莎士比亚风格的任务，模型生成的文本还未能完全达到目标。虽然句子表面上看起来具有古英语的风格，但它缺乏莎士比亚作品中的逻辑、语义和深度。

总结：

大模型在语言风格上做了一些尝试，使用了莎士比亚时代的词汇和表达，但生成的文本缺乏语义上的连贯性，句子结构也比较混乱。要提升效果，模型可能需要更好地理解上下文，并生成更符合人类阅读习惯的连贯文本。

以下是训练800个step情况（序列长度=256，使用RoPE）的生成结果：

GPT4的评估2

1. 语言风格的匹配性（权重：40%）：

• 本次评估：本次生成的文本继续模仿莎士比亚式的古英语风格，使用了词汇如“hath”、“thou”、“art”等。然而，与第一次相比，仍然有很多词汇组合显得不自然，且句子结构比较混乱。例如“sweet hath sleign all, that there and being to straight”这样的句子，语法不正确，句子的逻辑也不连贯。
• 对比第一次：与第一次生成相比，语言风格的模仿没有显著提升，模型依然局限于表面形式的仿古英语词汇堆砌，而无法真正模仿莎士比亚句子结构的优雅和严谨。

2. 语义连贯性（权重：30%）：

• 本次评估：与第一次相比，生成的文本仍然缺乏语义连贯性。比如句子“her beggar such and praise are at thy hopest from and throw men”显得无意义，单词之间的关联很弱，整体句子的逻辑性和可读性较差。并且许多句子是独立生成的，无法形成完整的情节或上下文。
• 对比第一次：本次生成的文本在语义连贯性上与第一次的表现相似，没有明显的改进。句子依旧是单独生成的，未能形成连贯的对话或叙述。

3. 创造力（权重：20%）：

• 本次评估：模型在生成文本时展现了创造性，继续生成了一些词汇和表达，虽然部分词汇组合不符合语法，但体现了一些语言的灵活性。例如“to so-prayed against the comparly”虽然没有意义，但模型在创造新的短语和组合上表现了一定的创造力。
• 对比第一次：与第一次相比，模型的创造力水平保持相近，生成了一些新的组合。不过这些创造性的表达并未带来实际的语义提升，仍然是“词汇实验”。

4. 模型对任务的适应性（权重：10%）：

• 本次评估：如果目标是模仿莎士比亚风格，模型本次生成的文本与第一次一样，仍然远未达到目标。尽管使用了古英语词汇，但句子逻辑和意义仍不连贯，无法形成有意义的段落。
• 对比第一次：与第一次相比，本次生成的文本在适应任务方面没有太大改进，依然主要表现在表面上的风格模仿，而不是语义上的理解和生成。

综合对比：

• 语言风格：两次生成的风格模仿相似，都使用了古英语词汇和句式，但缺乏句法的正确性和句子之间的逻辑关联。
• 语义连贯性：两次生成的文本语义上都缺乏连贯性，句子之间没有明确的主题或关联。
• 创造力：两次生成的文本都有一定的创造性，但这种创造性主要体现在无意义的词汇组合上，没有在语义上带来提升。
• 适应性：两次生成的文本都没有很好地完成模仿莎士比亚风格的任务。

总结：

总体而言，本次生成的效果与第一次生成相似，在语言风格和创造性上都有一定的表现，但在语义连贯性和适应性上仍有很大不足。两次生成的文本在莎士比亚风格的模仿上都仅停留在表面，未能深入到句子的逻辑结构和语义层面。

5 Grouped Query Attention (GQA)

5.1 讲义原文

图 1：注意力机制的示意图，展示了具有每个头单独键和值的多头注意力（Multi-head attention），具有分组查询共享公共键和值的分组查询注意力（Grouped-query attention），以及使用单一键和值的多查询注意力（Multi-query attention）。

在本部分中，你将实现分组查询注意力（Grouped Query Attention, GQA）（Ainslie 等, 2023）。

GQA简介

分组查询注意力（GQA）是一种用于神经网络架构的技术，它修改了模型（如Transformer）中使用的注意力机制。该机制将查询头划分为多个组，每组共享一个键和值头。这种方法可以在多查询注意力（MQA）和多头注意力（MHA）之间进行插值，提供了计算效率与模型质量之间的平衡 [图 1]。

令 $h_q$ 表示查询头的数量，$h_{kv}$ 表示键/值头的数量。我们假设 $h_q$ 可以被 $h_{kv}$ 整除，并且 $g=\frac{h_q}{h_{kv}}$ 是每组的大小（即每个键/值向量对应的查询向量数量）。

我们对 GQA 的参数矩阵大小保持一致：$W_q^{(g,i)},W_k^{(g)},W_v^{(g)}\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$，其中 $d_k=\frac{d_{\text{model}}}{h_q}$。然而，现在我们有不同数量的查询、键和值头：

$$X=[x_1,\dots ,x_T{]}^T$$

V ( i ) = X W v ( i ) , ∀ i ∈ { 1 , … , d k v } V^{(i)} = X W_v^{(i)}, \quad \forall i \in \{1, \dots, d_{kv}\} V(i)=XWv(i)​,∀i∈{1,…,dkv​}

K ( i ) = X W k ( i ) , ∀ i ∈ { 1 , … , d k v } K^{(i)} = X W_k^{(i)}, \quad \forall i \in \{1, \dots, d_{kv}\} K(i)=XWk(i)​,∀i∈{1,…,dkv​}

Q ( i , j ) = X W q ( i , j ) , ∀ i ∈ { 1 , … , d k v } , ∀ j ∈ { 1 , … , g } Q^{(i,j)} = X W_q^{(i,j)}, \quad \forall i \in \{1, \dots, d_{kv}\}, \forall j \in \{1, \dots, g\} Q(i,j)=XWq(i,j)​,∀i∈{1,…,dkv​},∀j∈{1,…,g}

上面，我们定义了比键/值向量多 $g$ 倍的查询向量。然后，我们计算每个查询向量 $(i,j)$ 与其对应键 $(i)$ 的缩放点积，并在每组内对查询求和，以得到相似度得分。使用这些相似度得分计算注意力矩阵，但只使用 $h_{kv}$ 个头：

$$S^{(i)}=\sum _{j=1}^g{Q}^{(i,j)}(K^{(i)}{)}^T/\sqrt{d_k}$$

$$A^{(i)}=\text{softmax}(S^{(i)})$$

$$X^{\prime (i)}=A^{(i)}{V}^{(i)}$$

实现细节：

你将在 mingpt/model.py 文件中的 GroupedQueryAttention 类中实现 GQA。你编写的代码大部分会与 CausalSelfAttention 类中的内容相似。

提示：你可能会发现，首先使用 einops.rearrange() 重新实现 CausalSelfAttention 类比使用 tensor.view() 或 tensor.transpose() 更加容易；此外，使用 einops.einsum() 替代 @ 操作符。如果你在这种形式下扩展实现 GroupedQueryAttention，可能会更直接。

• 初始化：

  • 熟悉初始化注意力机制的配置设置，包括查询头、键/值头的数量和嵌入维度。
  • 确保嵌入维度可以被查询头和键/值头的数量整除。

• 正则化：

  • 添加注意力和残差的 dropout 层，以防止过拟合。

• 维度和投影：

  • 实现查询、键和值的线性投影层，考虑维度约束和机制的分组性质。

• 旋转位置嵌入：

  • 如果启用了旋转位置嵌入（RoPE），将 RoPE 集成到查询和键的投影中。（注意：将 RoPE 和 GQA 结合是可选的，但实现相对简单。）

• 前向传播：

  • 在 forward 方法中，按照查询、键和值的投影变换输入。
  • 通过计算分组缩放点积注意力来应用注意力机制。
  • 对注意力进行掩码操作，以确保因果性（即避免未来的 tokens 被关注）。
  • 将注意力与值进行聚合，并将输出投影回嵌入维度。

• 内存效率：

  • 在注意力操作前后，监控并记录 CUDA 内存分配，以分析 GQA 的内存效率。在 CausalSelfAttention 类中有用于监控内存的参考代码。

---

5.2 GQA的代码实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class GroupedQueryAttention(nn.Module):
    """
    实现分组查询注意力机制。
    """

    def __init__(self, config):
        super().__init__()

        # 确保嵌入维度能够整除查询头和键/值头的数量
        assert config.n_embd % config.n_query_head == 0
        assert config.n_query_head % config.n_kv_head == 0

        # 分组数量g（每个键/值头对应多少个查询头）
        self.g = config.n_query_head // config.n_kv_head

        # 键和值的线性投影层，生成键和值的向量，大小为 2 * n_embd
        self.vk_attn = nn.Linear(config.n_embd, 2 * config.n_embd)

        # 查询的线性投影层，生成查询的向量，维度为 g * n_embd
        self.q_attn = nn.Linear(config.n_embd, self.g * config.n_embd)

        # 输出的线性投影层，用于将注意力输出映射回嵌入维度
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)

        # 正则化，防止过拟合
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)

        # 因果遮罩，确保注意力仅关注到当前位置及之前的序列
        self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                     .view(1, 1, config.block_size, config.block_size))

        # 存储键/值头数量和查询头数量
        self.n_head = config.n_kv_head
        self.n_query = config.n_query_head
        self.n_embd = config.n_embd
        
        # 是否使用RoPE（旋转位置嵌入）
        self.rope = config.rope
        if self.rope:
            # 如果使用RoPE，则初始化RoPE模块，维度为每个头的嵌入维度
            self.custom_rope = RotaryPositionalEmbeddings(self.n_embd // self.n_head)

    def forward(self, x):
        """
        前向传播，计算分组查询注意力输出。

        参数：
            x (torch.Tensor): 输入张量，形状为 (batch, seq_len, n_embd)。

        返回：
            torch.Tensor: 注意力输出，形状为 (batch, seq_len, n_embd)。
            int: 显存消耗。
        """
        # 获取输入的批次大小（B）、序列长度（T）和嵌入维度（C）
        B, T, C = x.size()

        # 通过线性层生成键和值，将嵌入维度分为两部分
        v, k = self.vk_attn(x).split(self.n_embd, dim=2)

        # 通过查询的线性层生成查询
        q = self.q_attn(x)

        # 将键和值 reshaped 为多头形式，维度为 (B, nh, T, d)，其中 nh 表示头的数量，d 表示每个头的维度
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, d)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, d)

        # 将查询 reshaped 为多头和分组形式，维度为 (B, nh, g, T, d)，g 表示查询的分组
        q = q.view(B, T, self.n_head, self.g, C // self.n_head).transpose(1, 2).transpose(2, 3)  # (B, nh, g, T, d)

        # 如果启用了 RoPE，则将旋转位置嵌入应用于查询和键
        if self.rope:
            q = self.custom_rope(q)
            k = self.custom_rope(k)

        # 清除缓存，确保准确计算显存消耗
        torch.cuda.empty_cache()
        start_memory = torch.cuda.memory_allocated()  # 获取开始时的显存消耗

        # 使用 einsum 计算注意力分数，维度为 (B, nh, g, T, T)
        att = torch.einsum('ijklm,ijmn->ijln', q, k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))

        # 应用因果遮罩，确保注意力只关注到当前位置及之前的序列
        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))

        # 使用 softmax 计算注意力权重，并应用 dropout 进行正则化
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)

        # 计算注意力输出，将注意力权重与值进行点积，得到 (B, nh, T, d)
        y = att @ v

        # 获取结束时的显存消耗
        end_memory = torch.cuda.memory_allocated()

        # 重塑注意力输出，将多头输出拼接，返回 (B, T, C)
        y = y.transpose(1, 2).contiguous().view(B, T, C)

        # 输出线性变换，并应用残差 dropout 进行正则化
        y = self.resid_dropout(self.c_proj(y))

        # 返回注意力输出和显存消耗差值
        return y, end_memory - start_memory

5.3 GQA问题

备注：下文5,6题我懒得去修改代码了，直接引用的CMU学生跑的结果图，第7问是我自己跑的结果图。

以下问题假设你使用的是绝对位置嵌入，而不是 RoPE。

5 （4 分）

绘制每次迭代中计算注意力所需的平均时间（毫秒）在不同K头数量 {1, 2, 3, 6} 下的变化。

6 （4 分）

绘制每次迭代的内存消耗（MB）在不同K头数量 {1, 2, 3, 6} 下的变化。

7 （4 分）

绘制你实现的 GQA 模型在 2 个K头下与原始（多头注意力）minGPT 在序列长度为 128 时训练 600 次迭代的训练损失图。

5.4 关于GQA跑实验的结论

从5.3可看出，使用QGA，降低K头数量，虽然计算时间降得很少，但是显存消耗降得非常多，并且训练速度或精度竟然还能提升！