本章内容包括:
- 编写一个类似 GPT 的大型语言模型(LLM),可以训练以生成类似人类的文本
- 规范化层激活以稳定神经网络训练
- 在深度神经网络中添加快捷连接
- 实现 transformer 块,以创建不同规模的 GPT 模型
- 计算 GPT 模型的参数数量和存储需求
您已经学习并编码了多头注意力机制,这是 LLM 的核心组件之一。现在,我们将编码 LLM 的其他构建块,并将它们组装成一个类似 GPT 的模型,我们将在下一章中训练该模型以生成类似人类的文本。
图 4.1 中引用的 LLM 架构由多个构建块组成。我们将首先从整体上看一下模型架构,然后再详细介绍各个组件。
编码 LLM 架构
像 GPT(生成预训练变换器)这样的 LLM 是大型深度神经网络架构,旨在逐字(或逐标记)生成新文本。然而,尽管它们的规模庞大,但模型架构并不像你想象的那样复杂,因为许多组件是重复的,正如我们稍后所见。图 4.2 展示了一个类似 GPT 的 LLM 的俯视图,突出显示了其主要组件。
我们已经涵盖了 LLM 架构的多个方面,例如输入的标记化和嵌入,以及掩蔽多头注意模块。现在,我们将实现 GPT 模型的核心结构,包括其变换器块,稍后我们将训练这些块以生成类似人类的文本。
之前,为了简单起见,我们使用了较小的嵌入维度,以确保概念和示例可以舒适地适应单页。现在,我们将扩展到小型 GPT-2 模型的大小,特别是124百万参数的最小版本,正如 Radford 等人所描述的《语言模型是无监督的多任务学习者》中所述(mng.bz/yoBq)。请注意,尽管原报告提到117百万参数,但后来已更正。在第六章中,我们将专注于将预训练权重加载到我们的实现中,并将其调整为更大参数的GPT-2 模型(345、762 和 1,542 百万参数)。
在深度学习和像 GPT 这样的 LLM 的上下文中,术语“参数”指的是模型的可训练权重。这些权重本质上是模型的内部变量,在训练过程中进行调整和优化,以最小化特定的损失函数。这种优化使模型能够从训练数据中学习。
例如,在由 2,048 × 2,048 维的权重矩阵(或张量)表示的神经网络层中,这个矩阵的每个元素都是一个参数。由于有 2,048 行和 2,048 列,因此该层的参数总数为 2,048 乘以 2,048,即 4,194,304 个参数。
GPT-2 与 GPT-3
请注意,我们专注于 GPT-2,因为 OpenAI 已将预训练模型的权重公开,这些权重将在第六章中加载到我们的实现中。GPT-3 在模型架构上与 GPT-2 基本相同,只是参数从 GPT-2 的 15 亿增加到 GPT-3 的 1750 亿,并且在更多数据上进行了训练。截至目前,GPT-3 的权重尚未公开。GPT-2 也是学习如何实现 LLM 的更好选择,因为它可以在单台笔记本电脑上运行,而 GPT-3 需要 GPU 集群进行训练和推理。根据 Lambda Labs 的说法(lambdalabs.com/),在单个V100 数据中心 GPU 上训练 GPT-3 需要 355 年,而在消费级 RTX 8000 GPU 上则需要 665 年。
我们通过以下 Python 字典指定小型 GPT-2 模型的配置,这将在后面的代码示例中使用:
GPT_CONFIG_124M = {
"vocab_size": 50257, # 词汇大小
"context_length": 1024, # 上下文长度
"emb_dim": 768, # 嵌入维度
"n_heads": 12, # 注意力头数量
"n_layers": 12, # 层数
"drop_rate": 0.1, # 随机失活率
"qkv_bias": False # 查询-键-值偏置
}
在 GPT_CONFIG_124M 字典中,我们使用简洁的变量名以提高清晰度,并避免代码行过长:
vocab_size指代由 BPE 分词器使用的 50,257 个单词的词汇(见第 2 章)。context_length表示模型通过位置嵌入能够处理的最大输入令牌数(见第 2 章)。emb_dim表示嵌入大小,将每个令牌转换为一个 768 维的向量。n_heads指示多头注意力机制中的注意力头数量(见第 3 章)。n_layers指定模型中的变换器块数量,我们将在接下来的讨论中涉及。drop_rate表示随机失活机制的强度(0.1 表示 10% 的隐藏单元随机失活),以防止过拟合(见第 3 章)。qkv_bias决定是否在多头注意力的线性层中为查询、键和值计算包含一个偏置向量。我们将最初禁用此项,遵循现代 LLM 的规范,但将在第 6 章中重新讨论,在那时我们将 OpenAI 的预训练 GPT-2 权重加载到模型中(见第 6 章)。
使用此配置,我们将实现一个 GPT 占位符架构(DummyGPTModel),如图 4.3 所示。这将为我们提供一个整体视图,了解一切如何结合在一起,以及我们需要编写其他组件来组装完整的 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( #1
*[DummyTransformerBlock(cfg) #1
for _ in range(cfg["n_layers"])] #1
) #1
self.final_norm = DummyLayerNorm(cfg["emb_dim"]) #2
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): #3
def __init__(self, cfg):
super().__init__()
def forward(self, x): #4
return x
class DummyLayerNorm(nn.Module): #5
def __init__(self, normalized_shape, eps=1e-5): #6
super().__init__()
def forward(self, x):
return x
- #1 使用占位符表示 TransformerBlock
- #2 使用占位符表示 LayerNorm
- #3 一个简单的占位符类,将来会被真正的 TransformerBlock 替代
- #4 该块不执行任何操作,只返回输入。
- #5 一个简单的占位符类,将来会被真正的 LayerNorm 替代
- #6 这里的参数只是为了模拟 LayerNorm 接口。
DummyGPTModel 类定义了一个使用 PyTorch 神经网络模块(nn.Module)的简化版 GPT 模型架构。该模型架构包括令牌和位置嵌入、随机失活、若干变换器块(DummyTransformerBlock)、最后的层归一化(DummyLayerNorm)以及线性输出层(out_head)。配置通过 Python 字典传入,例如我们之前创建的 GPT_CONFIG_124M 字典。
forward 方法描述了数据在模型中的流动:它为输入索引计算令牌和位置嵌入,应用随机失活,通过变换器块处理数据,应用归一化,最后通过线性输出层生成 logits。
列表 4.1 中的代码已经可以正常运行。不过,注意我们使用了占位符(DummyLayerNorm 和 DummyTransformerBlock)来表示变换器块和层归一化,稍后我们将开发它们。
接下来,我们将准备输入数据并初始化一个新的 GPT 模型,以说明其用法。在我们编码分词器的基础上(见第 2 章),现在让我们高层次地概述数据如何在 GPT 模型中流入和流出,如图 4.4 所示。
为了实现这些步骤,我们使用第 2 章中的 tiktoken 分词器对两个文本输入的批次进行分词,以供 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)
生成的两个文本的令牌 ID 如下:
tensor([[6109, 3626, 6100, 345], #1
[6109, 1110, 6622, 257]])
#1 第一行对应第一个文本,第二行对应第二个文本。
接下来,我们初始化一个新的 1.24 亿参数的 DummyGPTModel 实例,并将分词后的批次输入给它:
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>)
输出张量有两行,分别对应两个文本样本。每个文本样本包含四个令牌;每个令牌是一个 50,257 维的向量,大小与分词器的词汇量相匹配。
嵌入的维度为 50,257,因为这些维度中的每一个都对应词汇表中的一个唯一令牌。当我们实现后处理代码时,将把这些 50,257 维的向量转换回令牌 ID,然后可以将其解码为单词。
现在我们已经对 GPT 架构及其输入输出进行了总体概述,接下来我们将编码各个占位符,首先替换掉之前代码中的 DummyLayerNorm,实现真正的层归一化类。
使用层归一化规范化激活值
训练具有多个层的深度神经网络有时会面临挑战,例如梯度消失或梯度爆炸等问题。这些问题会导致训练动态不稳定,使网络难以有效调整其权重,这意味着学习过程在寻找一组参数(权重)以最小化损失函数时遇到困难。换句话说,网络在学习数据中的潜在模式时存在困难,无法达到准确预测或决策的程度。
注意 如果您对神经网络训练和梯度的概念不太熟悉,可以在附录 A 的 A.4 节找到这些概念的简要介绍。但是,深入理解梯度的数学知识并不是阅读本书内容所必需的。
现在,让我们实现层归一化,以提高神经网络训练的稳定性和效率。层归一化的主要思想是调整神经网络层的激活值(输出),使其均值为 0,方差为 1,也称为单位方差。这种调整加速了有效权重的收敛,并确保训练过程的一致性和可靠性。在 GPT-2 和现代变压器架构中,层归一化通常在多头注意模块之前和之后应用,并且,如我们在 DummyLayerNorm 占位符中看到的,也在最终输出层之前。图 4.5 提供了层归一化如何工作的视觉概述。
我们可以通过以下代码重现图 4.5 中的示例,其中我们实现一个具有五个输入和六个输出的神经网络层,并将其应用于两个输入示例:
torch.manual_seed(123)
batch_example = torch.randn(2, 5) #1
layer = nn.Sequential(nn.Linear(5, 6), nn.ReLU())
out = layer(batch_example)
print(out)
#1 创建两个具有五个维度(特征)的训练示例
这段代码输出以下张量,第一行列出了第一个输入的层输出,第二行列出了第二个输入的层输出:
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>)
我们编写的神经网络层由一个线性层和一个非线性激活函数 ReLU(整流线性单元)组成,ReLU 是神经网络中的标准激活函数。如果您对 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 × 1 的矩阵 [[0.1324], [0.2170]]。
dim 参数指定了在张量中计算统计量(在这里是均值或方差)时应沿着哪个维度进行计算。如图 4.6 所示,对于一个二维张量(如矩阵),在进行均值或方差计算时使用 dim=-1 与使用 dim=1 是相同的。这是因为 -1 指的是张量的最后一个维度,对应于二维张量中的列。稍后,在将层归一化添加到 GPT 模型时,该模型生成的张量形状为 [batch_size, num_tokens, embedding_size],我们仍然可以使用 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)
根据结果可以看到,归一化后的层输出现在也包含负值,均值为 0,方差为 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([[-5.9605e-08],
[1.9868e-08]], grad_fn=<MeanBackward1>)
Variance:
tensor([[1.],
[1.]], grad_fn=<VarBackward0>)
请注意,输出张量中的值 -5.9605e-08 是科学计数法,表示 -5.9605 × 10⁻⁸,换算为十进制形式为 -0.000000059605。这个值非常接近 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 模型中使用。
Listing 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),在归一化过程中添加到方差中以防止除以零。scale 和 shift 是两个可训练参数(与输入具有相同的维度),如果调整它们能够改善模型在训练任务上的表现,LLM 将在训练过程中自动调整它们。这使得模型能够学习适合其处理数据的适当缩放和偏移。
有偏方差
在我们的方差计算方法中,我们通过设置
unbiased=False来使用一个实现细节。对于那些好奇这意味着什么的人,在方差计算中,我们在方差公式中以输入数量n为分母。这种方法不应用贝塞尔修正,贝塞尔修正通常在样本方差估计中使用n - 1作为分母来调整偏差。这个选择导致了所谓的有偏方差估计。对于 LLMs 来说,嵌入维度n通常很大,使用n和n - 1之间的差异在实际应用中可以忽略不计。我选择这种方法是为了确保与 GPT-2 模型的归一化层兼容,并且它反映了用于实现原始 GPT-2 模型的 TensorFlow 的默认行为。使用类似的设置确保我们的方法与第六章中加载的预训练权重兼容。
让我们现在在实践中尝试 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 所示。接下来,我们将查看 GELU 激活函数,这是 LLMs 中使用的一种激活函数,而不是之前使用的传统 ReLU 函数。
层归一化与批归一化
如果你熟悉批归一化,这是一种常见且传统的神经网络归一化方法,你可能会想知道它与层归一化的比较。与批归一化不同,批归一化是在批次维度上进行归一化,而层归一化则是在特征维度上进行归一化。大型语言模型(LLMs)通常需要大量计算资源,且可用的硬件或特定使用案例可以决定训练或推理时的批次大小。由于层归一化独立于批次大小对每个输入进行归一化,因此在这些场景中提供了更多的灵活性和稳定性。这对于分布式训练或在资源受限的环境中部署模型尤其有益。
实现具有 GELU 激活的前馈网络
接下来,我们将实现一个小型神经网络子模块,该模块作为大型语言模型(LLM)中的变换器块的一部分。我们首先实现 GELU 激活函数,它在这个神经网络子模块中扮演着关键角色。
注意:有关在 PyTorch 中实现神经网络的更多信息,请参阅附录 A 的 A.5 节。
历史上,ReLU 激活函数因其简单性和在各种神经网络架构中的有效性而被广泛使用。然而,在 LLM 中,除了传统的 ReLU 外,还使用了其他几种激活函数。其中两个显著的例子是 GELU(高斯误差线性单元)和 SwiGLU(Swish 门控线性单元)。
GELU 和 SwiGLU 是更复杂且平滑的激活函数,分别结合了高斯和 sigmoid 门控线性单元。与简单的 ReLU 相比,它们为深度学习模型提供了更好的性能。
GELU 激活函数可以有多种实现方式;其确切版本定义为 GELU(x)=xΦ(x)\text{GELU}(x) = x \Phi(x)GELU(x)=xΦ(x),其中 Φ(x)\Phi(x)Φ(x) 是标准高斯分布的累积分布函数。然而,在实践中,常常实现一种计算成本更低的近似(原始的 GPT-2 模型也是使用这种通过曲线拟合找到的近似进行训练的):
在代码中,我们可以将这个函数实现为一个 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) #1
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()
#1 在范围 –3 到 3 中创建 100 个样本数据点
从图 4.8 的结果图中可以看出,ReLU(右侧)是一个分段线性函数,对于正输入直接输出该值;而对于负输入则输出零。GELU(左侧)是一个平滑的非线性函数,近似于 ReLU,但对于几乎所有负值(除了大约 x = –0.75)具有非零梯度。
GELU 的平滑特性可以在训练过程中带来更好的优化特性,因为它允许对模型参数进行更细致的调整。相比之下,ReLU 在零点处有一个锐角(图 4.18 右侧),这有时会使优化变得更加困难,尤其是在非常深或复杂的网络中。此外,与 ReLU 对任何负输入输出零不同,GELU 允许负值产生一个小的非零输出。这一特性意味着,在训练过程中,接收负输入的神经元仍然可以对学习过程产生贡献,尽管比正输入的贡献要小。
接下来,让我们使用 GELU 函数实现一个小型神经网络模块 FeedForward,稍后将在 LLM 的变换器块中使用。
代码示例 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 模块是一个小型神经网络,由两个线性层和一个 GELU 激活函数组成。在 124 百万参数的 GPT 模型中,它接收每个具有 768 维嵌入大小的输入批次,这一设置来自于 GPT_CONFIG_124M 字典,其中 GPT_CONFIG_124M["emb_dim"] = 768。图 4.9 展示了当我们传入一些输入时,嵌入大小在这个小型前馈神经网络中的处理方式。
接下来,我们将根据图 4.9 的示例,初始化一个新的 FeedForward 模块,令其具有 768 的令牌嵌入大小,并输入一个包含两个样本且每个样本有三个令牌的批次输入:
ffn = FeedForward(GPT_CONFIG_124M)
x = torch.rand(2, 3, 768) #1
out = ffn(x)
print(out.shape)
如我们所见,输出张量的形状与输入张量的形状相同:
torch.Size([2, 3, 768])
FeedForward 模块在增强模型学习和泛化数据的能力方面发挥了关键作用。尽管该模块的输入和输出维度相同,但它通过第一个线性层将嵌入维度扩展到更高维空间,如图 4.10 所示。这个扩展后是非线性的 GELU 激活,最后通过第二个线性变换收缩回原始维度。这种设计使得模型能够探索更丰富的表示空间。
此外,输入和输出维度的一致性简化了架构,使得可以堆叠多个层而无需调整它们之间的维度,从而使模型更具可扩展性。
如图 4.11 所示,我们现在已经实现了大多数 LLM 的构建模块。接下来,我们将讨论在神经网络不同层之间插入的快捷连接的概念,这对提高深度神经网络架构的训练性能非常重要。
添加快捷连接
让我们讨论一下快捷连接的概念,也称为跳过连接或残差连接。最初,快捷连接是在计算机视觉中的深度网络(特别是残差网络)中提出的,用于缓解梯度消失的问题。梯度消失问题是指梯度(在训练过程中指导权重更新的信号)在反向传播通过层时逐渐变小,导致难以有效训练前面的层。
让我们讨论一下快捷连接的概念,也称为跳过连接或残差连接。最初,快捷连接是在计算机视觉中的深度网络(特别是残差网络)中提出的,用于缓解梯度消失的问题。梯度消失问题是指梯度(在训练过程中指导权重更新的信号)在反向传播过程中逐渐变小,这使得有效训练前面的层变得困难。
图4.12显示,快捷连接通过跳过一个或多个层为梯度流动创造了一条替代的、更短的路径,这通过将一个层的输出添加到后续层的输出来实现。这就是这些连接被称为跳过连接的原因。它们在训练的反向传播过程中对于保持梯度流动起着至关重要的作用。
接下来,我们实现图4.12中的神经网络,以看看我们如何在前向传播方法中添加快捷连接。
Listing 4.5 一个神经网络示例,展示快捷连接
import torch
import torch.nn as nn
class ExampleDeepNeuralNetwork(nn.Module):
def __init__(self, layer_sizes, use_shortcut):
super().__init__()
self.use_shortcut = use_shortcut
self.layers = nn.ModuleList([ #1
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:
layer_output = layer(x) #2
if self.use_shortcut and x.shape == layer_output.shape: #3
x = x + layer_output
else:
x = layer_output
return x
在这段代码中,我们实现了一个具有五层的深度神经网络,每层由一个线性层和一个GELU激活函数组成。在前向传播中,我们逐层传递输入,并在self.use_shortcut属性设置为True时选择性地添加快捷连接。
接下来,我们初始化一个不带快捷连接的神经网络。每一层将被初始化,以接受具有三个输入值的示例并返回三个输出值。最后一层返回一个输出值:
layer_sizes = [3, 3, 3, 3, 3, 1]
sample_input = torch.tensor([[1., 0., -1.]])
torch.manual_seed(123) #1
model_without_shortcut = ExampleDeepNeuralNetwork(
layer_sizes, use_shortcut=False
)
接下来,我们实现一个计算模型反向传播中梯度的函数:
def print_gradients(model, x):
output = model(x) #1
target = torch.tensor([[0.]])
loss = nn.MSELoss()
loss = loss(output, target) #2
loss.backward() #3
for name, param in model.named_parameters():
if 'weight' in name:
print(f"{name} has gradient mean of {param.grad.abs().mean().item()}")
这段代码指定了一个损失函数,计算模型输出与用户指定目标(在这里为了简便设为0)之间的接近程度。然后,当调用loss.backward()时,PyTorch计算模型每层的损失梯度。我们可以通过model.named_parameters()遍历权重参数。
接下来,让我们使用print_gradients函数并将其应用于不带快捷连接的模型:
print_gradients(model_without_shortcut, sample_input)
输出结果显示,随着我们从最后一层(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.4)的梯度仍然大于其他层。然而,随着我们朝着第一层(layers.0)推进,梯度值稳定下来,并没有缩小到非常小的值。
总之,快捷连接对于克服深度神经网络中梯度消失问题的限制至关重要。快捷连接是非常大模型(如LLM)的核心构建块,它们将在我们在下一章训练GPT模型时促进更有效的训练。
接下来,我们将把之前讨论的所有概念(层归一化、GELU激活、前馈模块和快捷连接)结合在一起,形成一个变换器块,这是构建GPT架构所需的最后一个模块。
连接变换器块中的注意力层和线性层
现在,让我们实现变换器块,这是GPT和其他LLM架构的基本构建块。在124百万参数的GPT-2架构中,这个块重复了十多次,它结合了我们之前讨论的几个概念:多头注意力、层归一化、 dropout、前馈层和GELU激活。稍后,我们将把这个变换器块连接到GPT架构的其他部分。
图4.13显示了一个变换器块,结合了多个组件,包括掩码多头注意力模块(见第3章)和我们之前实现的前馈模块(见第4.3节)。当变换器块处理输入序列时,序列中的每个元素(例如,单词或子词标记)由一个固定大小的向量表示(在本例中为768维)。变换器块内的操作,包括多头注意力和前馈层,旨在以保留其维度的方式转换这些向量。
其核心思想是,多头注意力块中的自注意力机制识别和分析输入序列中元素之间的关系。相反,前馈网络在每个位置单独修改数据。这种结合不仅使得输入的理解和处理更为细致,也增强了模型处理复杂数据模式的整体能力。
我们可以在代码中创建变换器块。
Listing 4.6 GPT的变换器块组件
from chapter03 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):
shortcut = x
x = self.norm1(x)
x = self.att(x)
x = self.drop_shortcut(x)
x = x + shortcut
shortcut = x
x = self.norm2(x)
x = self.ff(x)
x = self.drop_shortcut(x)
x = x + shortcut
return x
在这段代码中,TransformerBlock类定义了一个在PyTorch中实现的变换器块,其中包括多头注意力机制(MultiHeadAttention)和前馈网络(FeedForward),这两个组件都是基于提供的配置字典(cfg)进行配置的,例如GPT_CONFIG_124M。
层归一化(LayerNorm)在这两个组件之前应用,而在它们之后应用dropout,以对模型进行正则化并防止过拟合。这也被称为预层归一化(Pre-LayerNorm)。较旧的架构,例如原始的变换器模型,通常在自注意力和前馈网络之后应用层归一化,称为后层归一化(Post-LayerNorm),这通常会导致训练动态不佳。
该类还实现了前向传递,其中每个组件后面都有一个快捷连接,将块的输入添加到其输出。这一关键特性有助于梯度在训练过程中通过网络流动,从而改善深度模型的学习(见4.4节)。
使用我们之前定义的GPT_CONFIG_124M字典,让我们实例化一个变换器块,并为其提供一些示例数据:
torch.manual_seed(123)
x = torch.rand(2, 4, 768) #1
block = TransformerBlock(GPT_CONFIG_124M)
output = block(x)
print("Input shape:", x.shape)
print("Output shape:", output.shape)
#1 创建形状为[batch_size, num_tokens, emb_dim]的示例输入
输出结果是:
Input shape: torch.Size([2, 4, 768])
Output shape: torch.Size([2, 4, 768])
正如我们所看到的,变换器块在输出中保持了输入的维度,这表明变换器架构在整个网络中处理数据序列时没有改变其形状。
在变换器块架构中保持形状不变并非偶然,而是其设计的关键方面。这种设计使其能够有效地应用于各种序列到序列的任务,其中每个输出向量直接对应于一个输入向量,保持一对一的关系。然而,输出是一个上下文向量,封装了整个输入序列的信息(见第3章)。这意味着,虽然序列的物理维度(长度和特征大小)在通过变换器块时保持不变,但每个输出向量的内容被重新编码,以整合整个输入序列的上下文信息。
随着变换器块的实现,我们现在拥有了构建GPT架构所需的所有构建块。正如图4.14所示,变换器块结合了层归一化、前馈网络、GELU激活和快捷连接。我们最终会看到,这个变换器块将构成GPT架构的主要组件。
编码GPT模型
我们在本章开始时对GPT架构进行了宏观概述,称之为DummyGPTModel。在这个DummyGPTModel的代码实现中,我们展示了GPT模型的输入和输出,但其构建块仍然是一个黑箱,使用DummyTransformerBlock和DummyLayerNorm类作为占位符。
现在,让我们用之前编码的真实TransformerBlock和LayerNorm类替换掉DummyTransformerBlock和DummyLayerNorm占位符,从而组装一个完全可用的原始124百万参数版本的GPT-2。在第5章中,我们将对GPT-2模型进行预训练,在第6章中,我们将加载来自OpenAI的预训练权重。
在我们编码组装GPT-2模型之前,先来看看其整体结构,如图4.15所示,其中包含了我们迄今为止覆盖的所有概念。可以看到,变换器块在GPT模型架构中重复了多次。在124百万参数的GPT-2模型中,它重复了12次,我们通过GPT_CONFIG_124M字典中的n_layers项来指定。对于最大规模的GPT-2模型(1,542百万参数),这个变换器块重复了48次。
最终变换器块的输出经过最后一层归一化步骤后,才到达线性输出层。该层将变换器的输出映射到高维空间(在此情况下为50,257维,对应于模型的词汇表大小),以预测序列中的下一个标记。
现在,让我们实现图4.15中的架构。
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)
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
由于TransformerBlock类的存在,GPTModel类相对简洁和紧凑。GPTModel类的__init__构造函数使用通过Python字典cfg传入的配置初始化了标记和位置嵌入层。这些嵌入层负责将输入的标记索引转换为密集向量,并添加位置信息(见第2章)。
接下来,__init__方法创建一个与cfg中指定层数相等的TransformerBlock模块的顺序堆叠。在变换器块之后,应用LayerNorm层,对变换器块的输出进行标准化,以稳定学习过程。最后,定义一个无偏置的线性输出头,将变换器的输出投影到标记器的词汇空间,以生成每个标记的logits。
forward方法接收一批输入标记索引,计算它们的嵌入,应用位置嵌入,将序列传递通过变换器块,归一化最终输出,然后计算logits,表示下一个标记的非归一化概率。我们将在下一节中将这些logits转换为标记和文本输出。
现在,让我们使用我们传入cfg参数的GPT_CONFIG_124M字典初始化124百万参数的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],
[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],因为我们传入了两个具有四个标记的输入文本。最后一个维度50257对应于标记器的词汇表大小。稍后,我们将看到如何将这些50,257维的输出向量转换回标记。
在我们继续编写将模型输出转换为文本的函数之前,让我们多花一点时间分析模型架构及其规模。使用numel()方法(元素数量的简称),我们可以收集模型参数张量中的参数总数:
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架构重用了标记嵌入层的权重到其输出层。为了更好地理解这一点,让我们查看我们在之前通过GPTModel初始化的标记嵌入层和线性输出层的形状:
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])
由于标记器词汇表中50257的行数,这两个嵌入层和输出层非常大。让我们根据权重绑定将输出层的参数计数从总GPT-2模型计数中去除:
total_params_gpt2 = (
total_params - sum(p.numel()
for p in model.out_head.parameters())
)
print(f"Number of trainable parameters "
f"considering weight tying: {total_params_gpt2:,}")
输出是:
Number of trainable parameters considering weight tying: 124,412,160
如我们所见,模型现在只有1.24亿参数,与原始GPT-2模型的大小相匹配。
权重绑定减少了模型的整体内存占用和计算复杂度。然而,根据我的经验,使用单独的标记嵌入和输出层可以带来更好的训练和模型性能;因此,我们在GPTModel实现中使用了单独的层。现代LLM也是如此。然而,当我们加载OpenAI的预训练权重时,我们将在第6章重新访问并实现权重绑定的概念。
练习 4.1 前馈和注意力模块中的参数数量
计算并比较前馈模块和多头注意力模块中包含的参数数量。
最后,让我们计算我们GPTModel对象中1.63亿参数的内存需求:
total_size_bytes = total_params * 4 #1
total_size_mb = total_size_bytes / (1024 * 1024) #2
print(f"Total size of the model: {total_size_mb:.2f} MB")
#1 计算总字节数(假设使用float32,每个参数占用4字节)
#2 转换为兆字节
结果是:
Total size of the model: 621.83 MB
总之,通过计算我们GPTModel对象中1.63亿参数的内存需求,并假设每个参数是占用4字节的32位浮点数,我们发现模型的总大小为621.83 MB,这说明了即使是相对较小的LLM也需要较大的存储容量。
现在,我们已经实现了GPTModel架构,并看到它输出形状为[batch_size, num_tokens, vocab_size]的数字张量,让我们编写代码将这些输出张量转换为文本。
练习 4.2 初始化更大的GPT模型
我们初始化了一个1.24亿参数的GPT模型,称为“GPT-2小型”。在不进行任何代码修改的情况下,只需更新配置文件,使用GPTModel类实现GPT-2中型(使用1024维嵌入、24个变换器块、16个多头注意力头)、GPT-2大型(1280维嵌入、36个变换器块、20个多头注意力头)和GPT-2 XL(1600维嵌入、48个变换器块、25个多头注意力头)。作为额外任务,计算每个GPT模型中的参数总数。
生成文本
现在,我们将实现将GPT模型的张量输出转换回文本的代码。在开始之前,让我们简要回顾一下像LLM这样的生成模型是如何一次生成一个单词(或标记)的。
图4.16展示了GPT模型在给定输入上下文(例如“你好,我是。”)时逐步生成文本的过程。随着每次迭代,输入上下文不断增长,使得模型能够生成连贯且符合上下文的文本。在第六次迭代时,模型构建了一个完整的句子:“你好,我是一个准备好帮助的模型。”我们看到当前的GPTModel实现输出形状为[batch_size, num_token, vocab_size]的张量。现在的问题是:GPT模型是如何从这些输出张量生成文本的?
GPT模型从输出张量到生成文本的过程涉及几个步骤,如图4.17所示。这些步骤包括解码输出张量、根据概率分布选择标记,以及将这些标记转换为人类可读的文本。
图4.17中详细描述的下一个标记生成过程展示了GPT模型在给定输入时生成下一个标记的单个步骤。在每一步中,模型输出一个矩阵,其中的向量表示潜在的下一个标记。对应下一个标记的向量被提取,并通过softmax函数转换为概率分布。在包含结果概率分数的向量中,找到最大值的索引,这个索引对应于标记ID。然后将这个标记ID解码回文本,生成序列中的下一个标记。最后,这个标记被附加到之前的输入中,形成下一次迭代的新输入序列。这一逐步过程使得模型能够顺序生成文本,从最初的输入上下文中构建连贯的短语和句子。
在实践中,我们会重复这个过程进行多次迭代,如图4.16所示,直到达到用户指定的生成标记数量。在代码中,我们可以按以下列表实现标记生成过程。
列表 4.8 GPT模型生成文本的函数
def generate_text_simple(model, idx, #1
max_new_tokens, context_size):
for _ in range(max_new_tokens):
idx_cond = idx[:, -context_size:] #2
with torch.no_grad():
logits = model(idx_cond)
logits = logits[:, -1, :] #3
probas = torch.softmax(logits, dim=-1) #4
idx_next = torch.argmax(probas, dim=-1, keepdim=True) #5
idx = torch.cat((idx, idx_next), dim=1) #6
return idx
#1 idx是当前上下文中标记索引的(批量,n_tokens)数组。
#2 如果当前上下文超过支持的上下文大小,则裁剪当前上下文,例如,如果LLM仅支持5个标记,且上下文大小为10,则仅使用最后5个标记作为上下文。
#3 仅关注最后一个时间步,因此(批量,n_token,vocab_size)变为(批量,vocab_size)。
#4 probas的形状为(批量,vocab_size)。
#5 idx_next的形状为(批量,1)。
#6 将采样索引附加到运行序列中,其中idx的形状为(批量,n_tokens+1)。
这段代码展示了使用PyTorch实现语言模型生成循环的简单实现。它针对要生成的新标记的指定数量进行迭代,裁剪当前上下文以适应模型的最大上下文大小,计算预测,然后根据最高概率预测选择下一个标记。
在编码generate_text_simple函数时,我们使用softmax函数将logits转换为概率分布,从中识别出具有最高值的位置,通过torch.argmax实现。softmax函数是单调的,这意味着它在转换为输出时保持输入的顺序。因此,实际上,softmax步骤是冗余的,因为softmax输出张量中最高得分的位置与logits张量中的位置相同。换句话说,我们可以直接对logits张量应用torch.argmax函数,并得到相同的结果。然而,我提供了转换代码以展示将logits转化为概率的完整过程,这可以提供额外的直觉,使得模型生成最可能的下一个标记,这被称为贪婪解码。
在下一章实现GPT训练代码时,我们将使用额外的采样技术来修改softmax输出,以便模型不总是选择最可能的标记。这引入了生成文本的变异性和创造性。
使用generate_text_simple函数逐个生成标记ID并将其附加到上下文中的过程在图4.18中进一步说明(每次迭代的标记ID生成过程详见图4.17)。我们以迭代方式生成标记ID。例如,在第一次迭代中,模型提供了对应“你好,我是”的标记,预测下一个标记(ID为257,即“a”),并将其附加到输入中。这个过程重复进行,直到模型在六次迭代后生成完整句子“你好,我是一个准备好帮助的模型”。
现在,让我们使用“你好,我是”上下文作为模型输入来尝试generate_text_simple函数。首先,我们将输入上下文编码为标记ID:
start_context = "Hello, I am"
encoded = tokenizer.encode(start_context)
print("encoded:", encoded)
encoded_tensor = torch.tensor(encoded).unsqueeze(0) #1
print("encoded_tensor.shape:", encoded_tensor.shape)
#1 添加批量维度
编码的ID是:
encoded: [15496, 11, 314, 716]
encoded_tensor.shape: torch.Size([1, 4])
接下来,我们将模型设置为.eval()模式。这会禁用随机组件,如dropout,这些组件仅在训练期间使用,然后在编码输入张量上使用generate_text_simple函数:
model.eval() #1
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]))
#1 禁用dropout,因为我们不在训练模型
生成的输出标记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
如我们所见,模型生成了无意义的文本,完全不像连贯的文本“你好,我是一个准备好帮助的模型”。发生了什么?模型无法生成连贯文本的原因是我们还没有对其进行训练。到目前为止,我们只实现了GPT架构,并使用初始随机权重初始化了一个GPT模型实例。模型训练本身是一个大主题,我们将在下一章中讨论。
练习 4.3 使用单独的dropout参数
在本章开始时,我们在GPT_CONFIG_124M字典中定义了一个全局的drop_rate设置,以便在GPTModel架构中的各个地方设置dropout率。请更改代码,为模型架构中的各个dropout层指定单独的dropout值。(提示:我们使用dropout层的三个不同地方分别是:嵌入层、快捷连接层和多头注意力模块。)
总结
- 层归一化通过确保每一层的输出具有一致的均值和方差来稳定训练。
- 快捷连接是跳过一个或多个层的连接,将一层的输出直接馈送到更深的层,这有助于缓解训练深度神经网络(如LLM)时的梯度消失问题。
- 变换器块是GPT模型的核心结构组件,结合了带有掩码的多头注意力模块和使用GELU激活函数的全连接前馈网络。
- GPT模型是具有许多重复变换器块的LLM,参数数量从数百万到数十亿不等。
- GPT模型有多种大小,例如1.24亿、3.45亿、7.62亿和15.42亿参数,我们可以使用相同的GPTModel Python类来实现。
- GPT类LLM的文本生成能力涉及通过根据给定的输入上下文顺序预测一个标记,将输出张量解码为人类可读的文本。
- 没有训练的GPT模型生成的是不连贯的文本,这强调了模型训练在生成连贯文本中的重要性。