Build-A-Large-Language-Model-CN/cn-Book/4.从零开始实现一个用于文本生成的 GPT 模型.md

本章涵盖以下内容：

+ **编写一个类 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. $$

我们可以编码将该函数实现为一个 PyTorch 模块，如下所示：

```python
# Listing 4.3 An implementation of the GELU activation function
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 函数进行对比，我们将这两个函数并排绘制：

```python
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 在 -3 到 3 的范围内生成 100 个样本数据点
```

如图 4.8 所示，ReLU 是一个分段线性函数，输入为正时输出输入值本身，否则输出零。而 GELU 是一种平滑的非线性函数，它近似于 ReLU，但在负值上也具有非零梯度。

<img src="../Image/chapter4/figure4.8.png" width="75%" />

如图 4.8 所示，GELU 的平滑性使其在训练过程中具有更好的优化特性，能够对模型参数进行更细微的调整。相比之下，ReLU 在零点处有一个拐角，这在网络深度较大或结构复杂时可能会增加优化难度。此外，与 ReLU 不同，ReLU 对所有负输入的输出为零，而 GELU 对负值允许一个小的非零输出。这意味着在训练过程中，接收负输入的神经元也能对学习过程产生一定的贡献，尽管贡献程度不及正输入。

接下来让我们使用 GELU 激活函数实现一个小型的神经网络模块 FeedForward，该模块稍后会应用在 LLM 的 transformer 模块中：

```python
# Listing 4.4 A feed forward neural network module
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 激活函数组成。在 1.24 亿参数的 GPT 模型中，该模块可以接收批量输入，每个输入 token 是一个 768 维的向量表示。这一嵌入维度大小通过 `GPT_CONFIG_124M` 配置字典中的 `GPT_CONFIG_124M["emb_dim"]` 参数指定。

图 4.9 展示了当我们输入数据后，这个前馈网络内部如何调整嵌入维度。

<img src="../Image/chapter4/figure4.9.png" width="75%" />

按照图 4.9 中的示例，我们初始化一个新的 FeedForward 模块，设置 token 嵌入维度为 768，并输入一个包含 2 个样本且每个样本有 3 个 token 的数据集：

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

#A 创建一个 batch 大小为 2 的示例输入
```

显然，输出张量的形状与输入张量相同：

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

我们在本节实现的 FeedForward 模块对模型能力的增强（主要体现在从数据中学习并泛化方面）起到了关键作用。尽管该模块的输入和输出维度相同，但在内部，它首先通过第一个线性层将嵌入维度扩展到一个更高维度的空间（如图 4.10 所示）。之后再接入非线性 GELU 激活，最后再通过第二个线性层变换回原始维度。这样的设计能够探索更丰富的表示空间。

> [!TIP]
>
> **个人思考：** 这段描述一笔带过了扩展和收缩嵌入维度为模型训练带来的好处，那到底该如何理解这样的设计能够探索更丰富的表示空间呢？
>
> 可以将扩展和收缩的过程类比为一种**数据解压缩与重新压缩**的机制：
>
> + **扩展（解压缩）**：假设我们有一段压缩的音乐文件（例如 MP3），里面包含了音频的基本信息。通过解压缩（扩展），我们把这个文件变成了一个更高质量的音频格式，允许我们看到（听到）更多的细节，比如乐器的细微声响和音调变化。
> + **特征提取**：接着，我们可以在这个高质量的音频文件中应用各种音频处理算法（相当于非线性激活），分析出更多细节，比如每种乐器的声音特点。
> + **收缩（压缩）**：最后，我们将音频再次压缩为一种更适合传输和存储的格式。虽然最终文件变小了，但这个文件已经包含了之前提取出的更多的声音细节。
>
> 将这种理解再应用到神经网络中，扩展后的高维空间可以让模型“看到”输入数据中更多的隐藏特征，提取出更丰富的信息。然后在收缩回低维度时，这些丰富的特征被整合到了输入的原始维度表示中，使模型最终的输出包含更多的上下文和信息。

<img src="../Image/chapter4/figure4.10.png" width="75%" />



此外，输入输出维度保持一致也有助于简化架构，方便堆叠多层（在后续的章节实现），无需调整各层维度，从而提升了模型的可扩展性。

如图4.11所示，我们目前已经实现了LLM的大部分组成模块。

<img src="../Image/chapter4/figure4.11.png" width="75%" />

下一节，我们将介绍“快捷连接”的概念，即在神经网络的不同层之间插入的连接结构，它对于提升深度神经网络架构的训练性能非常重要。

  

## 4.4 添加快捷连接

接下来，我们来讨论快捷连接（也称跳跃连接或残差连接）的概念。快捷连接最初是为计算机视觉中的深度网络（尤其是残差网络）提出的，用于缓解梯度消失问题。梯度消失是指在训练中指导权重更新的梯度在反向传播过程中逐渐减小，导致早期层（靠近输入端的网络层）难以有效训练，如图 4.12 所示。

<img src="../Image/chapter4/figure4.12.png" width="75%" />

如图 4.12 所示，快捷连接通过跳过一层或多层，为梯度提供一条更短的流动路径，这是通过将某层的输出加到后续层的输出上来实现的。因此，这种连接方式也称为跳跃连接。在反向传播中，快捷连接对保持梯度流动至关重要。

在以下代码示例中，我们将实现图 4.12 中所示的神经网络，以展示如何在前向传播方法中添加快捷连接：

```python
# Listing 4.5 A neural network to illustrate shortcut connections
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 层的深度神经网络，每层包括一个线性层和 GELU 激活函数。在前向传播中，我们将输入逐层传递，同时如果 `self.use_shortcut` 属性设置为 True，则会添加图 4.12 所示的快捷连接。

我们将使用以下代码初始化一个没有快捷连接的神经网络，其中每一层都被初始化为接受 3 个输入值并返回 3 个输出值。最后一层则返回一个单一的输出值：

```python
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
)
```

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

```python
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×3 的矩阵，那么这一层会有 3×3 的梯度值。然后我们打印出这 3×3 梯度值的绝对均值，以便得到每层的单一梯度值，从而更容易比较各层之间的梯度大小。

简而言之，`.backward()` 是 PyTorch 中一个便捷的方法，用于自动计算损失梯度，这在模型训练中是必要的。它让我们无需亲自实现梯度计算的数学过程，从而大大简化了深度神经网络的开发过程。如果您对梯度和神经网络训练不熟悉，建议参考附录 A 中的 A.4 节：**轻松实现自动微分** 和 A.7 节：**典型的训练循环**。

现在让我们使用 `print_gradients` 函数，并将其应用到没有跳跃连接的模型上：

```python
print_gradients(model_without_shortcut, sample_input)
```

输出如下：

```python
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）时逐渐减小，这种现象称为梯度消失问题。

我们再来创建一个带有跳跃连接的模型，看看它的表现如何：

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

输出如下：

```python
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 架构所需的最后一个模块。

> [!TIP]
>
> **个人思考：** 看到这里，不知各位读者是否真正理解了快捷连接在深度神经网络中的作用，这里其实涉及到快捷连接的两个重要的作用：
>
> + **保持信息（或者说是特征）流畅传递**
> + **缓解梯度消失问题**
>
> 让我们逐一解读，LLM 中的每个Transformer 模块通常包含两个重要组件（**可以先阅读完4.5节，再回头看这里的解读**）：
>
> 1. **自注意力层（Self-Attention Layer）**：计算每个 token 与其他 token 的关联，帮助模型理解上下文。
> 2. **前馈网络（Feed Forward Network）**：对每个 token 的嵌入（embedding）进行进一步的非线性转换，使模型能够提取更复杂的特征。
>
> 这两个部分都在层归一化（Layer Normalization）和快捷连接（Shortcut Connections）的配合下工作。
>
> 假设我们正在训练一个 LLM ，并希望它理解下面的句子：
>
> `The cat sat on the mat because it was tired.`
>
> 模型需要通过多个 Transformer 层来逐层处理该句子，使得每个词（token）在上下文中能被理解。为了达到这一目的，每个 token 的嵌入会在多层中进行注意力计算和前馈网络处理。
>
> 1. **没有快捷连接时的情况**
>
>    如果没有快捷连接，那么每个 Transformer 层的输出就直接传递到下一个层。这种情况下，网络中的信息流大致如下：
>
>    + **层间信息传递的局限**：假设当前层的注意力机制计算出了“it”和“cat”之间的关系，如果前馈网络进一步转换了这个信息，那么下一层就只能基于该层的输出，可能丢失一些最初的语义信息。
>    + **梯度消失**：在训练过程中，梯度从输出层逐层向回传播。如果层数过多，梯度会逐渐变小（即“梯度消失”），从而导致模型难以有效更新前面层的参数。
>
>    这种情况下，由于信息不能直接流动到更深层次的网络，可能会导致模型难以有效捕捉到前层的一些原始信息。
>
> 2. **加入快捷连接后的情况**
>
>    加入快捷连接后，信息可以在层与层之间**直接跳跃**。例如，假设在第 n 层，我们有输入 X<sub>n</sub>，经过注意力和前馈网络得到输出F(X<sub>n</sub>)。加入快捷连接后，这一层的输出可以表示为：
>
>    $$\text { 输出 }=X_{n}+F\left(X_{n}\right) $$
>
>    这意味着第 n 层的输出不仅包含了这一层的新信息 F(X<sub>n</sub>)，还保留了原始输入 X<sub>n </sub>的信息。下面是这样做的好处：
>
>    - **保留原始信息**
>
>      快捷连接让输入的原始信息直接传递到后续层，避免了在多层处理过程中丢失重要信息。例如，“it” 和 “cat” 之间的关系在较浅层中被捕捉到后，即使后面的层有进一步的处理，模型依然能够从快捷连接中获得最初的上下文信息。
>
>    - **减轻梯度消失**
>
>      假设我们有一个简单的三层网络，第三层的输出 O 是整个网络的输出。我们从损失函数 LLL 开始计算梯度：
>
>      - 根据反向传播的原理，**无快捷连接**时，梯度必须逐层传递，如下：
>
>        $$\frac{\partial L}{\partial X_{1}}=\frac{\partial L}{\partial X_{3}} \cdot \frac{\partial X_{3}}{\partial X_{2}} \cdot \frac{\partial X_{2}}{\partial X_{1}}$$
>
>        这里g，如果某一层的梯度值很小，那么梯度会被逐层缩小，导致梯度消失。
>
>      - **有快捷连接**时，假设我们在每一层之间都添加快捷连接，梯度的传播路径就多了一条直接路径：
>
>        $$\frac{\partial L}{\partial X_{1}}=\frac{\partial L}{\partial\left(X_{1}+F\left(X_{1}\right)\right)} \cdot\left(1+\frac{\partial F\left(X_{1}\right)}{\partial X_{1}}\right)$$
>
>        这样，即使 $` \frac{\partial F\left(X_{1}\right)}{\partial X_{1}} `$ 很小，梯度依然可以通过 111 这条路径直接传递到更前面的层。

  

## 4.5 在 Transformer 模块中连接注意力层与线性层

本节我们将实现 Transformer 模块，它是 GPT 和其他大语言模型架构的基本模块。这个在 124M 参数的 GPT-2 架构中重复了十几次的模块，结合了多头注意力、层归一化、dropout、前馈层和 GELU 激活等多个概念，详见图 4.13。在下一节中，我们将把这个 Transformer 模块连接到 GPT 架构的其余部分。

<img src="../Image/chapter4/figure4.13.png" width="75%" />

如图 4.13 所示，Transformer 模块结合了多个组件，包括第 3 章中的带掩码的多头注意力模块以及我们在 4.3 节中实现的前馈网络模块。

当 Transformer 模块处理输入序列时，序列中的每个元素（如单词或子词 token）都会被表示为固定大小的向量（如图 4.13 中为 768 维）。Transformer 模块中的操作，包括多头注意力和前馈层，旨在以维度不变的方式对这些向量进行转换。

之所以这样设计，是因为多头注意力模块中的自注意力机制用于识别和分析输入序列中各元素之间的关系，而前馈神经网络则对输入序列中每个位置的数据单独进行修改。这种组合不仅能够更细致地理解和处理输入信息，还增强了模型处理复杂数据模式的整体能力。

可以通过以下代码实现 Transformer 模块：

```python
# Listing 4.6 The transformer block component of GPT
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):
       shortcut = x                                          #A
       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），这种方式通常会导致较差的训练效果。

该类还实现了前向传播（`forward方法`），其中每个组件后面都设有一个快捷连接，将对应组件的输入添加到输出中。这一关键特性有助于在训练过程中促进梯度流动，从而提升 LLM 的学习能力，相关内容详见第 4.4 节。

现在使用我们之前定义的 `GPT_CONFIG_124M` 配置字典，实例化一个 Transformer 模块，并向其中输入一些示例数据：

```python
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 建一个形状为 [batch_size, num_tokens, emb_dim] 的输入张量
```

输出如下：

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

从代码输出可以看出，Transformer 模块的输出维度与输入维度保持一致，这说明 Transformer 架构在整个网络中处理序列数据时不会改变数据的形状。

Transformer 模块结构中保持数据形状不变并非偶然，而是其设计的一个关键特性。这种设计使 Transformer 擅长处理各种序列到序列任务，因为每个输出向量直接对应一个输入向量，保持一一对应关系。然而，输出向量是包含整个输入序列上下文信息的“上下文向量”。也就是说，尽管序列的物理维度（长度和特征维度）在经过 Transformer 模块时保持不变，但每个输出向量的内容会被重新编码，融合整个输入序列的上下文信息。

在本节完成了 Transformer 模块的实现后，我们已经具备了实现 GPT 架构所需的全部基础模块（如图 4.14 所示）。

<img src="../Image/chapter4/figure4.14.png" width="75%" />

如图 4.14 所示，Transformer 模块由层归一化、带有 GELU 激活函数的前馈网络和快捷连接组成，这些内容在本章前面已经讨论过。正如我们将在接下来的章节中看到的，这个 Transformer 模块将构成我们要实现的 GPT 架构的核心部分。