大模型原理:实现GPT模型

介绍

在 大模型原理:构建GPT-2大语言模型 中,对GPT架构进行了部分实现,并将其称为 DummyGPTModel。虽然 DummyGPTModel 的代码实现中展示了GPT模型的输入和输出,但其包含的构建块仍然是一个“黑盒”,因此使用了 DummyTransformerBlock 类和 DummyLayerNorm 类作为占位

在本文中,把 DummyTransformerBlock 占位类和 DummyLayerNorm 占位类替换为之前编写的真正的TransformerBlock 类和 LayerNorm类 ,组装出一个完全可用且参数量为 1.24亿 的GPT-2模型

整体结果

在用代码构建GPT-2模型之前,先来看看它的整体结构,如下图所示:

执行流程

该图展示了GPT模型的数据流程:

  1. 从底部开始,词元化文本首先被转换成词元嵌入
  2. 然后用位置嵌入进行增强。
  3. 这些组合信息形成一个张量,然后通过中间所示的一系列 Transformer 块(每个块都包含多头注意力和前馈神经网络层,并带有dropout和层归一化功能),这些块相互堆叠并重复12次。参数量为 1.24亿 的GPT-2模型

可以通过GPT_CONFIG_124M 字典中的 n_layers 进行指定。

提示:在最大规模的GPT-2模型(参数量为 15.42亿)中,这个Transformer块被重复了48次。

输出流程

经过最后一步的 层归一化 处理,然后传递到 线性输出层。这个层会将Transformer的输出映射到一个高维空间(在本例中为 50 257维,对应模型的词汇表大小),以预测序列中的下一个词元。

GPT模型架构的实现

准备

从 大模型原理:全部代码 复制以下类或函数:

  • MultiHeadAttention
  • GELU
  • FeedForward
  • LayerNorm
  • TransformerBlock

创建 GPTModel 类

GPTModel 类的 __init__() 构造函数的解释:

  1. 通过字典 cfg 传递的配置,来初始化词元嵌入层和位置嵌入层。这些嵌入层负责将输入的词元索引转换为稠密向量,并添加位置信息。
  2. 接下来,__init__() 方法会创建一个 TransformerBlock 模块的顺序栈,其层数与cfg中指定的层数相同。
  3. Transformer块 之后会应用一个 LayerNorm 层,将 Transformer块 的输出标准化,以稳定学习过程。
  4. 最后,定义一个无偏置的线性输出头,将Transformer的输出投射到分词器的词汇空间,为词汇中的每个词元生成分数(logits)。

foward()方法解释:

  1. 首先接收一批输入的词元索引,然后计算它们的嵌入表示
  2. 接着应用位置嵌入,将序列通过一系列Transformer块传递,并对最终输出进行归一化处理。
  3. 最后一步是计算 logits,这些 logits 代表了下一个词元的非归一化概率
class GPTModel(nn.Module) :
    def __init__(self, cfg):
        super().__init__()
        # 词元嵌入层
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"] )
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        tbs = [TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
        self.trf_blocks = nn.Sequential(*tbs)
        self.final_norm = LayerNorm(cfg["emb_dim"])
        # 线性输出层
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        # device 的设置允许我们在CPU或GPU上训练模型,具体取决于输入数据所在的设备
        pos_embeds = self.pos_emb(
            torch.arange(seq_len, device=in_idx.device)
        )
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

执行

传入了 cfg 的 GPT_CONFIG_124M 初始化参数量为 1.24亿 的GPT模型,并向其输入之前创建的批次文本数据

GPT_CONFIG_124M = {
    "vocab_size": 50257,     # 词汇表大小
    "context_length": 1024,  # 上下文长度
    "emb_dim": 768,          # 嵌入维度
    "n_heads": 12,           # 注意力头的数量
    "n_layers": 12,          # 层数
    "drop_rate": 0.1,        # dropout率
    "qkv_bias": False        # 查询-键-值偏置
}


# 准备输入数据
import tiktoken

tokenizer = tiktoken.get_encoding("gpt2")
batch = []
txt1 = "Every effort moves you"
txt2 = "Every day holds a"

batch.append(torch.tensor(tokenizer.encode(txt1)))
batch.append(torch.tensor(tokenizer.encode(txt2)))
batch = torch.stack(batch, dim=0)
print("batch:\n", batch)

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)

out = model(batch)
print("Input batch:\n", batch)
print("\nOutput shape:", out.shape)
print(out)

执行结果:

batch:
 tensor([[6109, 3626, 6100,  345],  # 文本1的词元ID
        [6109, 1110, 6622,  257]])  # 文本2的词元ID
Input batch:
 tensor([[6109, 3626, 6100,  345],
        [6109, 1110, 6622,  257]])

Output shape: torch.Size([2, 4, 50257])
tensor([[[ 0.3613,  0.4222, -0.0711,  ...,  0.3483,  0.4661, -0.2838],
         [-0.1792, -0.5660, -0.9485,  ...,  0.0477,  0.5181, -0.3168],
         [ 0.7120,  0.0332,  0.1085,  ...,  0.1018, -0.4327, -0.2553],
         [-1.0076,  0.3418, -0.1190,  ...,  0.7195,  0.4023,  0.0532]],

        [[-0.2564,  0.0900,  0.0335,  ...,  0.2659,  0.4454, -0.6806],
         [ 0.1230,  0.3653, -0.2074,  ...,  0.7705,  0.2710,  0.2246],
         [ 1.0558,  1.0318, -0.2800,  ...,  0.6936,  0.3205, -0.3178],
         [-0.1565,  0.3926,  0.3288,  ...,  1.2630, -0.1858,  0.0388]]],
       grad_fn=<UnsafeViewBackward0>)

可以看到,输出张量的形状为 [2, 4, 50257],因为传入了2个输入文本,每个文本有4个词元。最后一个维度 50257 相当于分词器的词汇量。
在后面文章中,将这 50 257 维输出向量逐一转换回词元。

统计模型参数规模

通过 numel()(“number of elements”的缩写)方法可以统计模型参数张量的总参数量:

# 统计模型参数规模
print("\n\n统计模型参数规模--------")
total_params = sum(p.numel() for p in model.parameters())
print(f"模型参数总量: {total_params:,}")

执行结果:

模型参数总量: 163,009,536   # 1.63亿

提出问题

前面提到要初始化一个参数量为 1.24亿 的GPT模型,那为什么上面代码实际输出的参数量是 1.63亿 呢?

权重共享

原因在于原始GPT-2架构中使用了一个叫作 权重共享(weight tying)的概念。
也就是说,原始GPT-2架构是将 词元嵌入层 作为 输出层 重复使用的,即:词元嵌入层 使用两次,所以 其参数也算两次

查看GPTModel在模型上初始化的词元嵌入层和线性输出层的形状:

print("词元嵌入层 shape:", model.tok_emb.weight.shape)
print("线性输出层 shape:", model.out_head.weight.shape)

执行结果:

词元嵌入层 shape: torch.Size([50257, 768])
线性输出层 shape: torch.Size([50257, 768])

从打印输出可以看出,这两个层的权重张量具有相同的形状

减去重复计算的权重

根据权重共享的概念,需要从总的GPT-2模型参数计数中 减去 输出层的参数量:

total_params_gpt2 = (
    total_params - sum(p.numel()
    for p in model.out_head.parameters())
)
print(f"总的GPT-2模型参数减去输出层的参数量: {total_params_gpt2:,}"
)

执行结果:

总的GPT-2模型参数减去输出层的参数量: 124,412,160

现在模型只有 1.24亿 个参数,与GPT-2模型的原始大小相当

总结

权重共享可以 减少 模型的总体 内存占用 和 计算复杂度。
不过,使用 单独的词元嵌入层 和 输出层 可以获得更好的训练效果和模型性能。因此,本文在 GPTModel 实现中使用了 单独的层。
这一点在现代大语言模型中同样适用


原文出处:/show_1GW2eoJJVnpQ.html