Files
Build-A-Large-Language-Mode…/cn-Book/4.从零开始实现一个用于文本生成的 GPT 模型.md
T
2024-11-04 13:00:39 +08:00

405 lines
22 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
本章涵盖以下内容:
+ **编写一个类 GPT 的大语言模型(LLM),可以训练其生成类似人类的文本**
+ **对网络层的激活值进行归一化,以稳定神经网络的训练过程**
+ **在深度神经网络中添加快捷连接,以更高效地训练模型**
+ **通过实现 Transformer 模块来构建不同规模的 GPT 模型**
+ **计算 GPT 模型的参数数量和存储需求**
在上一章中,你学习并实现了多头注意力机制,这是大语言模型(LLM)的核心组件之一。本章将进一步实现 LLM 的其他构建模块,并将它们组装成一个类似 GPT 的模型。我们将在下一章中训练该模型,以生成类人文本,具体过程如图 4.1 所示。
图 4.1 展示了构建 LLM 的三个主要阶段的概念模型:在通用文本数据集上对 LLM 进行预训练,并在标注数据集上进行微调。本章重点在于实现 LLM 的架构,下一章中我们将对其进行训练。
<img src="../Image/chapter4/figure4.1.png" width="75%" />
大语言模型(LLM)架构(见图 4.1)由多个模块构成,我们将在本章中实现这些模块。接下来的部分中,我们将首先从整体视角介绍模型架构,然后详细讲解各个组件。
## 4.1 实现 LLM 的架构
LLM(如GPT,即生成式预训练 Transformer,Generative Pretrained Transformer)是一种大型深度神经网络架构,设计用于逐词(或逐 token)生成新文本。然而,尽管模型规模庞大,其结构却并没有想象中那么复杂,因为模型的许多组件是重复的(后文将对此展开说明)。图 4.2 展示了一个类 GPT 的 LLM 的整体视图,并突出了其主要组成部分。
<img src="../Image/chapter4/figure4.2.png" width="75%" />
如图 4.2 所示,我们已经在之前的章节中讲解过几个模块,如输入的分词和嵌入,以及掩码多头注意力模块。本章的重点是实现 GPT 模型的核心结构,包括其 Transformer 块。我们将在下一章对该模型进行训练,使其能够生成类人文本。
在前几章中,为了简单起见,我们使用了较小的嵌入维度,确保概念和示例能够更方便地展示在一页内。而在本章中,我们将逐步扩展模型规模,达到小型 GPT-2 模型的大小(拥有1.24 亿参数量的最小版本)。该模型在 Radford 等人的论文《Language Models are Unsupervised Multitask Learners.》中有详细描述。请注意,尽管最初的报告中提到参数量为 1.17 亿,但后来更正为 1.24 亿。
第 6 章将重点介绍如何将预训练权重加载到我们的实现中,并将其调整为更大的 GPT-2 模型版本(包括 3.45 亿、7.62 亿和 15.42 亿参数量规模)。在深度学习和 GPT 等大语言模型的背景下,‘参数’一词指的是模型的可训练权重。这些权重本质上是模型的内部变量,在训练过程中不断调整和优化,以最小化特定的损失函数。这种优化使得模型能够从训练数据中学习。
例如,在一个神经网络层中,其权重由一个 2,048 x 2,048 维的矩阵(或张量)表示,这个矩阵的每个元素都是一个参数。由于该矩阵有 2,048 行和 2,048 列,因此该层的总参数数量为 2,048 乘以 2,048,即 4,194,304 个参数。
> [!NOTE]
>
> **GPT-2 与 GPT-3 的比较**
>
> 我们之所以关注 GPT-2,是因为 OpenAI 已公开了其预训练模型的权重,这些权重将在第 6 章中加载到我们的实现中。GPT-3 的模型架构基本上与 GPT-2 相同,只是将参数规模从 GPT-2 的 15 亿增加到了 1750 亿,同时在更多的数据上进行了训练。截至本文撰写时,GPT-3 的权重尚未公开。对于学习如何实现LLM,GPT-2 是更好的选择,因为它可以在单台笔记本电脑上运行,而 GPT-3 的训练和推理则需要 GPU 集群。根据 Lambda Labs 的估算,在单块 V100 数据中心 GPU 上训练 GPT-3 需要 355 年,而在消费级的 RTX 8000 GPU 上则需要 665 年。
我们通过以下 Python 字典来定义小型 GPT-2 模型的配置,稍后将在代码示例中使用该配置:
```python
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`指的是第 2 章中 BPE 分词器使用的 50,257 个词汇的词表大小。
+ `context_length`表示模型所能处理的最大输入 token 数(在第 2 章介绍位置嵌入时讨论过)。
+ `emb_dim`表示嵌入维度,将每个 token 转换为 768 维的向量。
+ `n_layers`指定模型中 Transformer 块的层数,后续章节将对此详解。
+ `drop_rate`表示 dropout 机制的强度(例如,0.1 表示丢弃 10% 的隐藏单元),用于防止过拟合,具体内容请回顾第 3 章。
+ `qkv_bia 参数决定是否在多头注意力的查询、键和值的线性层中加入偏置向量。我们最初会禁用该选项,以遵循现代大语言模型的标准,之后在第 6 章加载 OpenAI 预训练的 GPT-2 权重时再重新考虑该设置。
使用上述配置,我们将从本章开始实现一个GPT占位架构(DummyGPTModel),如图4.3所示。这将为我们提供一个全局视图,了解所有组件如何组合在一起,以及在接下来的章节中需要编写哪些其他组件来组装完整的GPT模型架构。
<img src="../Image/chapter4/figure4.3.png" width="75%" />
图 4.3 中显示的编号框说明了我们编写最终 GPT 架构所需理解的各个概念的顺序。我们将从第 1 步开始,这是一个我们称之为 DummyGPTModel 的 GPT 占位架构:
```python
# Listing 4.1 A placeholder GPT model architecture class
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 一个简单的占位类,后续将被真正的 DummyLayerNorm 替换
#F 此处的参数仅用于模拟LayerNorm接口
```
此代码中的 DummyGPTModel 类使用 PyTorch 内置的神经网络模块(nn.Module)定义了一个简化版的类 GPT 模型。该类包括 token 嵌入、位置嵌入、dropout、多个 Transformer 块(DummyTransformerBlock)、最终的层归一化(DummyLayerNorm)以及线性输出层(out_head)。模型配置通过 Python 字典传入,稍后将传入我们之前创建的 GPT_CONFIG_124M 字典。
`forward`方法定义了数据在模型中的流动方式:计算输入索引的 token 嵌入和位置嵌入,应用 dropout,通过 transformer block 处理数据,应用归一化,最后通过线性输出层生成 logits。
上面的代码已经可以正常运行,不过需要先准备输入数据,在本节后面我们会看到运行效果。需要注意的是,目前代码中我们使用了 `DummyLayerNorm` 和 `DummyTransformerBlock` 作为 Transformer 块和层归一化的占位符,实际的实现会在后续部分详细介绍。
接下来,我们将准备输入数据并初始化一个新的 GPT 模型,以展示它的用法。基于第二章实现的分词器,图 4.4 展示了数据在 GPT 模型中流入和流出的整体流程概览。
<img src="../Image/chapter4/figure4.4.png" width="75%" />
根据图 4.4 的步骤,我们使用第 2 章介绍的 tiktoken 分词器对包含两个文本输入的批次进行分词,以供 GPT 模型使用:
```python
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 如下:
```python
tensor([[ 6109, 3626, 6100, 345], #A
[ 6109, 1110, 6622, 257]])
#A 第一行对应第一段文本,第二行对应第二段文本。
```
接下来,我们初始化一个拥有 1.24 亿参数的 DummyGPTModel 模型实例,并将分词后的数据批量输入到模型中:
```python
torch.manual_seed(123)
model = DummyGPTModel(GPT_CONFIG_124M)
logits = model(batch)
print("Output shape:", logits.shape)
print(logits)
```
模型输出(通常称为logits)如下:
```python
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 维的向量,维度大小与分词器的词汇表相同。
嵌入层的维度为 50,257,因为每个维度对应词汇表中的一个唯一 token。在之后的处理中,我们会将这些 50,257 维向量转换回 token ID,然后再解码成单词。
在对 GPT 架构及其输入输出进行了整体概览之后,接下来的章节中将编写各个占位模块的实现,首先从用真实的层归一化类替换之前代码中的 DummyLayerNorm 开始。
## 4.2 使用层归一化对激活值进行标准化
在训练深层神经网络时,梯度消失或梯度爆炸问题有时会带来挑战。这些问题会导致训练过程不稳定,使得网络难以有效调整权重,也就是说,模型难以找到一组能最小化损失函数的参数。换句话说,模型很难从数据中学习到足够准确的模式,以支持其做出准确的预测或决策。(如果您对神经网络训练或梯度概念不熟悉,可参考附录 A 的 A.4 节《自动微分入门》。但要理解本书内容,不需要对梯度概念有深刻的理解。)
> [!TIP]
>
> **个人思考:** 虽然对文本内容的理解并不需要深度掌握梯度的概念,但如果我们在学习过程中能习惯去发散,必然能帮助我们对所学知识理解的更深刻,下面我们就来聊一下梯度。
>
> 梯度本质上是一个**变化率**,描述了某个值(例如函数输出值)对另一个值(如输入变量)的变化趋势。简单来说,梯度告诉我们在当前位置上,朝哪个方向移动能让某个目标值增加或减少得更快。
>
> 举例:山坡上的爬山者
>
> 假设你站在一座山的某个位置,想要找到最快下山的路线。你会怎么做呢?首先你会注意到山坡的倾斜度(也就是梯度),倾斜越陡的地方,就意味着朝这个方向走可以让你更快地下降海拔。
>
> 在这个例子中:
>
> + **你的当前位置**代表模型当前的参数值。
> + **山坡的倾斜度**就是梯度,表示你在当前位置向下走的快慢和方向。
> + **往斜坡最陡的方向走**相当于使用梯度更新模型参数,使得海拔(也就是损失值)尽快下降。
>
> 而大模型在应用梯度的概念时,首先会设计一个损失函数,用来衡量模型的预测结果与目标结果的差距。在训练过程中,它通过梯度去帮助每个模型参数不断调整来快速减少损失函数的值,从而提高模型的预测精度。
本节中,我们将实现层归一化,以提高神经网络训练的稳定性和效率。
归一化的核心思想是将神经网络层的激活(输出)调整为均值为 0,方差为 1(即单位方差)。这种调整可以加速权重的收敛速度,确保训练过程的一致性和稳定性。正如上一节提到的,在 GPT-2 和现代 Transformer 架构中,层归一化通常应用于多头注意力模块的前后以及最终输出层之前。
在我们用代码实现层归一化之前,先通过图 4.5 了解一下层归一化的工作原理。
<img src="../Image/chapter4/figure4.5.png" width="75%" />
我们可以通过以下代码重现图 4.5 中的示例,其中实现了一个具有 5 个输入和 6 个输出的神经网络层,并将其应用于两个输入样本:
```python
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 创建2个训练样本,每个样本有5个维度(特征)
```
打印出的张量中,第一行表示第一个输入样本的网络层输出,第二行表示第二个输入样本的网络层输出:
```python
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,只需了解它的作用是将负值设为 0,确保输出层中没有负值。在 GPT 中,我们将使用另一种更复杂的激活函数,后续章节会介绍。
在对这些输出应用层归一化之前,我们先查看其均值和方差:
```python
mean = out.mean(dim=-1, keepdim=True)
var = out.var(dim=-1, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)
```
输出如下:
```python
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]`,而使用 `keepdim=True` 后,返回的张量则会是一个 `2×1` 的矩阵 `[[0.1324], [0.2170]]`。
`dim` 参数用于指定张量中进行统计计算(如均值或方差)的维度,具体如图 4.6 所示。
<img src="../Image/chapter4/figure4.6.png" width="75%" />
如图 4.6 所示,对于二维张量(如矩阵),在进行均值或方差计算等操作时,使用 `dim=-1` 等同于使用 `dim=1`,因为 `-1` 指的是张量的最后一个维度,即二维张量中的列。在后续对 GPT 模型加入层归一化时,模型会生成形状为 `[batch_size, num_tokens, embedding_size]` 的三维张量,我们依然可以使用 `dim=-1` 对最后一个维度进行归一化,而无需将 `dim=1` 改为 `dim=2`。
接下来,我们将对之前获得的层输出应用层归一化。该操作包括减去均值,并除以方差的平方根(即标准差):
```python
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:
```python
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 × 10^-8`的科学记数法表示,用十进制形式表示为`0.0000000298`。这个值虽然非常接近 0,但由于计算机表示数字的精度有限,会产生微小的数值误差,因此不完全等于 0。
为提高可读性,我们可以将 sci_mode 设置为 False,从而关闭张量值的科学计数法显示模式:
```python
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 模型中使用。
```python
# Listing 4.2 A layer normalization class
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)在训练中会自动调整这些参数,以改善模型在训练任务上的性能。这使得模型能够学习适合数据处理的最佳缩放和偏移方式。
> [!NOTE]
>
> **有偏方差**
>
> 我们在方差计算方法中选择了设置 `unbiased=False`。对于好奇其含义的读者,可以理解为我们在方差公式中直接用样本数 n 作为分母,不使用贝塞尔校正(通常分母使用 n−1 以校正样本方差估计中的偏差)。这种决定会导致所谓的有偏方差估计。对于大语言模型(LLM)来说,其嵌入维度 n 通常非常大,因此使用 n 和 n−1 的差异实际上可以忽略不计。我们选择这种方式是为了确保与 GPT-2 模型的归一化层兼容,并保持与 TensorFlow 的默认行为一致,后者用于实现最初的 GPT-2 模型。这种设置确保我们的方法与第 6 章中将加载的预训练权重兼容。
现在让我们在实践中尝试LayerNorm模块并将其应用于批量输入:
```python
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:
```python
Mean:
tensor([[ -0.0000],
[ 0.0000]], grad_fn=<MeanBackward1>)
Variance:
tensor([[1.0000],
[1.0000]], grad_fn=<VarBackward0>)
```
在本节中,我们介绍了实现 GPT 架构所需的一个基础模块,如图 4.7 所示。
<img src="../Image/chapter4/figure4.7.png" width="75%" />
在下一节中,我们将探讨大语言模型中使用的 GELU 激活函数,它将替代我们在本节使用的传统 ReLU 函数。
> [!NOTE]
>
> **层归一化与批量归一化的区别**
>
> 如果你熟悉批量归一化这种常见的传统神经网络归一化方法,可能会好奇它与层归一化的区别。与在数量维度上进行归一化的批量归一化不同,层归一化是在特征维度上进行归一化。LLM 通常需要大量计算资源,而可用的硬件资源或特定的使用场景可能会限制训练或推理过程中的批量大小。由于层归一化对每个输入的处理不依赖批量大小,因此在这些场景下提供了更高的灵活性和稳定性。这对于分布式训练或资源受限的环境中部署模型尤其有利。
## 4.3 实现带有 GELU 激活函数的前馈神经网络
在本节中,我们将实现一个小型神经网络子模块,作为 LLM 架构中的 Transformer 块的一部分。我们首先实现 GELU 激活函数,它将在这个神经网络子模块中起着至关重要的作用。(关于在 PyTorch 中实现神经网络的更多信息,请参考附录 A 的 A.5 节:实现多层神经网络)
过去,ReLU 激活函数因其简单且有效,常用于各种神经网络架构中。但在大语言模型中,除了传统的 ReLU,还使用了其他几种激活函数,其中两个典型的例子是 GELU(高斯误差线性单元)和 SwiGLU(Swish 门控线性单元)。
GELU 和 SwiGLU 是更复杂、平滑的激活函数,分别结合了高斯分布和 sigmoid 门控线性单元。与较简单的 ReLU 不同,这些激活函数能为深度学习模型提供更好的性能。
GELU 激活函数可以通过多种方式实现,其确切版本定义为 `GELU(x) = x ⋅ Φ(x)`,其中 Φ(x) 是标准正态分布的累积分布函数。然而在实践中,通常会采用计算开销更低的近似实现(最初的 GPT-2 模型也是用这种近似实现进行训练的):
$$ \text{GELU}(x) \approx 0.5 \cdot x \cdot\left(1+\tanh \left[\sqrt{(2 / \pi)} \cdot\left(x+0.044715 \cdot x^{3}\right]\right)\right. $$