前面章节内容会慢慢补上！！！

4 从零开始实现一个用于生成文本的GPT模型

本章内容涵盖：

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

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

图4.1 展示了编码大型语言模型($LLM$)的三个主要阶段的思维模型，包括在通用文本数据集上对$LLM$进行预训练，并在标注数据集上进行微调。本章重点在于实现$LLM$的架构，我们将在下一章对其进行训练。

$LLM$的架构（如图4.1所示）由多个构建模块组成，我们将在本章中逐步实现这些模块。在接下来的部分，我们将首先从整体上概述模型架构，然后再详细介绍各个组件。

from importlib.metadata import version

print("matplotlib version:", version("matplotlib"))
print("torch version:", version("torch"))
print("tiktoken version:", version("tiktoken"))

matplotlib version: 3.10.0
torch version: 2.5.1
tiktoken version: 0.9.0

4.1 编写$LLM$架构

大型语言模型($LLM$)，如$GPT$（即生成式预训练变换器），是一种大型深度神经网络架构，旨在逐词（或逐标记）生成新文本。然而，尽管这些模型规模庞大，它们的架构并不像你想象的那么复杂，因为许多组件是重复使用的，我们将在后续内容中看到这一点。图4.2提供了一个类似$GPT$的$LLM$的整体视图，并突出显示了其主要组件。

图4.2 展示了$GPT$模型的思维模型。除了嵌入层外，它由一个或多个包含掩码多头注意力模块的变换器块组成，而该模块正是我们在上一章中实现的。

如图4.2所示，我们已经涵盖了多个方面，例如输入tokenization和嵌入，以及掩码多头注意力模块。本章的重点将是实现$GPT$模型的核心结构，包括其变换器块，我们将在下一章训练该模型以生成类人文本。

在前几章中，为了简化概念和示例，使其能够轻松地适应单页内容，我们使用了较小的嵌入维度。而在本章，我们将模型规模扩展到一个小型$GPT$-2模型的级别，具体来说是$Radford$等人在论文《Language Models are Unsupervised Multitask Learners》中描述的最小版本，该模型具有$1.24$亿个参数。需要注意的是，尽管原始报告中提到$1.17$亿个参数，但这一数值后来被更正。

第6章将重点介绍如何将预训练权重加载到我们的实现中，并适配到更大规模的$GPT$-2模型，包括$3.45$亿、$7.62$亿和$15.42$亿个参数的版本。在深度学习和类似$GPT$的$LLM$中，“参数”指的是模型的可训练权重。这些权重本质上是模型的内部变量，在训练过程中不断调整和优化，以最小化特定的损失函数，从而使模型能够从训练数据中学习。

例如，在一个神经网络层中，假设权重由一个$2048\times 2048$维矩阵（或张量）表示，那么该矩阵的每个元素都是一个参数。由于该矩阵有$2048$行和$2048$列，因此该层的总参数数量为$2048\times 2048=4194304$个参数。

$GPT$-2与$GPT$-3的对比

需要注意的是，我们专注于$GPT$-2，因为$OpenAI$已公开提供了其预训练模型的权重，我们将在第6章将其加载到我们的实现中。从模型架构的角度来看，$GPT$-3在本质上与$GPT$-2相同，只是将参数规模从$GPT$-2的$15$亿扩展到$GPT$-3的$1750$亿，并且训练数据规模更大。截至撰写本文时，$GPT$-3的权重尚未公开。

此外，$GPT$-2也是学习如何实现$LLM$的更佳选择，因为它可以在单台笔记本电脑上运行，而$GPT$-3的训练和推理则需要$GPU$集群。据$LambdaLabs$估算，使用单块$V100$数据中心$GPU$训练$GPT$-3需要$355$年，而使用消费级$RTX8000$则需要$665$年。

我们通过以下$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" 指的是词汇量大小，共$50257$个单词，与第2章介绍的$BPE$标记器使用的词汇表一致。
• "context_length" 表示模型可以处理的最大输入标记数，这通过第2章讨论的位置嵌入实现。
• "emb_dim" 代表嵌入维度，将每个标记转换为$768$维向量。
• "n_heads" 指的是多头注意力机制中的注意力头数，该机制已在第3章实现。
• "n_layers" 指定了模型中的变换器块数量，我们将在接下来的章节中详细介绍。
• "drop_rate" 表示随机失活(dropout)的强度（$0.1$意味着隐藏单元有$10\%$的概率被丢弃），用于防止过拟合，这在第3章已介绍。
• "qkv_bias" 决定是否在多头注意力机制中的查询、键和值的线性层中包含偏置向量。最初，我们将禁用该选项，以遵循现代$LLM$的标准做法，但在第6章加载来自$OpenAI$的预训练$GPT$-2权重时，我们将重新讨论这一设置。

基于上述配置，我们将在本节实现一个$GPT$占位架构(DummyGPTModel)，如图4.3所示。这将帮助我们从整体上理解模型的架构，以及在接下来的章节中需要编写哪些组件，以最终组装完整的$GPT$模型架构。

图4.3 展示了$GPT$架构的编写顺序。在本章中，我们将从$GPT$的主干架构（占位架构）开始，然后逐步实现各个核心模块，并最终将它们整合到变换器块中，以构建完整的$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__()
        '''创建一个词嵌入层，将输入的 token 索引映射到一个固定维度的向量空间。'''
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        '''创建一个位置嵌入层，将每个位置映射到一个固定维度的向量。'''
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        '''构造一个 Dropout 层，用于在训练时随机丢弃部分神经元，防止过拟合。'''
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        
        # Use a placeholder for TransformerBlock
        '''
        * 操作符的作用是将列表中的元素拆开,为“解包”操作符。
        将这个列表中的所有元素“解包”，使它们变成 nn.Sequential 的单独参数。
        等效于写成（假设cfg["n_layers"] = 2）：
        self.trf_blocks = nn.Sequential(DummyTransformerBlock(cfg), DummyTransformerBlock(cfg))
        创建一个由多个 TransformerBlock 组成的序列，然后用 nn.Sequential 将它们串联在一起。
        '''
        self.trf_blocks = nn.Sequential(*[DummyTransformerBlock(cfg) for _ in range(cfg["n_layers"])])
        
        # Use a placeholder for LayerNorm
        '''
        创建一个层归一化层，通常用于稳定训练，但这里使用的是一个占位实现（DummyLayerNorm），实际不做任何操作。
        '''
        self.final_norm = DummyLayerNorm(cfg["emb_dim"])
        '''创建一个线性层，将最后得到的嵌入映射到词汇表大小的维度，输出 logits。'''
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

    def forward(self, in_idx):
        """定义模型的前向传播接口，输入为 token 索引组成的张量 in_idx。"""
        batch_size, seq_len = in_idx.shape
        '''
        将 token 索引转换为对应的嵌入向量。
        输出形状为 (batch_size, seq_len, emb_dim)，例如 (2, 6, 8)。
        '''
        tok_embeds = self.tok_emb(in_idx)
        '''
        生成位置序列 [0, 1, ..., seq_len-1] 并查找对应的位置信息向量。
        输出形状为 (seq_len, emb_dim)，例如 (6, 8)。
        '''
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        '''
        将 token 嵌入与位置嵌入相加，融合语义信息和位置信息。
        pos_embeds 的形状 (6, 8) 自动扩展到 (2, 6, 8) 与 tok_embeds 相加。
        1. tok_embeds[0](6,8)+pos_embeds(6, 8) ->[0](6,8)
        2. tok_embeds[1](6,8)+pos_embeds(6, 8) ->[1](6,8)
        变成(2, 6, 8)
        '''
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)
        '''
        将 x 依次传入所有 DummyTransformerBlock 层处理。
        '''
        x = self.trf_blocks(x)
        '''
        对 x 进行层归一化操作，稳定训练。
        '''
        x = self.final_norm(x)
        '''
        通过线性变换，将每个 token 的 8 维向量映射到词汇表大小维度（例如 100），得到每个词的 logits。
        '''
        logits = self.out_head(x)
        return logits


class DummyTransformerBlock(nn.Module):
    """定义一个占位用的 TransformerBlock 模块。"""
    def __init__(self, cfg):
        super().__init__()
        # A simple placeholder

    def forward(self, x):
        # This block does nothing and just returns its input.
        return x


class DummyLayerNorm(nn.Module):
    """定义一个占位用的 LayerNorm 层，用于模拟接口。"""
    def __init__(self, normalized_shape, eps=1e-5):
        super().__init__()
        # The parameters here are just to mimic the LayerNorm interface.

    def forward(self, x):
        # This layer does nothing and just returns its input.
        return x

pytorch小记（十四）：pytorch中 nn.Embedding 详解

DummyGPTModel类在这段代码中定义了一个使用$PyTorch$神经网络模块(nn.Module)的简化版$GPT$模型。DummyGPTModel类的模型架构包括标记嵌入和位置嵌入、随机失活、若干个变换器块(DummyTransformerBlock)、最终的层归一化(DummyLayerNorm)，以及一个线性输出层(out_head)。模型的配置通过$Python$字典传入，例如我们之前创建的GPT_CONFIG_124M字典。

forward方法描述了数据在模型中的流动过程：它计算输入索引的标记嵌入和位置嵌入，应用随机失活，然后将数据传递至变换器块进行处理，接着应用归一化，并最终通过线性输出层生成对数概率(logits)。

上述代码已经是一个可运行的模型，我们将在本节后续部分准备输入数据并进行测试。然而，目前需要注意的是，代码中我们使用了占位符(DummyLayerNorm和DummyTransformerBlock)来代替变换器块和层归一化，这些组件将在后续部分实现。

接下来，我们将准备输入数据并初始化一个新的$GPT$模型，以演示其使用方式。基于我们在第2章中编写标记器的内容，图4.4提供了数据在$GPT$模型中流入和流出的整体概览。

图4.4 展示了数据流入$GPT$模型的全局视图，包括输入数据的标记化、嵌入处理，并最终输入到$GPT$模型。需要注意，在我们之前编写的DummyGPTClass中，标记嵌入是在$GPT$模型内部处理的。在$LLM$中，嵌入输入标记的维度通常与输出维度相匹配。这里的输出嵌入代表了我们在第3章讨论的上下文向量。

为了实现图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], #A
        [6109, 1110, 6622,  257]])

接下来，我们初始化一个具有$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([[[-0.9289,  0.2748, -0.7557,  ..., -1.6070,  0.2702, -0.5888],
         [-0.4476,  0.1726,  0.5354,  ..., -0.3932,  1.5285,  0.8557],
         [ 0.5680,  1.6053, -0.2155,  ...,  1.1624,  0.1380,  0.7425],
         [ 0.0447,  2.4787, -0.8843,  ...,  1.3219, -0.0864, -0.5856]],

        [[-1.5474, -0.0542, -1.0571,  ..., -1.8061, -0.4494, -0.6747],
         [-0.8422,  0.8243, -0.1098,  ..., -0.1434,  0.2079,  1.2046],
         [ 0.1355,  1.1858, -0.1453,  ...,  0.0869, -0.1590,  0.1552],
         [ 0.1666, -0.8138,  0.2307,  ...,  2.5035, -0.3055, -0.3083]]],
       grad_fn=<UnsafeViewBackward0>)

Output shape: torch.Size([2, 4, 50257])：可以理解成2个batch，每个batch有4个单词（token），每个单词（token）用50257维来表示

输出张量包含两行，对应于两条文本样本。每个文本样本由$4$个标记组成，每个标记是一个$50257$维的向量，这与标记器的词汇表大小相匹配。

嵌入层的维度为$50257$，因为这些维度中的每一个都对应于词汇表中的一个唯一标记。在本章末尾，当我们实现后处理代码时，我们将把这些$50257$维的向量转换回标记$ID$，然后解码为单词。

现在，我们已经从整体上了解了$GPT$架构及其输入和输出，接下来将在后续部分逐步编写具体的占位组件，从实现真实的层归一化类开始，以替换之前代码中的DummyLayerNorm。

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

在训练具有多个层的深度神经网络时，可能会遇到梯度消失或梯度爆炸等问题。这些问题会导致训练过程不稳定，使得网络难以有效地调整其权重。这意味着学习过程难以找到一组能够最小化损失函数的参数（权重）。换句话说，网络难以学习数据中的潜在模式，从而无法做出准确的预测或决策。

（如果你对神经网络训练和梯度的概念不熟悉，可以参考附录A中的A.4节 《自动微分简明介绍》。然而，理解本教程内容并不需要对梯度有深奥的数学理解。）

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

层归一化的核心思想是调整神经网络层的激活值（输出），使其均值为$0$，方差为$1$，也称为单位方差。这种调整可以加速收敛，使网络更快地找到有效的权重，并确保训练过程的稳定性和可靠性。正如我们在上一节基于DummyLayerNorm占位符所看到的，在$GPT$-2及现代变换器架构中，层归一化通常应用于多头注意力模块的前后，以及最终输出层的前面。

在我们编写层归一化的代码之前，图4.5提供了层归一化的可视化概览。

图4.5展示了如何对$5$个层的输出（即激活值）进行归一化，使其均值为$0$，方差为$1$。

我们可以通过以下代码重新创建图4.5所示的示例，其中我们实现了一个具有$5$个输入和$6$个输出的神经网络层，并将其应用于两个输入示例：

torch.manual_seed(123)

# create 2 training examples with 5 dimensions (features) each
batch_example = torch.randn(2, 5) 
print(batch_example)
layer = nn.Sequential(nn.Linear(5, 6), nn.ReLU())
out = layer(batch_example)
print(out)

这将打印出以下张量，其中第一行列出了第一个输入的层输出，第二行列出了第二个输入的层输出：

tensor([[-0.1115,  0.1204, -0.3696, -0.2404, -1.1969],
        [ 0.2093, -0.9724, -0.7550,  0.3239, -0.1085]])
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$是神经网络中的标准激活函数。如果你对$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，返回的均值张量将是一个$2$维向量[0.1324, 0.2170]，而不是一个$2\times 1$维矩阵[[0.1324], [0.2170]]。

dim参数指定在张量的哪个维度上计算统计量（如均值或方差），其作用如图4.6所示。

图4.6 展示了在计算张量均值时dim参数的作用。例如，如果我们有一个$2D$张量（矩阵），其维度为$[行,列]$，那么使用dim=0将在行方向（纵向，见图底部）执行操作，得到的输出聚合了每列的数据。而使用dim=1或dim=-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)

从结果可以看出，归一化后的层输出现在包含了负值，并且均值为$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([[9.9341e-09],
        [0.0000e+00]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[1.0000],
        [1.0000]], grad_fn=<VarBackward0>)

需要注意的是，输出张量中的值9.9341e-08是科学计数法的表示方式，相当于$9.9341\times 10^{-9}$，即十进制形式的0.0000000099341。这个值非常接近$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.0000],
        [1.0000]], 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是一个很小的常数（$ϵ=10^{-5}$），用于防止归一化过程中出现除零错误。

scale和shift是两个可训练参数，它们的维度与输入相同。在$LLM$的训练过程中，模型会自动调整这些参数，以优化其在训练任务上的表现。这使得模型能够学习适合其处理数据的最佳缩放和偏移，从而提高训练效果。

有偏方差计算

在方差计算方法中，我们选择了一个实现细节，即将unbiased=False。对于好奇这一点的读者，在方差计算中，我们使用公式中的样本数量$n$作为分母进行计算，而不是使用$n-1$（即贝塞尔校正）。贝塞尔校正通常用于修正样本方差估计中的偏差，而我们的做法会导致所谓的有偏方差估计。

对于大型语言模型($LLM$)，由于嵌入维度$n$通常非常大，使用$n$和$n-1$之间的差异实际上可以忽略不计。我们选择这种方式，以确保与$GPT$-2模型的归一化层保持兼容，并且这样做也符合$TensorFlow$的默认行为（$GPT$-2的原始实现基于$TensorFlow$）。采用类似的设置可以确保我们的方法与第6章将要加载的预训练权重兼容。

现在，让我们在实践中测试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 展示了本章中实现$GPT$架构所需的各个构建模块的思维模型。

在下一节中，我们将介绍 GELU（高斯误差线性单元）激活函数，它是$LLM$中常用的激活函数之一，而不是本节中使用的传统$ReLU$函数。

层归一化 vs. 批量归一化

如果你熟悉 批量归一化(Batch Normalization)，这是一种常见的神经网络归一化方法，那么你可能会好奇它与 层归一化(Layer Normalization) 的对比。

与批量归一化不同，批量归一化是在 批量(batch) 维度 上进行归一化，而层归一化则是在 特征(feature) 维度 上进行归一化。

由于$LLM$通常需要大量计算资源，可用的硬件或特定的使用场景可能会影响训练或推理时的批量大小(batch size)。而层归一化对每个输入独立进行归一化，不受批量大小的影响，因此在这些情况下提供了更大的灵活性和稳定性。

这种特性在 分布式训练 或 计算资源受限的环境 中尤其有利，使得模型能够在不同的设备和计算环境下保持稳定的表现。

4.3 使用GELU激活函数实现前馈神经网络

在本节中，我们将实现一个小型神经网络子模块，该模块是$LLM$中变换器块的一部分。我们首先实现 GELU（高斯误差线性单元）激活函数，它在此神经网络子模块中起着关键作用。（如果需要额外了解如何在$PyTorch$中实现神经网络，请参考附录A的A.5节 《实现多层神经网络》。）

GELU vs. ReLU

历史上，ReLU（修正线性单元）因其简单性和在各种神经网络架构中的有效性而被广泛使用。然而，在$LLM$中，除了传统的$ReLU$之外，还使用了其他几种激活函数，其中两个重要的例子是：

• GELU（高斯误差线性单元）
• SwiGLU（Sigmoid加权线性单元）

GELU和SwiGLU比$ReLU$更复杂，它们分别结合了 高斯分布 和 Sigmoid门控线性单元，能够为深度学习模型提供更好的性能。

GELU激活函数的定义

GELU激活函数可以用多种方式实现，最精确的数学定义如下：

$$\text{GELU(x)}=\text{x}\cdot \Phi (\text{x})$$

其中，$\Phi (x)$ 是 标准高斯分布的累积分布函数(CDF)。

但在实际应用中，通常采用计算量更小的近似版本（$GPT$-2模型的原始训练也是基于此近似实现）：

$$\text{GELU(x)}\approx 0.5\cdot \text{x}\cdot (1+\tanh (\sqrt{2/\pi }\cdot (\text{x}+0.044715\cdot {\text{x}}^3)))$$

代码实现

我们可以将此函数实现为$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()

从 图4.8 所示的绘图结果可以看出，ReLU 是一个分段线性函数：当输入为正时，它直接输出输入值；当输入为负时，它输出$0$。

而 GELU 是一个平滑的非线性函数，它近似$ReLU$，但对于负值输入仍然保持了非零梯度。

图4.8 展示了使用matplotlib绘制的 GELU 和 ReLU 函数的输出。横轴($x$轴)表示函数输入，纵轴($y$轴)表示函数输出。

GELU 的平滑特性（如图4.8所示）可以在训练过程中带来更好的优化性能，因为它允许对模型参数进行更细微的调整。相比之下，ReLU 在$0$处存在一个尖角，这可能会在优化过程中带来困难，尤其是在深层网络或复杂架构中。

此外，与 ReLU 不同，ReLU 对任何负输入都会输出$0$，而 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)

print(GPT_CONFIG_124M["emb_dim"])

768

如上面的代码所示，FeedForward 模块是一个由 两个线性层(Linear layers) 和 GELU激活函数 组成的小型神经网络。在 具有1.24亿参数的$GPT$模型 中，该模块接收 输入批次(batches)，其中每个标记(token)的嵌入维度(embedding size) 为 768，这一配置由GPT_CONFIG_124M字典中的参数 GPT_CONFIG_124M["emb_dim"] = 768 确定。

图4.9 展示了当输入通过该 前馈神经网络(Feed Forward Neural Network) 时，嵌入维度如何在该小型网络内部进行变换。

图4.9 直观地展示了前馈神经网络各层之间的连接关系。值得注意的是，该神经网络可以处理不同的批量大小(batch size) 和 输入标记数(num_tokens)，但每个标记的嵌入维度(embedding size) 在初始化权重时即被确定，并保持固定。

按照 图4.9 中的示例，我们初始化一个新的 FeedForward 模块，其 标记嵌入维度(embedding size) 为 768，并向其输入一个包含 2个样本，每个样本包含3个标记 的批量输入：

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

如我们所见，输出张量的形状与输入张量相同：

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

在本节中实现的 FeedForward 模块在 增强模型的学习能力和泛化能力 方面发挥着重要作用。

尽管该模块的 输入和输出维度相同，但其 内部 通过第一层线性变换(Linear Layer) 将嵌入维度扩展到 更高维度空间，然后经过 GELU非线性激活，最终通过 第二层线性变换 将维度 收缩回原始尺寸，如 图4.10 所示。

这种设计使得模型能够探索更丰富的表示空间，进而提升对数据的理解能力。

图4.10 展示了 前馈神经网络(Feed Forward Neural Network) 中 层输出的扩展与收缩过程：1. 第一层线性变换：将输入 扩展4倍，即 从768维度扩展到3072维度。2. GELU非线性激活：对扩展后的向量进行非线性变换。3. 第二层线性变换：将3072维度压缩回768维度，以匹配原始输入的维度。

此外，输入和输出维度的一致性 简化了模型架构，使得我们可以 堆叠多个层（我们将在后续实现），而无需在层之间调整维度，从而提高模型的 可扩展性(scalability)。

如 图4.11 所示，我们已经实现了 大部分$LLM$的核心构建模块。

图4.11 展示了本章涵盖的主题，其中 黑色对勾标记表示我们已经完成的部分。

在下一节中，我们将介绍 捷径连接(shortcut connections) 的概念。这些连接被插入到神经网络的不同层之间，对于 提升深度神经网络架构的训练性能 至关重要。

4.4 添加捷径连接

接下来，我们讨论 捷径连接(shortcut connections) 的概念，这种连接也称为 跳跃连接(skip connections) 或 残差连接(residual connections)。

最初，捷径连接 是为 计算机视觉中的深度网络（尤其是残差网络(ResNet)）提出的，目的是 缓解梯度消失问题(vanishing gradient problem)。

梯度消失问题 指的是，在训练过程中，梯度（用于指导权重更新）在向后传播时逐渐变小，导致模型难以有效训练 靠近输入端的早期层，如 图4.12 所示。

图4.12 比较了 一个包含5层的深度神经网络，其中：

• 左侧 是 没有捷径连接 的网络。
• 右侧 是 添加了捷径连接 的网络。

捷径连接的核心思想是：将某一层的输入直接加到该层的输出上，从而 创建一个绕过某些层的额外路径。图1.1 中所示的 梯度(gradient) 代表 每一层的平均绝对梯度值，我们将在后续的代码示例中计算这些梯度。

如 图4.12 所示，捷径连接(shortcut connection) 创建了一条 额外的、较短的梯度传播路径，使梯度可以绕过一层或多层直接流经网络。这种 跳跃连接(skip connection) 通过 将前一层的输出加到后续层的输出上 实现，因此在训练的 反向传播(backward pass) 过程中，它们在 保持梯度流动 方面起到了至关重要的作用。

在下面的代码示例中，我们实现了 图4.12 中所示的神经网络，并探索如何在 前向传播(forward method) 中添加 捷径连接：

代码清单 4.5 实现带有捷径连接的深度神经网络

class ExampleDeepNeuralNetwork(nn.Module):
    def __init__(self, layer_sizes, use_shortcut):
        super().__init__()
        self.use_shortcut = use_shortcut
        self.layers = nn.ModuleList([
            # 5 层，每层包含一个线性层(Linear)和GELU激活函数
            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)
            # 如果启用了捷径连接，并且形状匹配，则应用捷径连接
            if self.use_shortcut and x.shape == layer_output.shape:
                x = x + layer_output
            else:
                x = layer_output
        return x

代码解析

• 该代码实现了一个 深度神经网络，包含 5层，每一层由：
  • 一个线性层(Linear layer)
  • GELU激活函数
• 在 前向传播(forward pass) 过程中：
  1. 逐层计算输入的输出
  2. 如果 self.use_shortcut=True 且当前层的输入与输出形状相同，则 应用捷径连接(x = x + layer_output)
  3. 否则，直接使用当前层的输出(x = layer_output)

初始化一个不带捷径连接的神经网络

我们首先初始化 一个没有捷径连接的神经网络。这里，每一层都被初始化为：

• 输入为3个数值
• 输出为3个数值
• 最后一层输出1个数值

layer_sizes = [3, 3, 3, 3, 3, 1]
sample_input = torch.tensor([[1., 0., -1.]])

torch.manual_seed(123)  # 指定随机种子，以确保可复现性
model_without_shortcut = ExampleDeepNeuralNetwork(
    layer_sizes, 
    use_shortcut=False
)

这样，我们创建了一个 不包含捷径连接的神经网络，其权重由 随机种子 设定，以确保实验结果可复现。

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

def print_gradients(model, x):
    # 前向传播
    output = model(x)
    target = torch.tensor([[0.]])
    
    # 计算损失，衡量模型输出与目标值（此处设为0）的接近程度
    loss = nn.MSELoss()
    loss = loss(output, target)

    # 反向传播计算梯度
    loss.backward()

    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() 迭代模型的权重参数 并获取其梯度信息。假设某一层的权重参数矩阵是 $3\times 3$，那么该层将会有 $3\times 3$ 个梯度值。为了更方便地比较不同层的梯度，我们 计算并打印这些梯度的平均绝对值。

关于 .backward()

• .backward() 方法 是 PyTorch 中的一个便捷函数，可 自动计算损失函数的梯度，无需手动实现梯度计算的数学公式。
• 这一方法使得 深度神经网络的训练更加高效和便捷。

如果你对 梯度(gradients) 和 神经网络训练过程 还不熟悉，建议阅读 附录A 中的以下内容：

• A.4: 《自动微分简明介绍(Auto Differentiation Made Easy)》
• A.7: 《典型训练循环(A Typical Training Loop)》

现在，我们使用 print_gradients 函数，并将其应用于 不带跳跃连接(skip connections) 的模型：

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) 逐渐减小，这种现象被称为 梯度消失问题(vanishing gradient problem)。

创建带跳跃连接的模型，并进行比较

接下来，我们 实例化一个包含跳跃连接(skip connections) 的模型，并观察梯度分布的变化：

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.20694106817245483
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) 传播的过程中趋于稳定，并且不会缩小到极小的值。

捷径连接(shortcut connections) 在深度神经网络中对于克服 梯度消失问题(vanishing gradient problem) 的限制至关重要。捷径连接是 超大规模模型(如$LLM$) 的核心构建模块，它们在训练$GPT$模型时确保梯度能够稳定地 在各层之间流动，从而 提高训练的有效性。

在介绍完 捷径连接 之后，我们将在 下一节 中 整合 之前讨论的所有核心概念：

• 层归一化(Layer Normalization)
• GELU激活函数
• 前馈模块(Feed Forward Module)
• 捷径连接(Shortcut Connections)

最终，我们将 在变换器块(Transformer Block) 中实现这些组件，这是 $GPT$架构的最后一个构建模块。

4.5 在变换器块中连接注意力机制和线性层

在本节中，我们将实现 变换器块(Transformer Block)，这是 GPT 及其他 $LLM$架构 的 核心构建模块。

在 $GPT$-2(1.24亿参数) 的架构 中，该变换器块被 重复多次，并整合了我们之前介绍的多个概念：

• 多头注意力机制(Multi-Head Attention)
• 层归一化(Layer Normalization)
• 随机失活(Dropout)
• 前馈层(Feed Forward Layers)
• GELU激活函数

这一结构如 图4.13 所示。

在下一节中，我们将 把该变换器块与$GPT$架构的其余部分连接，完成最终的 完整$GPT$模型。

图4.13 展示了 变换器块(Transformer Block) 的组成结构：

• 底部 显示了 输入标记(tokens)，它们已经被 转换为768维的向量（即 嵌入表示）。
• 每一行(row) 对应 一个标记的向量表示。
• 变换器块的输出 与 输入的维度相同，因此可以进一步输入到 后续的$LLM$层 进行处理。

如 图4.13 所示，变换器块(Transformer Block) 结合了多个关键组件，包括：

• 第3章介绍的掩码多头注意力模块(Masked Multi-Head Attention Module)
• 第4.3节实现的前馈模块(FeedForward Module)

当变换器块处理 输入序列 时，序列中的每个元素（如单词或子词标记） 都会被表示为一个 固定大小的向量（在 图4.13 中，该向量维度为 768）。变换器块的操作（包括 多头注意力机制和前馈层）不会改变输入向量的维度，而是通过 变换这些向量的表示方式 来增强模型的处理能力。

核心思想

• 多头自注意力机制(Multi-Head Self-Attention) 负责 识别和分析输入序列中元素之间的关系。
• 前馈神经网络(Feed Forward Network, FFN) 在 每个位置单独修改数据。

这种 组合方式 既能 让模型更深入地理解输入序列，又能 增强模型处理复杂数据模式的能力。

代码实现

我们可以使用以下代码实现 TransformerBlock：

代码清单 4.6 $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"]
        )
        '''
        (1,4)@(4,16) -> (1,16)
        (1,16)@(16,4) -> (1,4)
        '''
        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 connection for attention block
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        # Shortcut connection for feed forward block
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        return x

下载 previous_chapters.py 到相同目录中

按照图4.13流程来走，很容易理解！

该代码定义了 TransformerBlock 类，这是一个 $PyTorch$ 实现的变换器块，其中包含：

• 多头注意力机制(MultiHeadAttention)
• 前馈网络(FeedForward)
• 层归一化(LayerNorm)
• 随机失活(Dropout)

这些组件的配置基于 提供的配置字典(cfg)，例如 GPT_CONFIG_124M。

Pre-LayerNorm vs. Post-LayerNorm

• 在 本实现 中，层归一化 (LayerNorm) 应用在多头注意力和前馈网络之前，这种方式被称为 “Pre-LayerNorm”。
• 在 早期的变换器架构（如原始$Transformer$模型） 中，层归一化是应用在这些模块之后，这种方式被称为 “Post-LayerNorm”。
• “Post-LayerNorm” 可能会导致 训练动态较差，因此现代$LLM$通常采用 “Pre-LayerNorm”。

前向传播 (Forward Pass) 该类还实现了 前向传播，其中：

• 每个组件 后面都加上捷径连接(shortcut connection)，即 将块的输入添加到其输出。
• 这种 捷径连接 能够 帮助梯度在训练过程中更顺畅地传播，从而 改进深度模型的学习能力（详见 4.4节）。

代码示例：实例化变换器块
使用 GPT_CONFIG_124M 配置，我们创建一个 变换器块 并向其输入一些样本数据：

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)

输出：

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

保持输入输出形状不变
从代码输出可以看出：

• 变换器块的输入和输出维度相同，即：
  $$\text{Input shape}=\text{Output shape}=(batch\_size,sequence\_length,embedding\_size)$$
• 这表明 变换器架构在处理序列数据时不会改变其形状，而是在相同的形状下进行信息转换。

为什么保持形状不变很重要？

• 这是 变换器架构的核心设计理念，使其能够适用于 各种序列到序列(Sequence-to-Sequence)任务。
• 每个输出向量 直接 对应一个输入向量，保持 一对一关系。
• 尽管物理维度(sequence length & feature size) 不变，但 输出向量的内容已经重新编码，并整合了整个输入序列的上下文信息（详见 第3章）。

完成 $GPT$ 架构的所有构建模块
在本节中，我们成功实现了 变换器块，至此，我们已经具备了 构建$GPT$架构所需的全部组件，如 图4.14 所示。

图4.14 展示了 本章已实现的所有核心概念，这些构建模块将在下一节 完整地组合为$GPT$模型。

如 图4.14 所示，变换器块(Transformer Block) 结合了多个关键组件，包括：

• 层归一化(Layer Normalization)
• 前馈网络(Feed Forward Network)（包含 GELU激活函数）
• 捷径连接(Shortcut Connections)

这些组件我们已经在本章的前面部分进行了介绍和实现。

在下一章中，我们将看到 变换器块 将作为 $GPT$架构的主要组成部分，并用于构建完整的$GPT$模型。

4.6 编写GPT模型

本章开始时，我们对 GPT架构 进行了整体概述，并实现了一个 占位版本(DummyGPTModel)。在 DummyGPTModel 代码实现中，我们展示了 GPT模型的输入和输出，但其内部构建模块仍然是一个 黑盒，使用了 DummyTransformerBlock 和 DummyLayerNorm 作为占位符。

在本节中，我们将 用真实的TransformerBlock和LayerNorm类 替换这些 占位符，从而 组装出完整可运行的$GPT$-2模型(1.24亿参数版本)。

在 第5章，我们将 预训练(Pretrain) $GPT$-2模型，而在 第6章，我们将 加载来自$OpenAI$的预训练权重。

GPT-2 模型整体结构

在正式编写 GPT-2模型代码 之前，让我们先看 图4.15，它展示了 GPT模型的整体架构，并汇总了本章介绍的所有核心概念。

图4.15：GPT模型架构概览

• 数据流(data flow) 从底部开始：

1. 标记化文本(Tokenized Text) 转换为 标记嵌入(Token Embeddings)。
2. 标记嵌入 再与 位置嵌入(Positional Embeddings) 结合，以编码序列信息。
3. 这些 组合信息 形成 张量(Tensor)，并输入到 变换器块(Transformer Blocks) 进行处理。

• 中间部分(核心计算)

4. 变换器块(Transformer Blocks) 由 多个层堆叠组成，每个变换器块包括：
   • 多头注意力(Multi-Head Attention)
   • 前馈神经网络(Feed Forward Neural Network, FFN)
   • 层归一化(Layer Normalization)
   • 随机失活(Dropout)
5. 变换器块被堆叠并重复12次（即$GPT$-2的12层架构）。

在接下来的代码部分，我们将正式实现 完整的$GPT$-2架构。

如 图4.15 所示，在 $GPT$模型架构 中，我们在 4.5节 编写的 变换器块(Transformer Block) 会被多次重复。

• 对于$GPT$-2的1.24亿参数版本，该 变换器块重复12次，这一数值由 GPT_CONFIG_124M["n_layers"] = 12 指定。
• 对于$GPT$-2的最大版本(15.42亿参数)，该 变换器块重复36次。

最终的输出层

如 图4.15 所示：

1. 最后一个变换器块的输出 经过 最终的层归一化(Layer Normalization)。
2. 随后进入最终的线性输出层(Linear Output Layer)，该层将 变换器的输出投影到高维空间：
   • 在 $GPT$-2 中，这个维度是 50,257，对应 模型的词汇表大小(vocabulary size)。
3. 最终输出预测序列中的下一个标记(token)。

现在，我们来 编写代码，实现 图4.15所示的$GPT$模型架构：

代码清单 4.7 $GPT$模型架构的实现

class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        '''创建一个词嵌入层，将输入的 token 索引映射到一个固定维度的向量空间。'''
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        '''创建一个位置嵌入层，将每个位置映射到一个固定维度的向量。'''
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        '''构造一个 Dropout 层，用于在训练时随机丢弃部分神经元，防止过拟合。'''
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        
        # Use a placeholder for TransformerBlock
        '''
        * 操作符的作用是将列表中的元素拆开,为“解包”操作符。
        将这个列表中的所有元素“解包”，使它们变成 nn.Sequential 的单独参数。
        等效于写成（假设cfg["n_layers"] = 2）：
        self.trf_blocks = nn.Sequential(TransformerBlock(cfg), TransformerBlock(cfg))
        创建一个由多个 TransformerBlock 组成的序列，然后用 nn.Sequential 将它们串联在一起。
        '''
        self.trf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])
        
        # Use a placeholder for LayerNorm
        '''
        创建一个层归一化层，通常用于稳定训练，但这里使用的是一个占位实现（LayerNorm），实际不做任何操作。
        '''
        self.final_norm = LayerNorm(cfg["emb_dim"])
        '''创建一个线性层，将最后得到的嵌入映射到词汇表大小的维度，输出 logits。'''
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

    def forward(self, in_idx):
        """定义模型的前向传播接口，输入为 token 索引组成的张量 in_idx。"""
        batch_size, seq_len = in_idx.shape
        '''
        将 token 索引转换为对应的嵌入向量。
        输出形状为 (batch_size, seq_len, emb_dim)，例如 (2, 6, 8)。
        '''
        tok_embeds = self.tok_emb(in_idx)
        '''
        生成位置序列 [0, 1, ..., seq_len-1] 并查找对应的位置信息向量。
        输出形状为 (seq_len, emb_dim)，例如 (6, 8)。
        '''
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        '''
        将 token 嵌入与位置嵌入相加，融合语义信息和位置信息。
        pos_embeds 的形状 (6, 8) 自动扩展到 (2, 6, 8) 与 tok_embeds 相加。
        1. tok_embeds[0](6,8)+pos_embeds(6, 8) ->[0](6,8)
        2. tok_embeds[1](6,8)+pos_embeds(6, 8) ->[1](6,8)
        变成(2, 6, 8)
        '''
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)
        '''
        将 x 依次传入所有 TransformerBlock 层处理。
        '''
        x = self.trf_blocks(x)
        '''
        对 x 进行层归一化操作，稳定训练。
        '''
        x = self.final_norm(x)
        '''
        通过线性变换，将每个 token 的 8 维向量映射到词汇表大小维度（例如 100），得到每个词的 logits。
        '''
        logits = self.out_head(x)
        return logits

由于我们在 4.5节 实现了 TransformerBlock 类，因此 GPTModel 类的代码相对紧凑且易于理解。

在 GPTModel 类的 __init__ 构造函数 中：

• 标记嵌入层(Token Embedding Layer) 和 位置嵌入层(Positional Embedding Layer) 由传入的 配置字典(cfg) 进行初始化。
  • 这些 嵌入层 负责 将输入的标记索引转换为密集向量，并 添加位置信息（详见 第2章）。
• 随后，创建一系列 TransformerBlock 变换器块，其数量等于 cfg["n_layers"] 中指定的层数。
• 在变换器块之后，应用层归一化(LayerNorm)，标准化变换器块的输出，以 稳定学习过程。
• 最后，定义一个无偏置(bias=False)的线性输出层，它 将变换器的输出投影到词汇空间，用于 计算词汇表中每个标记的对数概率(logits)。

前向传播(forward 方法)

• 接收一个批量输入(batch)，其中包含标记索引。
• 计算其标记嵌入，并应用位置嵌入。
• 通过变换器块(Transformer Blocks) 处理序列。
• 对最终输出进行层归一化。
• 计算对数概率(logits)，用于预测下一个标记。

在下一节中，我们将 把这些对数概率转换为实际的标记(tokens) 和文本输出(text output)。

初始化 1.24 亿参数的 GPT 模型

现在，我们使用 GPT_CONFIG_124M 配置字典初始化 $GPT$-2模型，并输入我们在本章开始时创建的批量文本数据：

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],  # 文本1的标记ID
        [ 6109, 1110, 6622, 257]])  # 文本2的标记ID

Output shape: torch.Size([2, 4, 50257])
tensor([[[ 0.1381,  0.0077, -0.1963,  ..., -0.0222, -0.1060,  0.1717],
         [ 0.3865, -0.8408, -0.6564,  ..., -0.5163,  0.2369, -0.3357],
         [ 0.6989, -0.1829, -0.1631,  ...,  0.1472, -0.6504, -0.0056],
         [-0.4290,  0.1669, -0.1258,  ...,  1.1579,  0.5303, -0.5549]],

        [[ 0.1094, -0.2894, -0.1467,  ..., -0.0557,  0.2911, -0.2824],
         [ 0.0882, -0.3552, -0.3527,  ...,  1.2930,  0.0053,  0.1898],
         [ 0.6091,  0.4702, -0.4094,  ...,  0.7688,  0.3787, -0.1974],
         [-0.0612, -0.0737,  0.4751,  ...,  1.2463, -0.3834,  0.0609]]],
       grad_fn=<UnsafeViewBackward0>)

---

解析输出

1. 输入 (batch) 形状为 [2, 4]，表示 两个文本样本，每个文本包含 4 个标记(token)。
2. 输出 (out) 形状为 [2, 4, 50257]：
   • 2：表示 批量大小(batch size)，即两个输入样本。
   • 4：表示 序列长度(sequence length)，即每个输入文本包含 4 个标记。
   • 50257：表示 模型词汇表大小(vocabulary size)，即 每个标记都有 50257 维的对数概率(logits)，用于预测 下一个可能的标记。
3. 最终输出的 tensor：
   • 每个标记对应一个 大小为 50257 的对数概率向量(logits vector)，其中数值最高的索引即为下一个预测的标记。

正如我们所见，输出张量的形状为 [2, 4, 50257]，因为我们输入了 2 条文本数据，每条包含 4 个标记(tokens)。

• 最后一维(50257) 对应于标记器的词汇表大小，表示 每个标记都有一个大小为 50257 的对数概率(logits) 向量。
• 在下一节，我们将 把这些 50257 维的输出向量转换回标记(tokens)。

计算 GPT-2 的参数总数

在继续转换模型输出为文本之前，我们先花些时间分析 GPT-2 的模型大小。使用 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

为什么 GPT-2 不是 1.24 亿参数，而是 1.63 亿参数？

• 我们之前提到 GPT-2 的参数量为 1.24 亿，但 实际参数数为 1.63 亿，为什么会有差异？
• 这是由于 “权重共享(Weight Tying)” 这一概念，该技术被用于 原始$GPT$-2架构。
• 在 GPT-2 中，输出层复用了标记嵌入层(Token Embedding Layer) 的权重，这意味着实际存储的参数量要少于模型结构计算出的总参数数。

验证 Token 嵌入层 和 输出层 形状

为了理解 权重共享，我们打印 模型的标记嵌入层(tok_emb) 和 输出层(out_head) 的形状：

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, 768]，其中：
  • 50257：标记器的 词汇表大小。
  • 768：标记嵌入的 维度。

这说明 out_head(输出层) 的参数实际上与 tok_emb(标记嵌入层) 的参数是等价的，可以 共享同一组权重。

计算去除输出层后，模型的实际可训练参数量

根据权重共享的原则，我们从总参数数中移除输出层的参数：

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 模型的大小匹配。

权重共享的意义

• “权重共享(weight tying)” 减少了模型的整体内存占用(memory footprint) 和计算复杂度(computational complexity)。
• 但在实践中，分开使用独立的 Token 嵌入层 和 输出层，通常会得到更好的训练效果和模型性能。
• 现代 $LLM$（大型语言模型） 通常不会使用权重共享，而是 使用独立的标记嵌入层和输出层，正如我们在 GPTModel 的实现 中所做的那样。

后续章节
我们将在 第6章 重新讨论 “权重共享(weight tying)”，并在加载 $OpenAI$ 预训练的$GPT$-2权重 时实现该技术。

---

练习 4.1：计算前馈网络和多头注意力模块的参数数量

计算并比较：

• 前馈网络(Feed Forward Module) 中的参数数量。
• 多头注意力模块(Multi-Head Attention Module) 中的参数数量。

答案

block = TransformerBlock(GPT_CONFIG_124M)
print(block)

total_params = sum(p.numel() for p in block.ff.parameters())
print(f"Total number of parameters in feed forward module: {total_params:,}")
total_params = sum(p.numel() for p in block.att.parameters())
print(f"Total number of parameters in attention module: {total_params:,}")

TransformerBlock(
  (att): MultiHeadAttention(
    (W_query): Linear(in_features=768, out_features=768, bias=False)
    (W_key): Linear(in_features=768, out_features=768, bias=False)
    (W_value): Linear(in_features=768, out_features=768, bias=False)
    (out_proj): Linear(in_features=768, out_features=768, bias=True)
    (dropout): Dropout(p=0.1, inplace=False)
  )
  (ff): FeedForward(
    (layers): Sequential(
      (0): Linear(in_features=768, out_features=3072, bias=True)
      (1): GELU()
      (2): Linear(in_features=3072, out_features=768, bias=True)
    )
  )
  (norm1): LayerNorm()
  (norm2): LayerNorm()
  (drop_shortcut): Dropout(p=0.1, inplace=False)
)
Total number of parameters in feed forward module: 4,722,432
Total number of parameters in attention module: 2,360,064

---

计算 $GPTModel$ 的存储需求
接下来，我们计算 $GPTModel$（1.63 亿参数）所需的内存：

total_size_bytes = total_params * 4  # A: 每个参数占用 4 字节（32-bit float）
total_size_mb = total_size_bytes / (1024 * 1024)  # B: 转换为 MB
print(f"Total size of the model: {total_size_mb:.2f} MB")

输出结果：

Total size of the model: 621.83 MB

结论

• 假设每个参数都是 32-bit 浮点数（即 4 字节），则 1.63 亿个参数的 $GPTModel$ 需要 621.83 MB 的存储空间。
• 这说明，即使是 较小的$LLM$，也需要相对较大的存储空间。

在本节中，我们：

• 实现了 $GPTModel$ 架构，并观察到其输出是形状为 [batch_size, num_tokens, vocab_size] 的数值张量。
• 计算了模型的存储需求，并理解了即使是 小规模的$LLM$，也需要较大的存储容量。
• 在 下一节，我们将 编写代码，将这些输出张量转换为文本。

---

练习 4.2：初始化更大规模的 GPT 模型

本章中，我们初始化了 1.24 亿参数的$GPT$模型（GPT-2 Small）。
现在，不修改代码结构，仅修改配置文件，使用 GPTModel 实现以下更大规模的 GPT-2 版本：

GPT-2 版本	嵌入维度	变换器块数量	多头注意力头数
GPT-2 Medium	1024	24	16
GPT-2 Large	1280	36	20
GPT-2 XL	1600	48	25

额外任务：计算 每个$GPT$模型的总参数量。

答案

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
}

def get_config(base_config, model_name="gpt2-small"):
    GPT_CONFIG = base_config.copy()

    if model_name == "gpt2-small":
        GPT_CONFIG["emb_dim"] = 768
        GPT_CONFIG["n_layers"] = 12
        GPT_CONFIG["n_heads"] = 12

    elif model_name == "gpt2-medium":
        GPT_CONFIG["emb_dim"] = 1024
        GPT_CONFIG["n_layers"] = 24
        GPT_CONFIG["n_heads"] = 16

    elif model_name == "gpt2-large":
        GPT_CONFIG["emb_dim"] = 1280
        GPT_CONFIG["n_layers"] = 36
        GPT_CONFIG["n_heads"] = 20

    elif model_name == "gpt2-xl":
        GPT_CONFIG["emb_dim"] = 1600
        GPT_CONFIG["n_layers"] = 48
        GPT_CONFIG["n_heads"] = 25

    else:
        raise ValueError(f"Incorrect model name {model_name}")

    return GPT_CONFIG


def calculate_size(model): # based on chapter code
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total number of parameters: {total_params:,}")

    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:,}")
    
    # Calculate the total size in bytes (assuming float32, 4 bytes per parameter)
    total_size_bytes = total_params * 4
    
    # Convert to megabytes
    total_size_mb = total_size_bytes / (1024 * 1024)
    
    print(f"Total size of the model: {total_size_mb:.2f} MB")

from gpt import GPTModel


for model_abbrev in ("small", "medium", "large", "xl"):
    model_name = f"gpt2-{model_abbrev}"
    CONFIG = get_config(GPT_CONFIG_124M, model_name=model_name)
    model = GPTModel(CONFIG)
    print(f"\n\n{model_name}:")
    calculate_size(model)

下载 gpt.py 到相同目录中。

---

4.7 生成文本

在本章的最后一节，我们将 实现代码，将$GPT$模型的张量输出转换回文本。在开始之前，让我们简要回顾 生成式模型(Generative Model)（如$LLM$）是如何 逐步生成文本的，如 图4.16 所示。

图4.16 展示了 $LLM$ 逐步生成文本的过程：

• 模型一次生成一个标记(token)。
• 以 初始输入上下文(initial input context) “Hello, I am” 为例：
  1. 第一步：模型预测 下一个标记，添加 “a”。
  2. 第二步：模型再预测 下一个标记，添加 “model”。
  3. 第三步：继续预测 下一个标记，添加 “ready”。
• 最终，模型逐步构建完整的句子。

图4.16 直观地展示了 $GPT$模型 在给定输入上下文（如 “Hello, I am”）后，逐步生成文本的过程。

• 每次迭代，输入上下文都会增长，使得模型能够 生成连贯且符合上下文的文本。
• 到 第6次迭代，模型已经 构造出完整的句子：
  “Hello, I am a model ready to help.”

在 上一节，我们看到 GPTModel 的当前实现 输出形状为 [batch_size, num_tokens, vocab_size] 的张量。现在的问题是：$GPT$模型如何从这些输出张量生成如图4.16所示的文本？

文本生成的关键步骤： $GPT$模型从输出张量到最终生成文本的过程涉及多个步骤，如 图4.17 所示：

1. 解码输出张量（将 logits 转换为概率分布）。
2. 基于概率分布选择标记（可采用贪心搜索、采样、Top-k采样等策略）。
3. 将选定的标记转换为可读文本。

图4.17 详细说明了 $GPT$模型文本生成的内部机制，展示了 单个标记生成过程(iteration)：

• 首先，将 输入文本编码为标记ID。
• 然后，将标记ID 输入到$GPT$模型。
• 最后，模型的输出被 转换回文本，并 追加到原始输入文本，形成新的上下文。

图4.17 详细展示了 GPT模型的“下一个标记生成”过程，即 模型如何在每一步预测下一个标记。

单步标记生成过程： 在 每个生成步骤 中：

1. $GPT$模型输出一个矩阵，其中包含表示 潜在下一个标记的向量。
2. 提取下一个标记对应的向量，并使用 Softmax函数 将其转换为 概率分布。
3. 在概率分布中找到最大值对应的索引，该索引即 预测的标记ID。
4. 将该标记ID解码回文本，得到 序列中的下一个标记。
5. 将该标记追加到之前的输入序列，形成 新的输入上下文，用于下一次迭代。

通过这种 逐步扩展输入序列 的方式，模型可以 顺序生成文本，从初始上下文构建 连贯的短语和句子。

循环迭代生成完整文本： 在实际应用中，我们会 重复此过程（如 图4.16 所示），直到：

• 达到用户设定的标记数(token count)。
• 遇到特殊终止标记(如<EOS>)。

我们可以使用以下代码 实现GPT模型的标记生成过程：

代码清单 4.8 GPT模型生成文本的函数

def generate_text_simple(model, idx, max_new_tokens, context_size):  # A
    """
    - model：GPT 语言模型。  
    - idx：当前输入标记序列（张量格式）。  
    - max_new_tokens：要生成的新标记数。  
    - context_size：模型最大可接受的上下文窗口。  
    """
    for _ in range(max_new_tokens):
        ''' 取输入序列的最后 context_size 个标记，确保输入序列不会超出模型最大上下文限制。 '''
        idx_cond = idx[:, -context_size:]  # B
        with torch.no_grad():
            '''
            model(idx_cond) 计算所有时间步的 logits，其中：
            - 形状为 [batch_size, num_tokens, vocab_size]。  
            '''
            logits = model(idx_cond)
        '''取最后一个token(-1) 的 logits ，即预测下一个标记的概率分布。'''
        logits = logits[:, -1, :]  # C
        '''将 logits 转换为概率分布(probas)。 '''
        probas = torch.softmax(logits, dim=-1)  # D
        '''
        选取最高概率的标记索引，即下一个预测标记ID。  
        - 该方法被称为 “贪心解码(Greedy Decoding)”。
        '''
        idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # E
        '''将新标记追加到输入序列，用于下一轮预测。'''
        idx = torch.cat((idx, idx_next), dim=1)  # F
    return idx

上面的代码片段展示了 使用$PyTorch$实现的语言模型生成循环。
该代码 迭代生成指定数量的新标记(tokens)，并且：

1. 裁剪当前上下文 以适配模型的 最大上下文窗口大小。
2. 计算预测(logits) 并选择 下一个标记。

详细步骤

• (A) generate_text_simple 函数：

  • 输入：
    • model：$GPT$ 语言模型。
    • idx：当前 输入标记序列（张量格式）。
    • max_new_tokens：要生成的新标记数。
    • context_size：模型最大可接受的上下文窗口。
  • 输出：扩展后的新标记序列。

• (B) idx[:, -context_size:]：

  • 取 输入序列的最后 context_size 个标记，确保 输入序列不会超出模型最大上下文限制。

• (C) logits[:, -1, :]：

  • model(idx_cond) 计算 所有时间步的 logits，其中：
    • 形状为 [batch_size, num_tokens, vocab_size]。
  • 取 最后一个token(-1) 的 logits，即 预测下一个标记的概率分布。

• (D) torch.softmax(logits, dim=-1)：

  • 将 logits 转换为概率分布(probas)。

• (E) torch.argmax(probas, dim=-1, keepdim=True)：

  • 选取 最高概率的标记索引，即 下一个预测标记ID。
  • 该方法被称为 贪心解码(Greedy Decoding)。

• (F) torch.cat((idx, idx_next), dim=1)：

  • 将新标记追加到输入序列，用于下一轮预测。

关于 Softmax 计算的额外说明
在 generate_text_simple 函数中，我们使用 torch.softmax() 将 logits 转换为概率分布，并用 torch.argmax() 选取最高概率的标记。

但实际上：

• Softmax 是单调函数(monotonic function)，即它不会改变输入的相对大小顺序。
• 因此，直接对 logits 应用 torch.argmax() 也能得到相同的结果。

尽管如此，我们仍然 显式地展示 Softmax 计算步骤，以便直观理解 logits → 概率 → 选取最高概率标记 这一完整流程。此外，在 下一章($GPT$训练代码部分)，我们会引入 更复杂的采样策略，比如：

• 随机采样(Stochastic Sampling)
• 温度缩放(Temperature Scaling)
• Top-k采样(Top-k Sampling)
• 核采样(Nucleus Sampling, Top-p Sampling)

这些策略可以 避免模型始终选择最高概率标记，从而 提高生成文本的多样性 和 创造性。

这个 逐步生成标记ID 并将其追加到上下文的过程，使用 generate_text_simple 函数实现，并在 图4.18 中进一步说明。
（每次迭代中的 标记ID生成过程 详见 图4.17。）

图4.18 展示了 六次迭代的标记预测循环，其中：

• 模型接收初始标记ID序列作为输入。
• 预测下一个标记 并 将其追加到输入序列，形成新的输入序列，用于下一次迭代。
• （为便于理解）标记ID 也被转换为对应的文本。

如 图4.18 所示，我们以迭代方式生成标记ID。例如，在 第1次迭代 时：

• 模型接收 代表 "Hello, I am" 的 输入标记。
• 预测下一个标记（ID为 257，对应 "a"）。
• 将预测出的标记追加到输入序列，形成新的输入。

该过程不断重复，直到 第6次迭代，模型最终生成完整的句子：
“Hello, I am a model ready to help.”

我们现在尝试 将 "Hello, I am" 作为模型输入，并使用 generate_text_simple 进行文本生成，如 图4.18 所示。

首先，将输入文本编码为标记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)

编码后的标记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]))

生成的标记ID如下：

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

将标记ID解码回文本

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

模型输出的文本如下：

Hello, I am Featureiman Byeswickattribute argue

---

为什么生成的文本是无意义的？
从上面的输出可以看出，模型生成的文本是杂乱无章的，完全不像 图4.18 中的连贯文本。

原因：模型尚未经过训练！

• 目前，我们只是实现了$GPT$架构，并实例化了一个$GPT$模型，但其参数仍然是随机初始化的。
• 模型训练是一个庞大的主题，我们将在下一章深入探讨如何训练$GPT$模型。

---

练习 4.3：使用不同的 dropout 参数

在本章的开头，我们在 GPT_CONFIG_124M 字典中定义了 全局 drop_rate 参数，用于设置 $GPTModel$ 结构中的所有 dropout 层。

修改代码，使不同的 dropout 层使用独立的 dropout 值：
（提示：模型架构中有三个不同的 dropout 层）：

1. 嵌入层(Embedding Layer) 的 dropout。
2. 捷径连接(Shortcut Layer) 的 dropout。
3. 多头注意力模块(Multi-Head Attention Module) 的 dropout。

任务：

• 修改 GPT_CONFIG_124M，为不同 dropout 位置指定不同的值。
• 更新 GPTModel 代码，使其分别使用这些不同的 dropout 设置。

答案

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate_emb": 0.1,        # NEW: dropout for embedding layers
    "drop_rate_attn": 0.1,       # NEW: dropout for multi-head attention  
    "drop_rate_shortcut": 0.1,   # NEW: dropout for shortcut connections  
    "qkv_bias": False
}
import torch.nn as nn
from gpt import MultiHeadAttention, LayerNorm, FeedForward


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_attn"], # NEW: dropout for multi-head attention
            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_shortcut"])

    def forward(self, x):
        # Shortcut connection for attention block
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        # Shortcut connection for feed-forward block
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        return x


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_emb"]) # NEW: dropout for embedding layers

        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  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits
import torch

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

4.8 总结

• 层归一化(Layer Normalization) 通过确保每一层的输出具有一致的 均值(mean) 和 方差(variance) 来 稳定训练过程。
• 捷径连接(Shortcut Connections) 通过 跳过一个或多个层，将某一层的输出直接传递到更深的层，有助于缓解深度神经网络（如$LLM$）的梯度消失问题(vanishing gradient problem)。
• 变换器块(Transformer Blocks) 是 $GPT$模型的核心结构组件，结合了掩码多头注意力模块(masked multi-head attention) 和 使用GELU激活函数的全连接前馈网络。
• $GPT$模型是由多个重复的变换器块组成的$LLM$，参数量从百万到数十亿。
• $GPT$模型有多个不同的规模，例如 124、345、762 和 1542 百万参数，这些都可以使用相同的 GPTModel Python 类进行实现。
• $GPT$类$LLM$的文本生成能力 依赖于 逐步预测单个标记(token) 并解码输出张量为可读文本，基于 给定的输入上下文 进行生成。
• 未经训练的$GPT$模型生成的文本是无意义的，这强调了 模型训练对生成连贯文本的重要性，这一主题将在后续章节中深入探讨。