Build a Large Language Model (From Scratch) · Chapter 4 中文译文

从零实现用于文本生成的 GPT 模型

译自 Build a Large Language Model (From Scratch) MEAP V08。本文保留章节结构、代码、图、表与公式标记;图像从原 PDF 页段抽取并嵌入。

第 4 章 从零实现用于生成文本的 GPT 模型

本章涵盖以下内容

  • 编写一个类似 GPT 的大型语言模型(LLM),它可以经过训练来生成类似人类的文本
  • 对层激活值进行归一化,以稳定神经网络训练
  • 在深度神经网络中添加快捷连接,以更有效地训练模型
  • 实现 transformer 块,以创建不同规模的 GPT 模型
  • 计算 GPT 模型的参数数量和存储需求

在上一章中,你学习并编写了多头注意力机制,这是 LLM 的核心组件之一。在本章中,我们现在将编写 LLM 的其他构建块,并把它们组装成一个类似 GPT 的模型;在下一章中,我们将训练该模型来生成类似人类的文本,如图 4.1 所示。

图 4.1
图 4.1 关于编写 LLM、在通用文本数据集上预训练 LLM、以及在带标签数据集上微调 LLM 这三个主要阶段的思维模型。本章重点实现 LLM 架构;我们将在下一章训练该架构。

图 4.1 中所指的 LLM 架构由多个构建块组成,我们将在本章逐一实现。在更详细地介绍各个组件之前,我们会先在下一节从顶层视角观察模型架构。

4.1 编写 LLM 架构

LLM,例如 GPT(Generative Pretrained Transformer,即生成式预训练 Transformer),是大型深度神经网络架构,设计目标是一次生成一个词(或 token)的新文本。不过,尽管它们规模庞大,模型架构并没有你想象的那么复杂,因为其中很多组件都会重复出现,后文我们会看到这一点。图 4.2 从顶层展示了一个类似 GPT 的 LLM,并突出显示其主要组件。

图 4.2
图 4.2 GPT 模型的思维模型。除嵌入层外,它由一个或多个 transformer 块组成,其中包含我们在上一章实现的带掩码多头注意力模块。

如图 4.2 所示,我们已经介绍过其中若干方面,例如输入 token 化和嵌入,以及带掩码多头注意力模块。本章的重点是实现 GPT 模型的核心结构,包括其 transformer 块;随后我们将在下一章训练它来生成类似人类的文本。

在前几章中,为了简单起见,我们使用了较小的嵌入维度,以确保概念和示例可以舒适地放在单页中。现在,在本章中,我们将扩展到一个小型 GPT-2 模型的规模,具体来说是最小版本,包含 \(124\) 百万个参数;这与 Radford 等人的论文 “Language Models are Unsupervised Multitask Learners” 中描述的版本相对应。请注意,虽然原始报告提到的是 \(117\) 百万个参数,但后来对此进行了更正。

第 6 章将重点介绍如何把预训练权重加载到我们的实现中,并将其适配到包含 \(345\)、\(762\) 和 \(1,542\) 百万个参数的更大 GPT-2 模型。在深度学习以及 GPT 这类 LLM 的语境中,“参数”一词指模型中可训练的权重。这些权重本质上是模型的内部变量,会在训练过程中被调整和优化,以最小化特定的损失函数。这个优化过程使模型能够从训练数据中学习。

例如,在一个由 \(2,048 \times 2,048\) 维权重矩阵(或张量)表示的神经网络层中,矩阵的每个元素都是一个参数。由于有 \(2,048\) 行和 \(2,048\) 列,该层中的参数总数就是 \(2,048\) 乘以 \(2,048\),等于 \(4,194,304\) 个参数。

我们通过下面的 Python 字典指定小型 GPT-2 模型的配置,并将在后续代码示例中使用它:

GPT_CONFIG_124M = {
     "vocab_size": 50257,     # Vocabulary size
     "context_length": 1024,         # Context length
     "emb_dim": 768,          # Embedding dimension
     "n_heads": 12,           # Number of attention heads
     "n_layers": 12,          # Number of layers
     "drop_rate": 0.1,        # Dropout rate
     "qkv_bias": False        # Query-Key-Value bias
}

GPT_CONFIG_124M 字典中,我们使用简洁的变量名以保持清晰,并避免代码行过长:

  • "vocab_size" 指词汇表中有 \(50,257\) 个词,与第 2 章中的 BPE tokenizer 使用的词汇表一致。
  • "context_length" 通过第 2 章讨论过的位置嵌入,表示模型能够处理的最大输入 token 数。
  • "emb_dim" 表示嵌入大小,会把每个 token 转换为一个 \(768\) 维向量。
  • "n_heads" 表示多头注意力机制中的注意力头数量,该机制已在第 3 章实现。
  • "n_layers" 指定模型中 transformer 块的数量,后续章节会进一步说明。
  • "drop_rate" 表示 dropout 机制的强度(\(0.1\) 意味着丢弃 \(10\%\) 的隐藏单元),用于防止过拟合,如第 3 章所述。
  • "qkv_bias" 决定是否在多头注意力中用于计算查询、键和值的 Linear 层里包含偏置向量。遵循现代 LLM 的惯例,我们一开始会禁用它;但在第 6 章把 OpenAI 的预训练 GPT-2 权重加载到模型中时,会重新讨论这一点。

使用上面的配置,我们将在本节先实现一个 GPT 占位符架构(DummyGPTModel),如图 4.3 所示。这将为我们提供一个整体视图,说明所有部分如何组合在一起,以及为了组装完整的 GPT 模型架构,我们还需要在后续章节编写哪些其他组件。

图 4.3
图 4.3 概述 GPT 架构编写顺序的思维模型。在本章中,我们会先从 GPT 主干开始,也就是一个占位符架构;然后再进入各个核心部件,最终把它们组装到 transformer 块中,形成最终的 GPT 架构。

图 4.3 中的编号方框展示了我们处理编写最终 GPT 架构所需各个概念的顺序。我们会从第 1 步开始,即一个占位符 GPT 主干,我们称之为 DummyGPTModel

代码清单 4.1 一个占位符 GPT 模型架构类
import torch
import torch.nn as nn


class DummyGPTModel(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"])
          self.trf_blocks = nn.Sequential(
               *[DummyTransformerBlock(cfg) for _ in range(cfg["n_layers"])]) #A
          self.final_norm = DummyLayerNorm(cfg["emb_dim"])                             #B
          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)
          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


class DummyTransformerBlock(nn.Module):                                                #C
     def __init__(self, cfg):
          super().__init__()


     def forward(self, x):                                                             #D
          return x


class DummyLayerNorm(nn.Module):                                                       #E
     def __init__(self, normalized_shape, eps=1e-5):                                   #F
          super().__init__()


     def forward(self, x):
          return x
  • #A 使用 TransformerBlock 的占位符
  • #B 使用 LayerNorm 的占位符
  • #C 一个简单的占位符类,稍后会被真正的 TransformerBlock 替换
  • #D 这个块什么也不做,只返回输入。
  • #E 一个简单的占位符类,稍后会被真正的 TransformerBlock 替换
  • #F 这里的参数只是为了模拟 LayerNorm 接口。

这段代码中的 DummyGPTModel 类使用 PyTorch 的神经网络模块(nn.Module)定义了一个类似 GPT 模型的简化版本。DummyGPTModel 类中的模型架构包括 token 嵌入和位置嵌入、dropout、一系列 transformer 块(DummyTransformerBlock)、最终的层归一化(DummyLayerNorm),以及一个线性输出层(out_head)。配置通过 Python 字典传入,例如我们前面创建的 GPT_CONFIG_124M 字典。

forward 方法描述了数据在模型中的流动方式:它为输入索引计算 token 嵌入和位置嵌入,应用 dropout,让数据经过 transformer 块,应用归一化,最后通过线性输出层产生 logits。

上面的代码已经可以运行;在本节稍后准备好输入数据后,我们会看到这一点。不过,现在请注意,在上面的代码中,我们为 transformer 块和层归一化使用了占位符(DummyLayerNormDummyTransformerBlock),它们将在后续章节中开发。

接下来,我们将准备输入数据并初始化一个新的 GPT 模型,以说明它的用法。基于第 2 章中我们编写 tokenizer 时看到的图示,图 4.4 从高层次概述了数据如何流入和流出 GPT 模型。

图 4.4
图 4.4 展示输入数据如何被 token 化、嵌入并送入 GPT 模型的整体概览。请注意,在前面编写的 DummyGPTClass 中,token 嵌入是在 GPT 模型内部处理的。在 LLM 中,嵌入后的输入 token 维度通常与输出维度匹配。这里的输出嵌入表示我们在第 3 章讨论过的上下文向量。

为了实现图 4.4 中所示的步骤,我们使用第 2 章介绍的 tiktoken tokenizer,对一个由两段文本输入组成的批次进行 token 化,作为 GPT 模型的输入:

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)

这两段文本得到的 token ID 如下:

tensor([[ 6109,       3626,    6100,     345],                                        #A
          [ 6109,     1110,    6622,     257]])
  • #A 第一行对应第一段文本,第二行对应第二段文本

接下来,我们初始化一个新的、包含 \(124\) 百万个参数的 DummyGPTModel 实例,并把 token 化后的批次输入其中:

torch.manual_seed(123)
model = DummyGPTModel(GPT_CONFIG_124M)
logits = model(batch)
print("Output shape:", logits.shape)
print(logits)

模型输出通常称为 logits,如下所示:

Output shape: torch.Size([2, 4, 50257])
tensor([[[-1.2034,        0.3201, -0.7130,       ..., -1.5548, -0.2390, -0.4667],
           [-0.1192,      0.4539, -0.4432,       ...,    0.2392,     1.3469,    1.2430],
           [ 0.5307,      1.6720, -0.4695,       ...,    1.1966,     0.0111,    0.5835],
           [ 0.0139,      1.6755, -0.3388,       ...,    1.1586, -0.0435, -1.0400]],


          [[-1.0908,      0.1798, -0.9484,       ..., -1.6047,       0.2439, -0.4530],
           [-0.7860,      0.5581, -0.0610,       ...,    0.4835, -0.0077,       1.6621],
           [ 0.3567,      1.2698, -0.6398,       ..., -0.0162, -0.1296,         0.3717],
           [-0.2407, -0.7349, -0.5102,           ...,    2.0057, -0.3694,       0.1814]]],
         grad_fn=<UnsafeViewBackward0>)

输出张量有两行,对应两个文本样本。每个文本样本由 \(4\) 个 token 组成;每个 token 是一个 \(50,257\) 维向量,这与 tokenizer 的词汇表大小一致。

嵌入有 \(50,257\) 个维度,是因为这些维度中的每一个都指向词汇表中的一个唯一 token。在本章末尾实现后处理代码时,我们会把这些 \(50,257\) 维向量转换回 token ID,然后再把它们解码成词。

现在,我们已经从顶层视角看过 GPT 架构及其输入和输出;在后续章节中,我们将编写各个占位符,首先是真正的层归一化类,它将替换前面代码中的 DummyLayerNorm

4.2 使用层归一化对激活进行归一化

训练包含许多层的深度神经网络有时会很有挑战性,原因包括梯度消失或梯度爆炸等问题。这些问题会导致训练动态不稳定,并使网络难以有效调整自身权重;这意味着学习过程很难为神经网络找到一组能最小化损失函数的参数(权重)。换句话说,网络很难把数据中的底层模式学习到足以做出准确预测或决策的程度。(如果你刚接触神经网络训练和梯度概念,可以在附录 A“PyTorch 简介”的 A.4 节“轻松理解自动微分”中找到这些概念的简要介绍。不过,阅读本书内容并不需要对梯度有深入的数学理解。)

在本节中,我们将实现层归一化,以提高神经网络训练的稳定性和效率。

层归一化背后的主要思想,是调整神经网络层的激活(输出),使其均值为 \(0\)、方差为 \(1\),也称为单位方差。这种调整会加快向有效权重的收敛,并确保训练过程一致、可靠。正如我们在上一节基于 DummyLayerNorm 占位符所看到的,在 GPT-2 和现代 transformer 架构中,层归一化通常应用在多头注意力模块之前和之后,以及最终输出层之前。

在用代码实现层归一化之前,图 4.5 先直观概述了层归一化的工作方式。

图 4.5
图 4.5 层归一化示意图,其中 5 个层输出(也称为激活)被归一化,使它们具有零均值和方差 \(1\)。

我们可以通过以下代码重现图 4.5 中显示的例子:这里实现了一个有 5 个输入和 6 个输出的神经网络层,并将其应用于两个输入样本:

torch.manual_seed(123)
batch_example = torch.randn(2, 5)                                                     #A
layer = nn.Sequential(nn.Linear(5, 6), nn.ReLU())
out = layer(batch_example)
print(out)

#A create 2 training examples with 5 dimensions (features) each

这会打印以下张量,其中第一行列出第一个输入的层输出,第二行列出第二行输入的层输出:

tensor([[0.2260, 0.3470, 0.0000, 0.2216, 0.0000, 0.0000],
          [0.2133, 0.2394, 0.0000, 0.5198, 0.3297, 0.0000]],
         grad_fn=<ReluBackward0>)

我们刚编写的神经网络层由一个 Linear 层后接一个非线性激活函数 ReLU 组成;ReLU 是 Rectified Linear Unit(修正线性单元)的缩写,是神经网络中的标准激活函数。如果你不熟悉 ReLU,它只是把负输入截断为 \(0\),从而确保某一层只输出正值,这也解释了为什么得到的层输出不包含任何负值。(注意,在 GPT 中我们会使用另一种更复杂的激活函数,下一节会介绍。)

在对这些输出应用层归一化之前,先检查均值和方差:

mean = out.mean(dim=-1, keepdim=True)
var = out.var(dim=-1, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)

输出如下:

Mean:
  tensor([[0.1324],
            [0.2170]], grad_fn=<MeanBackward1>)
Variance:
  tensor([[0.0231],
            [0.0398]], grad_fn=<VarBackward0>)

上面均值张量中的第一行包含第一个输入行的均值,第二个输出行包含第二个输入行的均值。

在均值或方差计算等操作中使用 keepdim=True,可以确保输出张量保留与输入张量相同的维数,即使该操作会沿着通过 dim 指定的维度缩减张量。例如,如果不使用 keepdim=True,返回的均值张量会是一个二维向量 [0.1324, 0.2170],而不是一个 \(2 \times 1\) 维矩阵 [[0.1324], [0.2170]]

dim 参数指定在张量中沿哪个维度计算统计量(这里是均值或方差),如图 4.6 所示。

图 4.6
图 4.6 计算张量均值时 dim 参数的示意图。例如,如果有一个维度为 [行, 列] 的 \(2D\) 张量(矩阵),使用 dim=0 会跨行执行操作(垂直方向,如底部所示),得到聚合每一列数据的输出。使用 dim=1dim=-1 会跨列执行操作(水平方向,如顶部所示),得到聚合每一行数据的输出。

如图 4.6 所解释的,对于 \(2D\) 张量(如矩阵),在均值或方差计算等操作中使用 dim=-1 与使用 dim=1 相同。这是因为 \(-1\) 指的是张量的最后一个维度,在 \(2D\) 张量中对应列。之后,当我们把层归一化加入 GPT 模型时,该模型会产生形状为 \([batch\_size, num\_tokens, embedding\_size]\) 的 \(3D\) 张量;我们仍然可以使用 dim=-1 沿最后一个维度做归一化,而不必把 dim=1 改为 dim=2

接下来,将层归一化应用到前面得到的层输出。这个操作包括减去均值,并除以方差的平方根(也称为标准差):

out_norm = (out - mean) / torch.sqrt(var)
mean = out_norm.mean(dim=-1, keepdim=True)
var = out_norm.var(dim=-1, keepdim=True)
print("Normalized layer outputs:\n", out_norm)
print("Mean:\n", mean)
print("Variance:\n", var)

从结果可以看到,归一化后的层输出现在也包含负值,它们的均值为零,方差为 \(1\):

Normalized layer outputs:
 tensor([[ 0.6159,    1.4126, -0.8719,    0.5872, -0.8719, -0.8719],
        [-0.0189,    0.1121, -1.0876,    1.5173,   0.5647, -1.0876]],
        grad_fn=<DivBackward0>)
Mean:
 tensor([[2.9802e-08],
        [3.9736e-08]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[1.],
        [1.]], grad_fn=<VarBackward0>)

注意,输出张量中的值 2.9802e-08 是 \(2.9802 \times 10^{-8}\) 的科学计数法,也就是十进制形式的 \(0.0000000298\)。这个值非常接近 \(0\),但由于计算机表示数字时使用有限精度,会积累微小的数值误差,因此它并不恰好等于 \(0\)。

为了提高可读性,我们还可以通过把 sci_mode 设置为 False,在打印张量值时关闭科学计数法:

torch.set_printoptions(sci_mode=False)
print("Mean:\n", mean)
print("Variance:\n", var)
Mean:
 tensor([[     0.0000],
        [    0.0000]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[1.],
        [1.]], grad_fn=<VarBackward0>)

到目前为止,在本节中,我们已经以逐步方式编写并应用了层归一化。现在把这个过程封装进一个 PyTorch 模块中,以便稍后在 GPT 模型中使用:

清单 4.2 一个层归一化类
class LayerNorm(nn.Module):
     def __init__(self, emb_dim):
          super().__init__()
          self.eps = 1e-5
          self.scale = nn.Parameter(torch.ones(emb_dim))
          self.shift = nn.Parameter(torch.zeros(emb_dim))


     def forward(self, x):
          mean = x.mean(dim=-1, keepdim=True)
          var = x.var(dim=-1, keepdim=True, unbiased=False)
          norm_x = (x - mean) / torch.sqrt(var + self.eps)
          return self.scale * norm_x + self.shift

这个特定的层归一化实现作用于输入张量 x 的最后一个维度,该维度表示嵌入维度(emb_dim)。变量 eps 是一个加到方差上的小常数(epsilon),用于在归一化期间防止除以零。scaleshift 是两个可训练参数(与输入具有相同维度);如果 LLM 判断这样做能改善模型在训练任务上的性能,就会在训练期间自动调整它们。这使模型能够学习最适合其所处理数据的缩放和平移。

现在让我们实际试用 LayerNorm 模块,并将它应用到批次输入上:

ln = LayerNorm(emb_dim=5)
out_ln = ln(batch_example)
mean = out_ln.mean(dim=-1, keepdim=True)
var = out_ln.var(dim=-1, unbiased=False, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)

从结果可以看到,层归一化代码按预期工作,会对两个输入各自的值进行归一化,使它们的均值为 \(0\)、方差为 \(1\):

Mean:
 tensor([[          -0.0000],
          [         0.0000]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[1.0000],
          [1.0000]], grad_fn=<VarBackward0>)

在本节中,我们介绍了实现 GPT 架构所需的一个构建块,如图 4.7 的心智模型所示。

图 4.7
图 4.7 一个心智模型,列出了本章中为组装 GPT 架构而实现的不同构建块。

在下一节中,我们将考察 GELU 激活函数;它是 LLM 中使用的激活函数之一,用来替代本节使用的传统 ReLU 函数。

4.3 使用 GELU 激活实现前馈网络

在本节中,我们将实现一个小型神经网络子模块,它会作为 LLM 中 transformer 块的一部分使用。我们首先实现 GELU 激活函数;它在这个神经网络子模块中起着关键作用。(关于在 PyTorch 中实现神经网络的更多信息,请参见附录 A 的 A.5 节“实现多层神经网络”。)

从历史上看,ReLU 激活函数因为简单且在各种神经网络架构中有效,被深度学习广泛使用。然而,在 LLM 中,除了传统 ReLU 之外,还会使用其他几种激活函数。两个值得注意的例子是 GELU(Gaussian Error Linear Unit,高斯误差线性单元)和 SwiGLU(Swish-Gated Linear Unit,Swish 门控线性单元)。

GELU 和 SwiGLU 是更复杂且更平滑的激活函数,分别结合了高斯和 sigmoid 门控线性单元。与更简单的 ReLU 不同,它们能为深度学习模型提供更好的性能。

GELU 激活函数可以用多种方式实现;其精确版本定义为 \(\mathrm{GELU}(x)=x \cdot \Phi(x)\),其中 \(\Phi(x)\) 是标准高斯分布的累积分布函数。不过在实践中,通常会实现一个计算成本更低的近似式(原始 GPT-2 模型也使用这个近似式进行训练):

\(\mathrm{GELU}(x) \approx 0.5x \left(1 + \tanh\left[\sqrt{2/\pi}(x + 0.044715x^3)\right]\right)\)

在代码中,我们可以把这个函数实现为如下 PyTorch 模块:

清单 4.3 GELU 激活函数的一个实现
class GELU(nn.Module):
     def __init__(self):
          super().__init__()


     def forward(self, x):
          return 0.5 * x * (1 + torch.tanh(
               torch.sqrt(torch.tensor(2.0 / torch.pi)) *
               (x + 0.044715 * torch.pow(x, 3))
          ))

接下来,为了了解这个 GELU 函数的形状,以及它与 ReLU 函数有何不同,让我们把这两个函数并排绘制出来:

import matplotlib.pyplot as plt
gelu, relu = GELU(), nn.ReLU()


x = torch.linspace(-3, 3, 100)                                                  #A
y_gelu, y_relu = gelu(x), relu(x)
plt.figure(figsize=(8, 3))
for i, (y, label) in enumerate(zip([y_gelu, y_relu], ["GELU", "ReLU"]), 1):
     plt.subplot(1, 2, i)
     plt.plot(x, y)
     plt.title(f"{label} activation function")
     plt.xlabel("x")
     plt.ylabel(f"{label}(x)")
     plt.grid(True)
plt.tight_layout()
plt.show()

#A Create 100 sample data points in the range -3 to 3

从图 4.8 中得到的图形可以看到,ReLU 是一个分段线性函数:如果输入为正,它就直接输出输入;否则输出零。GELU 是一个平滑的非线性函数,它近似 ReLU,但在负值处仍有非零梯度。

图 4.8
图 4.8 使用 matplotlib 绘制的 GELU 和 ReLU 曲线输出。x 轴表示函数输入,y 轴表示函数输出。

如图 4.8 所示,GELU 的平滑性可以在训练期间带来更好的优化性质,因为它允许对模型参数进行更细致的调整。相比之下,ReLU 在零点有一个尖角,这有时会让优化更加困难,尤其是在非常深或架构复杂的网络中。此外,与 ReLU 对任何负输入都输出零不同,GELU 会允许负值对应一个很小但非零的输出。这个特性意味着,在训练过程中,接收负输入的神经元仍然可以参与学习过程,尽管贡献程度低于正输入。

接下来,让我们使用 GELU 函数来实现小型神经网络模块 FeedForward,稍后我们将在 LLM 的 transformer 块中使用它:

清单 4.4 一个前馈神经网络模块
class FeedForward(nn.Module):
     def __init__(self, cfg):
          super().__init__()
          self.layers = nn.Sequential(
               nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
               GELU(),
               nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
          )


     def forward(self, x):
          return self.layers(x)

从前面的代码可以看到,FeedForward 模块是一个小型神经网络,由两个 Linear 层和一个 GELU 激活函数组成。在拥有 1.24 亿参数的 GPT 模型中,它会通过 GPT_CONFIG_124M 字典接收包含词元的输入批次,其中每个词元的嵌入大小为 \(768\),也就是 GPT_CONFIG_124M["emb_dim"] = 768

图 4.9 展示了当我们把一些输入传入这个小型前馈神经网络时,嵌入大小会如何在其内部被操作。

图 4.9
图 4.9 前馈神经网络各层之间连接的直观概述。需要注意的是,这个神经网络可以容纳可变的批次大小和输入词元数量。不过,每个词元的嵌入大小会在初始化权重时确定并固定。

按照图 4.9 中的例子,让我们用 \(768\) 的词元嵌入大小初始化一个新的 FeedForward 模块,并向它输入一个包含 2 个样本、每个样本 3 个词元的批次输入:

ffn = FeedForward(GPT_CONFIG_124M)
x = torch.rand(2, 3, 768)                                                       #A
out = ffn(x)
print(out.shape)

#A create sample input with batch dimension 2

可以看到,输出张量的形状与输入张量相同:

torch.Size([2, 3, 768])

本节实现的 FeedForward 模块在增强模型从数据中学习并泛化的能力方面起着关键作用。虽然该模块的输入维度和输出维度相同,但如图 4.10 所示,它会在内部通过第一个线性层把嵌入维度扩展到更高维空间。这个扩展后接一个非线性 GELU 激活,然后通过第二个线性变换收缩回原始维度。这样的设计允许探索更丰富的表示空间。

图 4.10
图 4.10 前馈神经网络中层输出扩展和收缩的示意图。首先,输入按 \(4\) 倍从 \(768\) 个值扩展到 \(3072\) 个值。然后,第二层把这 \(3072\) 个值压缩回 \(768\) 维表示。

此外,输入和输出维度的一致性简化了架构,使我们之后能够堆叠多个层,而不需要在层与层之间调整维度,从而让模型更具可扩展性。

如图 4.11 所示,我们现在已经实现了 LLM 的大部分构建块。

图 4.11
图 4.11 一个心智模型,展示了本章涵盖的主题,其中黑色对勾表示我们已经介绍过的内容。

在下一节中,我们将讲解插入到神经网络不同层之间的快捷连接概念;这些连接对于提升深度神经网络架构的训练性能很重要。

4.4 添加捷径连接

接下来,我们讨论捷径连接背后的概念,它也称为跳跃连接或残差连接。最初,捷径连接是在计算机视觉中的深度网络(具体来说,是残差网络)里提出的,用来缓解梯度消失这一挑战。梯度消失问题指的是:梯度(它们在训练期间指导权重更新)在逐层反向传播时会逐渐变小,使得有效训练较早的层变得困难,如图 4.12 所示。

图 4.12
图 4.12 对比一个由 \(5\) 层组成的深度神经网络在没有捷径连接(左)和有捷径连接(右)时的情况。捷径连接会把某一层的输入加到它的输出上,从而有效创建一条绕过某些层的替代路径。图 1.1 中所示的梯度表示每一层的平均绝对梯度,我们将在下面的代码示例中计算它。

如图 4.12 所示,捷径连接通过跳过一层或多层,为梯度流经网络创建了一条替代的、更短的路径;这是通过把某一层的输出加到后续某一层的输出上实现的。这就是这些连接也被称为跳跃连接的原因。它们在训练的反向传播过程中对保持梯度流动起着关键作用。

在下面的代码示例中,我们实现图 4.12 所示的神经网络,看看如何在 forward 方法中添加捷径连接:

代码清单 4.5 一个用于说明捷径连接的神经网络
class ExampleDeepNeuralNetwork(nn.Module):
     def __init__(self, layer_sizes, use_shortcut):
          super().__init__()
          self.use_shortcut = use_shortcut
          self.layers = nn.ModuleList([
               # Implement 5 layers
               nn.Sequential(nn.Linear(layer_sizes[0], layer_sizes[1]), GELU()),
               nn.Sequential(nn.Linear(layer_sizes[1], layer_sizes[2]), GELU()),
               nn.Sequential(nn.Linear(layer_sizes[2], layer_sizes[3]), GELU()),
               nn.Sequential(nn.Linear(layer_sizes[3], layer_sizes[4]), GELU()),
               nn.Sequential(nn.Linear(layer_sizes[4], layer_sizes[5]), GELU())
          ])


     def forward(self, x):
          for layer in self.layers:
               # Compute the output of the current layer
               layer_output = layer(x)
               # Check if shortcut can be applied
               if self.use_shortcut and x.shape == layer_output.shape:
                     x = x + layer_output
               else:
                     x = layer_output
          return x

这段代码实现了一个包含 \(5\) 层的深度神经网络,每一层都由一个 Linear 层和一个 GELU 激活函数组成。在前向传播中,我们迭代地将输入传过各层,并且如果 self.use_shortcut 属性被设为 True,就可选择性地添加图 4.12 中描绘的捷径连接。

我们先用这段代码初始化一个没有捷径连接的神经网络。这里,每一层都会被初始化为接受一个包含 \(3\) 个输入值的样本,并返回 \(3\) 个输出值。最后一层返回单个输出值:

layer_sizes = [3, 3, 3, 3, 3, 1]
sample_input = torch.tensor([[1., 0., -1.]])
torch.manual_seed(123) # specify random seed for the initial weights for reproducibility
model_without_shortcut = ExampleDeepNeuralNetwork(
     layer_sizes, use_shortcut=False
)

接下来,我们实现一个函数,用于计算模型反向传播过程中的梯度:

def print_gradients(model, x):
    # Forward pass
    output = model(x)
    target = torch.tensor([[0.]])


    # Calculate loss based on how close the target
    # and output are
    loss = nn.MSELoss()
    loss = loss(output, target)


    # Backward pass to calculate the gradients
    loss.backward()


    for name, param in model.named_parameters():
          if 'weight' in name:
               # Print the mean absolute gradient of the weights
               print(f"{name} has gradient mean of {param.grad.abs().mean().item()}")

在前面的代码中,我们指定了一个损失函数,用来计算模型输出与用户指定目标(这里为简单起见,目标值为 \(0\))之间有多接近。然后,当调用 loss.backward() 时,PyTorch 会为模型中的每一层计算损失梯度。我们可以通过 model.named_parameters() 迭代权重参数。假设某一层有一个 \(3 \times 3\) 的权重参数矩阵,那么这一层会有 \(3 \times 3\) 个梯度值;我们打印这些 \(3 \times 3\) 个梯度值的平均绝对梯度,从而为每一层得到单个梯度值,以便更容易地比较各层之间的梯度。

简而言之,.backward() 是 PyTorch 中一个便捷的方法,它可以计算模型训练期间所需的损失梯度,而不需要我们自己实现梯度计算的数学细节,因此让深度神经网络的使用更容易上手。如果你不熟悉梯度和神经网络训练的概念,我建议阅读附录 A 中的 A.4 节“让自动微分变得简单”和 A.7 节“一个典型的训练循环”。

现在,让我们使用 print_gradients 函数,并将它应用到没有跳跃连接的模型上:

print_gradients(model_without_shortcut, sample_input)

输出如下:

layers.0.0.weight has gradient mean of 0.00020173587836325169
layers.1.0.weight has gradient mean of 0.0001201116101583466
layers.2.0.weight has gradient mean of 0.0007152041653171182
layers.3.0.weight has gradient mean of 0.001398873864673078
layers.4.0.weight has gradient mean of 0.005049646366387606

print_gradients 函数的输出可以看到,当我们从最后一层(layers.4)推进到第一层(layers.0)时,梯度会变小,这种现象称为梯度消失问题。

现在,让我们实例化一个带有跳跃连接的模型,并看看它的表现如何:

torch.manual_seed(123)
model_with_shortcut = ExampleDeepNeuralNetwork(
     layer_sizes, use_shortcut=True
)
print_gradients(model_with_shortcut, sample_input)

输出如下:

layers.0.0.weight has gradient mean of 0.22169792652130127
layers.1.0.weight has gradient mean of 0.20694105327129364
layers.2.0.weight has gradient mean of 0.32896995544433594
layers.3.0.weight has gradient mean of 0.2665732502937317
layers.4.0.weight has gradient mean of 1.3258541822433472

可以看到,根据输出,最后一层(layers.4)的梯度仍然比其他层更大。不过,当我们朝第一层(layers.0)推进时,梯度值会趋于稳定,并不会缩小到几乎消失的微小值。

总之,捷径连接对于克服深度神经网络中梯度消失问题带来的限制非常重要。捷径连接是 LLM 等超大模型的核心构建块;在下一章训练 GPT 模型时,它们会通过确保各层之间一致的梯度流动,帮助实现更有效的训练。

介绍完捷径连接之后,在下一节中,我们会把前面介绍过的所有概念(层归一化、GELU 激活、前馈模块和捷径连接)连接到一个 transformer 块中;这是我们编写 GPT 架构代码所需的最后一个构建块。

4.5 在 transformer 块中连接注意力层和线性层

在本节中,我们将实现 transformer 块,它是 GPT 和其他 LLM 架构的基本构建块。在拥有 1.24 亿参数的 GPT-2 架构中,这个块会重复十几次;它结合了我们之前已经介绍过的多个概念:多头注意力、层归一化、dropout、前馈层以及 GELU 激活,如图 4.13 所示。然后,在下一节中,我们会把这个 transformer 块连接到 GPT 架构的其余部分。

图 4.13
图 4.13 transformer 块的示意图。图底部展示的是已经被嵌入为 \(768\) 维向量的输入 token。每一行对应一个 token 的向量表示。transformer 块的输出是与输入维度相同的向量,随后可以送入 LLM 中的后续层。

如图 4.13 所示,transformer 块组合了多个组件,包括第 3 章中的掩码多头注意力模块,以及我们在第 4.3 节中实现的 FeedForward 模块。

当 transformer 块处理输入序列时,序列中的每个元素(例如一个词或子词 token)都由固定大小的向量表示(在图 4.13 中是 \(768\) 维)。transformer 块内部的操作,包括多头注意力和前馈层,都被设计为以一种保持维度不变的方式转换这些向量。

其思想是,多头注意力块中的自注意力机制会识别并分析输入序列中各元素之间的关系。相比之下,前馈网络会在每个位置上单独修改数据。这种组合不仅使模型能够更细致地理解和处理输入,也增强了模型处理复杂数据模式的整体能力。

在代码中,我们可以按如下方式创建 TransformerBlock

代码清单 4.6 GPT 的 transformer 块组件
from previous_chapters import MultiHeadAttention


class TransformerBlock(nn.Module):
    def __init__(self, cfg):
         super().__init__()
         self.att = MultiHeadAttention(
              d_in=cfg["emb_dim"],
              d_out=cfg["emb_dim"],
              context_length=cfg["context_length"],
              num_heads=cfg["n_heads"],
              dropout=cfg["drop_rate"],
              qkv_bias=cfg["qkv_bias"])
         self.ff = FeedForward(cfg)
         self.norm1 = LayerNorm(cfg["emb_dim"])
         self.norm2 = LayerNorm(cfg["emb_dim"])
         self.drop_shortcut = nn.Dropout(cfg["drop_rate"])


    def forward(self, x):
                                                                               #A
         shortcut = x
         x = self.norm1(x)
         x = self.att(x)
         x = self.drop_shortcut(x)
         x = x + shortcut      # Add the original input back


         shortcut = x                                                          #B
         x = self.norm2(x)

          x = self.ff(x)
          x = self.drop_shortcut(x)
          x = x + shortcut                                                    #C
          return x
  • #A 注意力块的捷径连接
  • #B 前馈块的捷径连接
  • #C 将原始输入加回来

给出的代码在 PyTorch 中定义了一个 TransformerBlock 类,它包含一个多头注意力机制(MultiHeadAttention)和一个前馈网络(FeedForward),二者都基于给定的配置字典(cfg)进行配置,例如 GPT_CONFIG_124M

层归一化(LayerNorm)会在这两个组件之前应用,而 dropout 会在它们之后应用,用来对模型进行正则化并防止过拟合。这也称为 Pre-LayerNorm(前置层归一化)。更早的架构,例如原始 transformer 模型,则是在自注意力网络和前馈网络之后应用层归一化,称为 Post-LayerNorm(后置层归一化),这通常会导致更差的训练动态。

这个类还实现了前向传播,其中每个组件后面都会跟随一个捷径连接,将块的输入加到它的输出上。正如第 4.4 节所解释的,这一关键特性有助于梯度在训练期间流经网络,并改善深层模型的学习效果。

使用我们之前定义的 GPT_CONFIG_124M 字典,让我们实例化一个 transformer 块,并向它输入一些样本数据:

torch.manual_seed(123)
x = torch.rand(2, 4, 768)                                                        #A
block = TransformerBlock(GPT_CONFIG_124M)
output = block(x)


print("Input shape:", x.shape)
print("Output shape:", output.shape)

#A Create sample input of shape [batch_size, num_tokens, emb_dim]
  • #A 创建形状为 \([\mathrm{batch\_size}, \mathrm{num\_tokens}, \mathrm{emb\_dim}]\) 的样本输入

输出如下:

Input shape: torch.Size([2, 4, 768])
Output shape: torch.Size([2, 4, 768])

从代码输出可以看到,transformer 块在输出中保持了输入维度,这表明 transformer 架构在整个网络中处理数据序列时不会改变其形状。

在整个 transformer 块架构中保持形状不变并非偶然,而是其设计中的一个关键方面。这一设计使它能够有效应用于各种序列到序列任务,在这些任务中,每个输出向量都直接对应一个输入向量,从而保持一一对应关系。不过,如第 3 章所学,输出是一个上下文向量,它封装了来自整个输入序列的信息。这意味着,当序列经过 transformer 块时,虽然序列的物理维度(长度和特征大小)保持不变,但每个输出向量的内容会被重新编码,以整合来自整个输入序列的上下文信息。

通过本节实现的 transformer 块,我们现在已经拥有了图 4.14 所示的所有构建块,这些构建块是下一节实现 GPT 架构所需要的。

图 4.14
图 4.14 本章到目前为止我们已经实现的不同概念的心智模型。

如图 4.14 所示,transformer 块结合了层归一化、包含 GELU 激活的前馈网络,以及我们在本章前面已经介绍过的捷径连接。正如我们将在接下来的章节中看到的,这个 transformer 块将构成我们要实现的 GPT 架构的主要组件。

4.6 编写 GPT 模型代码

本章开头,我们从一个 GPT 架构的总体概览开始,并把它称为 DummyGPTModel。在这个 DummyGPTModel 代码实现中,我们展示了 GPT 模型的输入和输出,但它的构建块仍然是黑盒,使用 DummyTransformerBlockDummyLayerNorm 类作为占位符。

在本节中,我们现在用本章稍后编写的真正 TransformerBlockLayerNorm 类,替换 DummyTransformerBlockDummyLayerNorm 占位符,以组装一个完全可运行的原始 1.24 亿参数版本 GPT-2。在第 5 章中,我们将预训练一个 GPT-2 模型;在第 6 章中,我们将加载来自 OpenAI 的预训练权重。

在用代码组装 GPT-2 模型之前,我们先看一下图 4.15 中的整体结构,它结合了本章到目前为止介绍过的所有概念。

图 4.15
图 4.15 GPT 模型架构概览。本图展示了数据在 GPT 模型中的流动方式。从底部开始,分词后的文本首先被转换为 token 嵌入,然后再与位置嵌入相加。这些组合后的信息形成一个张量,它会传入图中央所示的一系列 transformer 块(每个块都包含多头注意力层和前馈神经网络层,并带有 dropout 和层归一化),这些 transformer 块相互堆叠并重复 \(12\) 次。

如图 4.15 所示,我们在第 4.5 节中编写的 transformer 块会在整个 GPT 模型架构中重复多次。对于 1.24 亿参数的 GPT-2 模型,它会重复 \(12\) 次,我们通过 GPT_CONFIG_124M 字典中的 "n_layers" 条目来指定这一点。对于最大的、拥有 15.42 亿参数的 GPT-2 模型,这个 transformer 块会重复 \(36\) 次。

如图 4.15 所示,最终 transformer 块的输出随后会经过最后一次层归一化步骤,然后到达线性输出层。该层把 transformer 的输出映射到一个高维空间(在这里是 \(50,257\) 维,对应模型的词表大小),用来预测序列中的下一个 token。

现在,让我们用代码实现图 4.15 中看到的架构:

代码清单 4.7 GPT 模型架构实现
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"])


            self.trf_blocks = nn.Sequential(
                *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])


            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)
                                                                                    #A
            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
  • #A 设备设置将允许我们根据输入数据所在的设备,在 CPU 或 GPU 上训练模型

得益于我们在第 4.5 节中实现的 TransformerBlock 类,GPTModel 类相对较小且紧凑。

这个 GPTModel 类的 __init__ 构造函数会使用通过 Python 字典 cfg 传入的配置来初始化 token 嵌入层和位置嵌入层。正如第 2 章所讨论的,这些嵌入层负责把输入 token 索引转换为稠密向量,并加入位置信息。

接下来,__init__ 方法会创建一个由 TransformerBlock 模块组成的顺序堆栈,其数量等于 cfg 中指定的层数。在 transformer 块之后,会应用一个 LayerNorm 层,对 transformer 块的输出进行标准化,以稳定学习过程。最后,定义一个不带偏置的线性输出头,它将 transformer 的输出投影到分词器的词表空间中,为词表中的每个 token 生成 logits。

forward 方法接收一批输入 token 索引,计算它们的嵌入,应用位置嵌入,将序列传过 transformer 块,对最终输出进行归一化,然后计算 logits;这些 logits 表示下一个 token 的未归一化概率。我们将在下一节中把这些 logits 转换为 token 和文本输出。

现在,让我们使用传入 cfg 参数的 GPT_CONFIG_124M 字典初始化这个 1.24 亿参数的 GPT 模型,并把本章开头创建的批量文本输入送入模型:

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


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

前面的代码会先打印输入批次的内容,然后打印输出张量:

Input batch:
 tensor([[ 6109,      3626,   6100,   345], # token IDs of text 1
          [ 6109,    1110,    6622,   257]]) # token IDs of text 2


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\) 个 token。最后一个维度 \(50,257\) 对应分词器的词表大小。在下一节中,我们将看到如何把这些 \(50,257\) 维输出向量中的每一个转换回 token。

在进入下一节并编写将模型输出转换为文本的函数之前,我们先花一点时间继续研究模型架构本身,并分析它的规模。

使用 numel() 方法(“number of elements”的缩写),我们可以统计模型参数张量中的参数总数:

total_params = sum(p.numel() for p in model.parameters())
print(f"Total number of parameters: {total_params:,}")

结果如下:

Total number of parameters: 163,009,536

现在,细心的读者可能会注意到一个差异。前面我们说初始化的是一个 1.24 亿参数的 GPT 模型,那么为什么前面代码输出中显示的实际参数数量是 \(1.63\) 亿呢?

原因在于原始 GPT-2 架构中使用了一个称为权重绑定的概念,也就是说,原始 GPT-2 架构会在输出层中复用 token 嵌入层的权重。为了理解这意味着什么,让我们看看前面通过 GPTModel 在模型上初始化的 token 嵌入层和线性输出层的形状:

print("Token embedding layer shape:", model.tok_emb.weight.shape)
print("Output layer shape:", model.out_head.weight.shape)

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

Token embedding layer shape: torch.Size([50257, 768])
Output layer shape: torch.Size([50257, 768])

由于分词器词表中 \(50,257\) 个条目对应的行数,token 嵌入层和输出层都非常大。根据权重绑定,我们从 GPT-2 模型总参数数量中移除输出层的参数数量:

total_params_gpt2 =    total_params - sum(p.numel() for p in model.out_head.parameters())
print(f"Number of trainable parameters considering weight tying: {total_params_gpt2:,}")

输出如下:

Number of trainable parameters considering weight tying: 124,412,160

可以看到,现在模型规模只有 1.24 亿参数,与原始 GPT-2 模型的大小相匹配。

权重绑定会降低模型的整体内存占用和计算复杂度。不过,根据我的经验,使用独立的 token 嵌入层和输出层会带来更好的训练效果和模型性能;因此,我们在 GPTModel 实现中使用独立的层。现代 LLM 也是如此。不过,在第 6 章中,当我们加载来自 OpenAI 的预训练权重时,会重新讨论并实现权重绑定这一概念。

最后,让我们计算 GPTModel 对象中 \(1.63\) 亿个参数所需的内存:

total_size_bytes = total_params * 4                                              #A
total_size_mb = total_size_bytes / (1024 * 1024)                                 #B
print(f"Total size of the model: {total_size_mb:.2f} MB")

#A Calculate the total size in bytes (assuming float32, 4 bytes per parameter)
#B Convert to megabytes
  • #A 计算总大小(单位为字节,假设使用 float32,即每个参数 \(4\) 字节)
  • #B 转换为兆字节

结果如下:

Total size of the model: 621.83 MB

总之,通过计算 GPTModel 对象中 \(1.63\) 亿个参数的内存需求,并假设每个参数都是占用 \(4\) 字节的 32 位浮点数,我们发现模型总大小达到 \(621.83\) MB,这说明即使是相对较小的 LLM,也需要相对较大的存储容量来容纳。

在本节中,我们实现了 GPTModel 架构,并看到它会输出形状为 \([\mathrm{batch\_size}, \mathrm{num\_tokens}, \mathrm{vocab\_size}]\) 的数值张量。在下一节中,我们将编写代码,把这些输出张量转换为文本。

4.7 生成文本

在本章最后一节中,我们将实现把 GPT 模型的张量输出转换回文本的代码。在开始之前,我们先简要回顾一下像 LLM 这样的生成式模型是如何一次生成一个词(或 token)的,如图 4.16 所示。

图 4.16
图 4.16 本图展示了 LLM 一次生成一个 token 的分步过程。从初始输入上下文(“Hello, I am”)开始,模型在每次迭代中预测后续 token,并将其追加到输入上下文中,用于下一轮预测。如图所示,第一次迭代添加 “a”,第二次添加 “model”,第三次添加 “ready”,逐步构建出句子。

图 4.16 从宏观层面说明了 GPT 模型在给定输入上下文(例如 “Hello, I am,”)时生成文本的分步过程。每次迭代都会让输入上下文变长,使模型能够生成连贯且符合上下文的文本。到第 6 次迭代时,模型已经构造出一个完整句子:“Hello, I am a model ready to help.”

在上一节中,我们看到当前的 GPTModel 实现会输出形状为 \([\mathrm{batch\_size}, \mathrm{num\_token}, \mathrm{vocab\_size}]\) 的张量。现在的问题是,GPT 模型如何从这些输出张量得到图 4.16 所示的生成文本?

GPT 模型从输出张量到生成文本的过程包含几个步骤,如图 4.17 所示。这些步骤包括解码输出张量、基于概率分布选择 token,以及把这些 token 转换为人类可读的文本。

图 4.17
图 4.17 本图通过展示 token 生成过程中的一次迭代,详细说明了 GPT 模型中文本生成的机制。该过程从把输入文本编码为 token ID 开始,这些 token ID 随后被送入 GPT 模型。模型的输出随后被转换回文本,并追加到原始输入文本后面。

图 4.17 详细展示的下一个 token 生成过程说明了一个单独步骤:GPT 模型在给定输入的情况下生成下一个 token。

在每一步中,模型都会输出一个矩阵,其中的向量表示潜在的下一个 token。我们提取与下一个 token 对应的向量,并通过 softmax 函数将其转换为概率分布。在包含所得概率分数的向量中,找到最大值所在的索引,这个索引会转换为 token ID。然后将这个 token ID 解码回文本,从而生成序列中的下一个 token。最后,将该 token 追加到先前输入之后,形成用于后续迭代的新输入序列。这个逐步过程使模型能够按顺序生成文本,从初始输入上下文开始构建连贯的短语和句子。

在实践中,我们会在许多次迭代中重复这一过程,例如前面图 4.16 所示,直到达到用户指定的生成 token 数量。

在代码中,我们可以按如下方式实现 token 生成过程:

代码清单 4.8 用于让 GPT 模型生成文本的函数
def generate_text_simple(model, idx, max_new_tokens, context_size): #A
     for _ in range(max_new_tokens):
          idx_cond = idx[:, -context_size:]                                           #B
          with torch.no_grad():
                logits = model(idx_cond)


          logits = logits[:, -1, :]                                                   #C
          probas = torch.softmax(logits, dim=-1)                                      #D
          idx_next = torch.argmax(probas, dim=-1, keepdim=True)                       #E
          idx = torch.cat((idx, idx_next), dim=1)                                     #F


     return idx
  • #A idx 是当前上下文中索引组成的 \((\mathrm{batch}, n_{\mathrm{tokens}})\) 数组。
  • #B 如果当前上下文超过了支持的上下文大小,就裁剪当前上下文。例如,如果 LLM 只支持 5 个 token,而上下文大小为 10,那么只使用最后 5 个 token 作为上下文。
  • #C 只关注最后一个时间步,因此 \((\mathrm{batch}, n_{\mathrm{token}}, \mathrm{vocab\_size})\) 变为 \((\mathrm{batch}, \mathrm{vocab\_size})\)。
  • #D probas 的形状为 \((\mathrm{batch}, \mathrm{vocab\_size})\)。
  • #E idx_next 的形状为 \((\mathrm{batch}, 1)\)。
  • #F 将采样得到的索引追加到运行中的序列,其中 idx 的形状为 \((\mathrm{batch}, n_{\mathrm{tokens}}+1)\)。

在前面的 generate_text_simple 函数代码中,我们使用 softmax 函数把 logits 转换为概率分布,然后通过 torch.argmax 找到值最大的位置。softmax 函数是单调的,也就是说,它在把输入转换为输出时会保持输入的顺序。因此在实践中,softmax 这一步是多余的,因为 softmax 输出张量中分数最高的位置与 logit 张量中的对应位置相同。换句话说,我们可以直接把 torch.argmax 函数应用到 logits 张量上,并得到完全相同的结果。不过,我们仍然写出了这一转换过程,是为了展示把 logits 转换为概率的完整流程;这也能带来额外的直观理解,例如模型会生成最可能的下一个 token,这称为贪婪解码。

在下一章实现 GPT 训练代码时,我们还会引入其他采样技术;这些技术会修改 softmax 输出,使模型不会总是选择最可能的 token,从而为生成文本引入变化性和创造性。

这种一次生成一个 token ID 并用 generate_text_simple 函数将其追加到上下文中的过程,在图 4.18 中有进一步说明。(每次迭代的 token ID 生成过程在图 4.17 中详细说明。)

图 4.18
图 4.18 本图展示了 token 预测循环的六次迭代:模型把一串初始 token ID 作为输入,预测下一个 token,并将该 token 追加到输入序列中,作为下一次迭代的输入。(为便于理解,token ID 也被翻译为对应文本。)

如图 4.18 所示,我们以迭代方式生成 token ID。例如,在第 1 次迭代中,模型获得与 “Hello , I am” 对应的 token,预测下一个 token(ID 为 257,即 “a”),并将其追加到输入中。重复这一过程,直到模型在六次迭代后生成完整句子 “Hello, I am a model ready to help.”。

现在,让我们按照图 4.18 所示,在实践中用 “Hello, I am” 这个上下文作为模型输入,试用 generate_text_simple 函数。

首先,我们把输入上下文编码为 token ID:

start_context = "Hello, I am"
encoded = tokenizer.encode(start_context)
print("encoded:", encoded)
encoded_tensor = torch.tensor(encoded).unsqueeze(0)                                  #A
print("encoded_tensor.shape:", encoded_tensor.shape)
  • #A 添加批次维度。

编码后的 ID 如下:

encoded: [15496, 11, 314, 716]
encoded_tensor.shape: torch.Size([1, 4])

接下来,我们把模型切换到 .eval() 模式,这会禁用 dropout 等随机组件;这些组件只在训练期间使用。然后在编码后的输入张量上使用 generate_text_simple 函数:

model.eval()                                                                     #A
out = generate_text_simple(
     model=model,
     idx=encoded_tensor,
     max_new_tokens=6,
     context_size=GPT_CONFIG_124M["context_length"]
)
print("Output:", out)
print("Output length:", len(out[0]))
  • #A 禁用 dropout,因为我们没有在训练模型。

得到的输出 token ID 如下:

Output: tensor([[15496,           11,    314,     716, 27018, 24086, 47843, 30961, 42348,   7267]])
Output length: 10

使用 tokenizer 的 .decode 方法,我们可以把 ID 转换回文本:

decoded_text = tokenizer.decode(out.squeeze(0).tolist())
print(decoded_text)

文本格式的模型输出如下:

Hello, I am Featureiman Byeswickattribute argue

可以看到,根据前面的输出,模型生成的是乱码,完全不像图 4.18 中展示的连贯文本。发生了什么?模型无法生成连贯文本的原因是我们还没有训练它。到目前为止,我们只是实现了 GPT 架构,并用初始随机权重初始化了一个 GPT 模型实例。

模型训练本身是一个很大的主题,我们将在下一章处理它。

4.8 小结

  • 层归一化通过确保每一层的输出具有一致的均值和方差来稳定训练。
  • 快捷连接是跳过一层或多层的连接方式,它会把某一层的输出直接馈送到更深的层;这有助于缓解训练深度神经网络(例如 LLM)时的梯度消失问题。
  • Transformer 块是 GPT 模型的核心结构组件,它把掩码多头注意力模块与使用 GELU 激活函数的全连接前馈网络结合起来。
  • GPT 模型是由许多重复 transformer 块组成的 LLM,拥有从数百万到数十亿不等的参数。
  • GPT 模型有不同规模,例如 1.24 亿、3.45 亿、7.62 亿和 15.42 亿参数;我们可以用同一个 GPTModel Python 类来实现这些规模。
  • GPT 类 LLM 的文本生成能力涉及把输出张量解码为人类可读的文本,方式是基于给定输入上下文一次顺序预测一个 token。
  • 未经训练时,GPT 模型会生成不连贯的文本,这凸显了模型训练对于连贯文本生成的重要性;这也是后续章节的主题。

章节范围:原 PDF 物理第 113-154 页。图片从同一页段抽取并嵌入本文。