从零实现ChatGPT:第三章实现大型语言模型（1）

准备深入学习transformer，并参考一些资料和论文实现一个大语言模型，顺便做一个教程，今天是第三部分。

本系列禁止转载，主要是为了有不同见解的同学可以方便联系我，我的邮箱 fanzexuan135@163.com

实现大型语言模型

这一章我们将从头实现一个大型语言模型。我们将从下一节开始，先从模型架构的顶层视图开始，然后再详细介绍各个组件。

编码一个LLM架构

LLMs如GPT(Generative Pre-trained Transformer)是大型深度神经网络架构，旨在生成新文本，一次生成一个单词(或token)。然而，尽管它们的规模很大，但模型架构并没有你想象的那么复杂，因为它的许多组件是重复的，我们将在后面看到。图1提供了一个类似GPT的LLM的自顶向下视图，并突出显示了它的主要组件。

图1：GPT模型的心智模型。在嵌入层旁边，它由一个或多个transformer块组成，其中包含我们在前一章实现的带掩码的多头注意力模块。

如图1所示，我们已经涵盖了几个方面，如输入tokenization和嵌入，以及带掩码的多头注意力模块。本章的重点将是实现GPT模型的核心结构，包括它的transformer块，我们将在下一章中对其进行训练以生成类人文本。

在前面的章节中，为了简单起见，我们使用了较小的嵌入维度，以确保概念和示例能够舒适地放在一页上。现在在本章中，我们正在扩展到一个小型GPT模型的大小，特别是具有1.1亿参数的最小版本，如Radford等人的论文"Language Models are Unsupervised Multitask Learners"中所述。请注意，虽然原始报告提到1.17亿个参数，但这后来被更正了。

第6章将重点介绍将预训练的权重加载到我们的实现中，并使其适应具有3.45亿和7.62亿参数的更大的GPT模型。在深度学习和像GPT这样的LLMs的上下文中，术语"参数"是指模型的可训练权重。这些权重本质上是模型的内部变量，在训练过程中进行调整和优化，以最小化特定的损失函数。这种优化允许模型从训练数据中学习。

例如，在由2048 x 4096维矩阵(或张量)表示的神经网络层中，该矩阵的每个元素都是一个参数。因为有2048行和4096列，所以该层中的总参数数是2048乘以4096，等于8,388,608个参数。

GPT 与 GPT-3

请注意，我们专注于GPT，因为OpenAI已经公开提供了预训练模型的权重，我们将在第6章中将其加载到我们的实现中。GPT-3在模型架构方面基本上是相同的，只是从GPT-2的15亿参数扩展到GPT-3的1750亿参数，并在更多数据上进行训练。截至撰写本文时，GPT-3的权重尚未公开。GPT也是学习如何实现LLMs的更好选择，因为它可以在单个笔记本电脑上运行，而GPT-3需要GPU集群进行训练和推理。根据Lambda Labs的说法，在单个V100数据中心GPU上训练GPT-3需要355年，在消费级的RTX 3090 GPU上需要665年。

我们通过以下Python字典指定小型GPT模型的配置，我们将在后面的代码示例中使用它：

GPT_CONFIG_SM = {
    "vocab_size": 50257,  # 词汇表大小
    "context_length": 1024,  # 上下文长度 
    "emb_dim": 768,  # 嵌入维度
    "n_heads": 12,  # 注意力头数
    "n_layers": 12,  # 层数
    "dropout_rate": 0.1,  # Dropout率
    "qkv_bias": False  # Query-Key-Value偏置
}

在GPT_CONFIG_SM字典中，我们使用简洁的变量名以保持清晰和防止代码行过长。

• vocab_size指的是第4章中BPE分词器使用的50,257个单词的词汇表。
• context_length表示模型可以通过第3章讨论的位置嵌入处理的输入token的最大数量。
• emb_dim表示将每个token转换为768维向量的嵌入大小。
• n_heads表示第5章实现的多头注意力机制中的注意力头数。
• n_layers指定模型中的transformer块数，这将在后面的章节中详细阐述。
• dropout_rate表示dropout机制的强度(0.1意味着隐藏单元的10%被丢弃)，以防止过拟合，如第4章所述。
• qkv_bias确定是否在多头注意力的Linear层中为query、key和value计算包含偏置向量。我们最初将禁用此功能，遵循现代LLMs的规范，但在第6章中，当我们将OpenAI的预训练GPT-2权重加载到我们的模型中时，我们将重新讨论它。

使用上面的配置，我们将在本节开始实现一个GPT占位符架构(DummyGPTModel)，如图2所示。这将为我们提供一个全局视图，说明一切如何组合在一起，以及我们需要在接下来的部分中编写什么其他组件来组装完整的GPT模型架构。

图2：概述了我们在本章中编码GPT架构的顺序。在本章中，我们将从GPT主干占位符架构开始，然后转到各个核心部分，最终在最终的GPT架构中组装它们。

图2中显示的编号框说明了我们处理本章中实现最终GPT架构所需的各个概念的顺序。我们将从步骤1开始，一个我们称之为DummyGPTModel的GPT主干占位符。

代码清单：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['dropout_rate'])
        self.trf_blocks = nn.Sequential(
            *[DummyTransformerBlock(cfg) for _ in range(cfg['n_layers'])])
        self.final_norm = DummyLayerNorm(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
        
class DummyTransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
    
    def forward(self, x):
        return x
        
class DummyLayerNorm(nn.Module): 
    def __init__(self, normalized_shape, eps=1e-05):
        super().__init__()
            
    def forward(self, x):
        return x

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

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

上面的代码已经是可运行的，正如我们稍后在本节中准备好输入数据后将看到的那样。但是，请注意，在上面的代码中，我们使用了占位符(DummyLayerNorm和DummyTransformerBlock)来表示transformer块和层归一化，我们将在后面的部分中开发这些组件。

接下来，我们将准备输入数据并初始化一个新的GPT模型来说明它的用法。基于我们在第4章中看到的图，其中我们编码了分词器，图3提供了关于数据如何流入和流出GPT模型的高级概述。

图3：一个大图概述，显示输入数据如何被分词、嵌入并馈送到GPT模型中。请注意，在我们之前编码的DummyGPTClass中，token嵌入是在GPT模型内部处理的。在LLMs中，嵌入的输入token维度通常与输出维度相匹配。此处的输出嵌入表示我们在第3章中讨论的上下文向量。

为了实现图3所示的步骤，我们使用第4章中介绍的tiktoken分词器对两个文本输入的批次进行分词，以用于GPT模型。

import tiktoken

tokenizer = tiktoken.get_encoding("gpt2")

batch = []
txt0 = "Every effort moves you"
txt1 = "Every day holds a"

batch.append(torch.tensor(tokenizer.encode(txt0)))
batch.append(torch.tensor(tokenizer.encode(txt1)))
batch = torch.stack(batch, dim=0)
print(batch) 

这两个文本的结果token ID如下：

tensor([[11298,  1909,  8720,   379],
        [11298,  1391,  6410, 10851]])

接下来，我们初始化一个新的1.1亿参数DummyGPTModel实例，并将分词后的批次输入其中：

torch.manual_seed(1)
model = DummyGPTModel(GPT_CONFIG_SM)
logits = model(batch)

print("Output shape:", logits.shape)
print(logits)

模型输出(通常称为logits)如下：

Output shape: torch.Size([2, 4, 50257])
tensor([[[ -0.0104,  -0.0123,  -0.0040,  ...,  -0.0017,   0.0145,  -0.0003],
         [ -0.0019,  -0.0044,   0.0025,  ...,  -0.0049,   0.0038,  -0.0129],
         [ -0.0052,   0.0010,   0.0104,  ...,  -0.0147,   0.0010,   0.0025],
         [ -0.0012,  -0.0026,  -0.0091,  ...,   0.0018,   0.0127,  -0.0082]],

        [[ -0.0039,   0.0124,  -0.0132,  ...,   0.0093,   0.0039,  -0.0137],
         [ -0.0014,   0.0078,   0.0145,  ...,  -0.0103,  -0.0062,  -0.0099],
         [  0.0086,  -0.0138,  -0.0028,  ...,   0.0175,   0.0049,  -0.0091],
         [ -0.0050,   0.0096,   0.0013,  ...,   0.0040,  -0.0030,   0.0139]]],
       grad_fn=<UnsafeViewBackward0>)

输出张量有两行，对应于两个文本样本。每个文本样本由4个token组成，每个token是一个50,257维的向量，与分词器的词汇表大小相匹配。

这个嵌入是50,257维的，因为这些维度中的每一个都指的是词汇表中的一个唯一token。在本章结束时，当我们实现后处理代码时，我们将把这些50,257维向量转换回token ID，然后我们可以将其解码成单词。

现在我们已经从顶层看了GPT架构及其输入和输出，我们将在接下来的部分中编写各个占位符，从上一个代码中的DummyLayerNorm替换为真正的层归一化类开始。

用层归一化来规范化激活

训练具有许多层的深度神经网络有时可能具有挑战性，由于诸如梯度消失或爆炸等问题。这些问题导致不稳定的训练动态，并使网络难以有效地调整其权重，这意味着学习过程很难找到一组神经网络的参数(权重)集，以最小化损失函数。换句话说，网络难以学习数据中的潜在模式，以使其能够做出准确的预测或决策。(如果你不熟悉神经网络训练和梯度的概念，可以在附录A中找到对这些概念的简要介绍：A.4 自动微分变得简单。然而，对梯度的深入数学理解并不是理解本书内容所必需的。)

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

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

在我们用代码实现层归一化之前，图4提供了一个关于层归一化如何工作的视觉概述。

图4：层归一化的图解，其中层输出(也称为激活)被归一化，使其具有零均值和1的方差。

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

torch.manual_seed(42)

batch_example = torch.randn(2, 5)

layer = nn.Sequential(nn.Linear(5, 3), nn.ReLU())
out = layer(batch_example)
print(out)

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

tensor([[1.1575, 0.0000, 0.9309],
        [0.1584, 0.0000, 0.0202]], grad_fn=<ReluBackward0>)

我们编码的神经网络层由一个Linear层组成，后面是一个非线性激活函数ReLU(Rectified Linear Unit的缩写)，它是神经网络中的标准激活函数。如果你不熟悉ReLU，它简单地将负输入阈值化为0，确保层只输出正值，这解释了为什么结果层输出不包含任何负值。(请注意，我们将在GPT中使用另一个更复杂的激活函数，我们将在下一节中介绍它。)

在我们将层归一化应用于这些输出之前，让我们检查均值和方差：

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

输出如下：

Mean:
 tensor([[0.7295],
        [0.0595]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[0.3869],
        [0.0101]], grad_fn=<VarBackward0>)

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

在诸如均值或方差计算等操作中使用keepdim=True可确保输出张量保留与输入张量相同的形状，即使该操作通过dim指定的维度缩小了张量。例如，如果没有keepdim=True，返回的均值张量将是一个一维向量[0.7295, 0.0595]而不是二维矩阵[[0.7295], [0.0595]]。

dim参数指定在张量中执行计算(这里是均值或方差)的维度，如图5所示。

图5：计算张量均值时dim参数的图解。例如，如果我们有一个2D张量(矩阵)，维度为[rows, columns]，使用dim=0将在行上执行操作(如底部所示，垂直地)，导致输出聚合每列的数据。使用dim=1或dim=-1将在列上执行操作(如顶部所示，水平地)，导致输出聚合每行的数据。

如图5所解释的，对于2D张量(如矩阵)，使用dim=1进行均值或方差计算等操作与使用dim=-1相同。这是因为-1指的是张量的最后一个维度，对应于2D张量中的列。稍后，当我们将层归一化添加到生成3D张量的GPT模型中时，形状为[batch_size, num_tokens, embedding_size]，我们仍然可以使用dim=-1进行最后一个维度上的归一化，避免从dim=2更改为dim=-1。

接下来，让我们将层归一化应用于我们之前获得的层输出。该操作包括减去均值并除以方差的平方根(也称为标准差)：

out_norm = (out - mean) / torch.sqrt(var + 1e-5)

mean = out_norm.mean(dim=-1, keepdim=True)  
var = out_norm.var(dim=-1, unbiased=False, keepdim=True)
print("Normalized layer outputs:\n", out_norm)
print("Mean:\n", mean)
print("Variance:\n", var)  

正如我们从结果中看到的，经过归一化的层输出(现在也包含负值)具有0的均值和1的方差：

Normalized layer outputs:
 tensor([[ 1.1010, -1.8991,  0.5231],
        [ 0.9820, -0.5861, -0.3959]])
Mean:
 tensor([[-0.0920],
        [-0.0000]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[1.0000],
        [1.0000]], grad_fn=<VarBackward0>)

请注意，输出张量中的-0.0920值是-9.202499580672643e-08的科学记数法，以十进制形式为-0.00000009202499580672643。这个值非常接近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>)

到目前为止，在本节中，我们以一步一步的过程编码和应用了层归一化。现在让我们将这个过程封装在一个我们稍后可以在GPT模型中使用的PyTorch模块中。

代码清单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。这种方法不应用Bessel校正，通常在分母中使用n-1而不是n来调整样本方差估计中的偏差。这个决定导致了所谓的方差的有偏估计。对于嵌入维度n显著较大的大规模语言模型(LLMs)，使用n和n-1之间的差异实际上可以忽略不计。我们选择这种方法是为了确保与GPT-2模型的归一化层的兼容性，并且因为它反映了TensorFlow的默认行为，TensorFlow用于实现原始GPT-2模型。使用类似的设置可确保我们的方法与第6章中将加载的预训练权重兼容。

让我们现在在实践中尝试LayerNorm模块，并将其应用于我们之前的批次输入：

ln = LayerNorm(emb_dim=3)
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架构的构建块之一，如图2中的思维模型所示。

层归一化与批归一化

如果你熟悉批归一化，一种用于神经网络的常见和传统的归一化方法，你可能想知道它与层归一化相比如何。与跨批次维度归一化的批归一化不同，层归一化跨特征维度归一化。LLMs通常需要大量的计算资源，可用的硬件或特定的使用情况可能会在训练或推理期间决定批次大小。由于层归一化独立于批次大小归一化每个输入，因此在这些场景中提供了更大的灵活性和稳定性。当在资源受限的环境中进行分布式训练或部署模型时，这一点尤其有益。

用GELU激活实现前馈网络

在本节中，我们实现了一个小型神经网络子模块，用作LLMs中transformer块的一部分。我们首先实现GELU激活函数，它在这个神经网络子模块中扮演着关键角色。(有关在PyTorch中实现神经网络的更多信息，请参见A.3节：在PyTorch中实现多层神经网络。)

历史上，ReLU激活函数由于其简单性和在各种神经网络架构中的有效性而在深度学习中广泛使用。然而，在LLMs中，除了传统的ReLU之外，还采用了其他几种激活函数。两个值得注意的例子是GELU(高斯误差线性单元)和SwiGLU(Sigmoid加权线性单元)。GELU和SwiGLU是更复杂和平滑的激活函数，分别合并了高斯和sigmoid门控线性单元。与简单的ReLU不同，它们为深度学习模型提供了改进的性能。

GELU激活函数可以通过多种方式实现，精确版本定义为GELU(x) = x · Φ(x)，其中Φ(x)是标准高斯分布的累积分布函数。但是，在实践中，通常实现计算成本更低的近似(原始GPT-2模型也使用这种近似进行训练)：

GELU(x) ≈ 0.5x · (1 + tanh[√(2/π) · (x + 0.044715x^3)])

在代码中，我们可以将这个函数实现为PyTorch模块，如下所示：

代码清单3：GELU激活函数的实现

class GELU(nn.Module):
    def __init__(self):
        super().__init__()

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

接下来，为了了解这个GELU函数的样子以及它与ReLU函数的比较，让我们并排绘制这些函数：

import matplotlib.pyplot as plt

gelu_relu = GELU(), nn.ReLU()

x = torch.linspace(-4, 4, 100)  
y_gelu, y_relu = gelu_relu[0](x), gelu_relu[1](x)

plt.figure(figsize=(12, 4))
for i, (y, label) in enumerate(zip([y_gelu, y_relu], ['GELU', 'ReLU'])):
    plt.subplot(1, 2, i+1)
    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()

正如我们在图6中看到的结果图，ReLU是一个分段线性函数，如果输入为正，则直接输出输入，否则输出零。GELU是一个平滑的非线性函数，近似ReLU，但对负值具有非零梯度。

图6：使用matplotlib绘制的GELU和ReLU输出。x轴显示函数输入，y轴显示函数输出。

如图6所示，GELU的平滑性可以在训练期间带来更好的优化特性，因为它允许对模型的参数进行更细微的调整。相比之下，ReLU在零处有一个尖角，这有时会使优化变得更加困难，特别是在网络非常深或具有复杂架构的情况下。此外，与ReLU对任何负输入输出零不同，GELU允许对负值有一个小的非零输出。这个特性意味着，在训练过程中，接收负输入的神经元仍然可以为学习过程做出贡献，尽管程度不如正输入那么大。

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

代码清单4：前馈神经网络模块

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg['emb_dim'], 4 * cfg['emb_dim']),
            GELU(),
            nn.Linear(4 * cfg['emb_dim'], cfg['emb_dim']),
        )
        
    def forward(self, x):
        return self.layers(x)

如前面的代码所示，FeedForward模块是一个小型神经网络，由两个Linear层和一个GELU激活函数组成。在1.1亿参数的GPT模型中，它接收具有768个token的输入批次，每个token具有768的嵌入大小，通过GPT_CONFIG_SM字典，其中GPT_CONFIG_SM['emb_dim'] = 768。

图7显示了当我们向这个小型前馈神经网络输入一些输入时，嵌入大小是如何在其内部被操纵的。

图7提供了前馈神经网络中层输出之间连接的可视化概述。需要注意的是，这个神经网络可以处理可变的批次大小和输入中的token数量。然而，每个token的嵌入大小是在初始化权重时确定和固定的。

按照图7中的示例，让我们用768的token嵌入大小初始化一个新的FeedForward模块，并向其输入一个包含2个样本和3个token的批次输入：

ffn = FeedForward(GPT_CONFIG_SM)
x = torch.rand(2, 3, 768)  
out = ffn(x)
print("out.shape:", out.shape)

正如我们从代码输出中看到的，transformer块维持了其输出中的输入维度：

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

我们在本节中实现的FeedForward模块在增强模型从数据中学习和泛化的能力方面发挥着至关重要的作用。尽管这个模块的输入和输出维度相同，但它在内部通过第一个线性层将嵌入维度扩展到更高维的空间，如图8所示。这种扩展之后是非线性GELU激活，然后通过第二个线性变换收缩回原始维度。这样的设计允许探索更丰富的表示空间。

图8：前馈神经网络中层输出的扩展和收缩的图解。首先，输入通过一个因子4从768扩展到3,072个值。然后，第二层将3,072个值压缩回768维表示。

此外，输入和输出维度的一致性简化了架构，允许堆叠多个层(我们稍后会这样做)，而无需在它们之间调整维度，从而使模型更具可扩展性。

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

图9：本章中我们实现的不同概念的思维模型，黑色勾号表示我们已经介绍过的概念。

如图9所示，transformer块结合了层归一化、包括GELU激活的前馈网络和快捷连接，我们在本章前面已经介绍过了。正如我们将在下一章看到的，这个transformer块将构成我们将实现的GPT架构的主要组成部分。

在下一节中，我们将讨论在神经网络的不同层之间插入快捷连接的概念，这对于提高深度神经网络架构中的训练性能很重要。

添加快捷连接

接下来，让我们讨论快捷连接背后的概念，也称为跳跃或残差连接。最初，快捷连接是为计算机视觉中的深度网络(特别是残差网络)提出的，以缓解梯度消失的挑战。梯度消失问题是指，随着梯度在层中向后传播，梯度(指导训练期间的权重更新)逐渐变小，使得难以有效地训练早期层，如图10所示。

图10：比较了没有快捷连接(左)和有快捷连接(右)的由5层组成的深度神经网络。快捷连接涉及将一个层的输入添加到其输出中，有效地创建一条绕过某些层的替代路径。图中所示的梯度表示每一层的平均绝对梯度，我们将在后面的代码示例中计算这个值。

如图10所示，快捷连接通过跳过一个或多个层，为梯度创建了一条穿过网络的替代(更短)路径，这是通过将一个层的输出添加到后面层的输出来实现的。这就是为什么这些连接也被称为跳跃连接。它们在训练期间的反向传递中起着关键作用，以保持梯度的流动。

在下面的代码示例中，我们实现了图10中所示的神经网络，看看我们如何在forward方法中添加快捷连接。

代码清单5：一个神经网络来说明快捷连接

class ExampleDeepNeuralNetwork(nn.Module):
    def __init__(self, layer_sizes, use_shortcut):
        super().__init__()
        self.use_shortcut = use_shortcut
        self.layers = nn.ModuleList([
            # 实现层
            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层和一个GELU激活函数组成。在前向传递中，我们迭代地将输入通过这些层，并可选地添加图10中所示的快捷连接，如果self.use_shortcut属性设置为True。

让我们使用这段代码首先初始化一个没有快捷连接的神经网络。这里，每个层将被初始化为接受具有5个输入值的样本，并返回3个输出值。最后一层返回单个输出值。

layer_sizes = [5, 4, 4, 3, 2, 1]
sample_input = torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5]])

torch.manual_seed(1)  # 指定初始权重的随机种子以实现可重复性
model_without_shortcut = ExampleDeepNeuralNetwork(
    layer_sizes, use_shortcut=False)

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

def print_gradients(model, x):
    # 前向传递
    output = model(x)
    target = torch.tensor([[0.5]])  
    
    # 基于目标和输出的接近程度计算损失
    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.5)计算它们有多接近。然后，当调用loss.backward()时，PyTorch为模型中的每一层计算损失梯度。我们可以通过model.named_parameters()迭代权重参数。假设我们有一个给定层的权重参数矩阵，在这种情况下，这一层将有4*3=12个梯度值，我们打印这12个梯度值的平均绝对梯度，以获得每层一个梯度值，以便更容易地比较层之间的梯度。

简而言之，backward()方法是PyTorch中的一个方便方法，它计算模型训练所需的损失梯度，而无需我们自己实现梯度计算的数学，从而使使用深度神经网络变得更加容易。如果你不熟悉梯度和神经网络训练的概念，我建议阅读附录A中的A.4自动微分变得容易和A.5典型的训练循环部分。

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

print_gradients(model_without_shortcut, sample_input)

输出如下：

layers.0.0.weight has gradient mean of 0.016092808544635773
layers.1.0.weight has gradient mean of 0.0016744751185178757
layers.2.0.weight has gradient mean of 0.00018377752811554074
layers.3.0.weight has gradient mean of 2.1612200699942932e-05
layers.4.0.weight has gradient mean of 2.6072029490950517e-06

正如我们从print_gradients函数的输出中看到的，梯度在从最后一层(layers.4)到第一层(layers.0)的过程中变得越来越小，这是一种称为梯度消失问题的现象。

现在让我们实例化一个带有跳跃连接的模型，看看它是如何比较的：克服深度神经网络中梯度消失问题带来的限制很重要。快捷连接是非常大的模型(如LLMs)的核心构建块，它们将有助于通过在跨层的训练期间确保一致的梯度流来促进更有效的训练，我们将在下一章训练GPT-2模型时看到这一点。

在本节介绍了快捷连接之后，我们将在下一节中将之前介绍的所有概念(层归一化、GELU激活、前馈模块和快捷连接)联系在一个transformer块中，这是我们需要编写的最后一个构建块，以组装GPT-2架构。

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

在本节中，我们正在实现transformer块，它是GPT-2和其他LLM架构的基本构建块。这个块在1.1亿参数的GPT-2架构中重复了十几次，结合了我们之前介绍过的几个概念：多头注意力、层归一化、dropout、前馈层和GELU激活，如图11所示。在下一节中，我们将把这个transformer块连接到GPT-2架构的其余部分。

图11：transformer块的图解。图的底部显示了已经嵌入到768维向量中的输入token。每一行对应于一个token的向量表示。transformer块的输出是与输入维度相同的向量，然后可以馈送到LLM中的后续层。

如图11所示，transformer块结合了几个组件，包括第5章中的带掩码的多头注意力模块和我们在3.2节中实现的FeedForward模块。

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

这个想法是，多头注意力块中的自注意力机制识别和分析输入序列中元素之间的关系。相比之下，前馈网络在每个位置单独修改数据。这种组合不仅能够对输入进行更细致入微的理解和处理，而且还提高了模型处理复杂数据模式的整体能力。

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

代码清单6：GPT-2的transformer块组件

from previous_chapters import MultiHeadAttention

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg['emb_dim'], 
            d_out=cfg['emb_dim'],
            block_size=cfg['context_length'],
            num_heads=cfg['n_heads'],
            dropout=cfg['dropout_rate'], 
            qkv_bias=cfg['qkv_bias'])
        self.ffn = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg['emb_dim'])
        self.norm2 = LayerNorm(cfg['emb_dim'])
        self.drop_resid = nn.Dropout(cfg['dropout_rate'])

    def forward(self, x):
        # A
        shortcut = x 
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_resid(x)
        x = x + shortcut  # 添加原始输入

        shortcut = x  # B
        x = self.norm2(x)
        x = self.ffn(x)
        x = self.drop_resid(x)
        x = x + shortcut  # C
        return x

给定的代码定义了一个TransformerBlock类，使用PyTorch，包括一个多头注意力机制(MultiHeadAttention)和一个前馈网络(FeedForward)，两者都根据提供的配置字典(cfg)进行配置，例如GPT_CONFIG_SM。

在这两个组件之前应用层归一化(LayerNorm)，之后应用dropout，以正则化模型并防止过拟合。这也被称为Pre-LayerNorm。较旧的架构，如原始的transformer模型，在自注意力和前馈网络之后应用层归一化，称为Post-LayerNorm，这通常会导致更差的训练动态。

该类还实现了前向传递，其中每个组件之后是一个快捷连接，将块的输入添加到其输出。这个关键特性有助于梯度在训练期间流过网络，并改善深度模型的学习，如3.3节所述。

使用我们之前定义的GPT_CONFIG_SM字典，让我们实例化一个transformer块并输入一些样本数据：

torch.manual_seed(1)
x = torch.rand(2, 3, 768)  # 示例输入

block = TransformerBlock(GPT_CONFIG_SM)
output = block(x)

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

输出如下：

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

如代码输出所示，transformer块在其输出中维护输入维度，表明transformer架构在整个网络中处理数据序列而不改变其形状。

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

通过本节中实现的transformer块，我们现在拥有了如图12所示的所有构建块，需要在下一节中实现GPT-2架构。

图12：我们目前为止在本章中实现的不同概念的思维模型。

如图12所示，transformer块结合了层归一化、包括GELU激活的前馈网络和快捷连接，我们在本章前面已经介绍过了。正如我们将在下一章看到的，这个transformer块将构成我们将实现的GPT-2架构的主要组成部分。

编写GPT-2模型

我们在本章开始时对GPT架构进行了高层次的概述，我们称之为DummyGPTModel。在这个DummyGPTModel代码实现中，我们展示了GPT-2模型的输入和输出，但它的构建块仍然是一个黑盒，使用DummyTransformerBlock和DummyLayerNorm类作为占位符。

在本节中，我们现在将DummyTransformerBlock和DummyLayerNorm占位符替换为本章后面编写的真实的TransformerBlock和LayerNorm类，以组装一个完全可操作的原始1.1亿参数版本的GPT-2。在第6章，我们将预训练一个GPT-2模型，在第7章，我们将加载来自OpenAI的预训练权重。

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

图13：GPT-2模型架构概述。该图说明了数据通过GPT-2模型的流动。从底部开始，分词后的文本首先转换为token嵌入，然后用位置嵌入增强。这个组合信息形成一个张量，通过一系列的transformer块(显示在中心，每个包含多头注意力和前馈神经网络层，有dropout和层归一化)，transformer块堆叠在一起并重复12次。

如图13所示，我们在3.4节中编写的transformer块在GPT-2模型架构中重复多次。在1.1亿参数的GPT-2模型的情况下，它重复了12次，我们通过GPT_CONFIG_SM字典中的n_layers条目指定。在具有15.5亿参数的最大GPT-2模型的情况下，这个transformer块重复了48次。

如图13所示，来自最后一个transformer块的输出然后经过最终的层归一化步骤，然后到达线性输出层。这一层将transformer的输出映射到一个高维空间(在这种情况下是50,257维，对应于模型的词汇表大小)，以预测序列中的下一个token。

现在让我们用代码实现我们在图13中看到的架构。

代码清单7：GPT-2模型架构实现

class GPT2Model(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['dropout_rate'])
        
        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg['n_layers'])])
        
        self.final_norm = LayerNorm(cfg['emb_dim'])
        self.out_head = nn.Linear(
            cfg['emb_dim'], cfg['vocab_size'], bias=False)
        
    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)  # A 
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x) 
        return logits

多亏了我们在3.4节中实现的TransformerBlock类，GPT2Model类相对较小和紧凑。

这个GPT2Model类的__init__构造函数使用通过Python字典cfg传递的配置初始化token和位置嵌入层。这些嵌入层负责将输入token索引转换为密集向量并添加位置信息，如第3章所述。

接下来，__init__方法创建一个等于cfg中指定的层数的TransformerBlock模块的序列堆栈。在transformer块之后，应用一个LayerNorm层，标准化来自transformer块的输出，以稳定学习过程。最后，定义了一个没有偏置的线性输出头，它将transformer的输出投影到分词器的词汇空间，以生成每个词汇表中token的logits。

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

让我们现在使用我们作为cfg参数传入的GPT_CONFIG_SM字典初始化1.1亿参数的GPT-2模型，并用我们在本章开头创建的批处理文本输入对其进行训练：

torch.manual_seed(1)  
model = GPT2Model(GPT_CONFIG_SM)

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

前面的代码打印输入批次的内容，然后是输出张量：

Input batch:
 tensor([[50256, 25461, 20496,   379],  # 文本1的token ID
        [50256, 11298,  6410, 10851]])  # 文本2的token ID

Output shape: torch.Size([2, 4, 50257])
tensor([[[ 0.0042,  0.0046,  0.0010,  ..., -0.0040,  0.0027,  0.0050],
         [-0.0007,  0.0073,  0.0027,  ..., -0.0052,  0.0013,  0.0052],
         [ 0.0090, -0.0012, -0.0084,  ..., -0.0043,  0.0021,  0.0063],
         [-0.0002,  0.0021,  0.0028,  ..., -0.0048,  0.0034,  0.0035]],

        [[ 0.0016,  0.0066,  0.0018,  ..., -0.0034,  0.0023,  0.0039],
         [-0.0011,  0.0082,  0.0029,  ..., -0.0048,  0.0019,  0.0047],
         [-0.0007,  0.0083,  0.0023,  ..., -0.0047,  0.0020,  0.0043],
         [-0.0018,  0.0085,  0.0023,  ..., -0.0043,  0.0021,  0.0040]]],

```python
torch.manual_seed(1)  
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.016092808544635773
layers.1.0.weight has gradient mean of 0.005305760353803635
layers.2.0.weight has gradient mean of 0.00530576407909393
layers.3.0.weight has gradient mean of 0.0017695879936218262
layers.4.0.weight has gradient mean of 0.0026072029490950513

正如我们从输出中看到的，最后一层(layers.4)仍然具有比其他层更大的梯度。然而，当我们向第一层(layers.0)前进时，梯度值稳定下来，并没有缩小到消失的小值。

快捷连接对于克服深度神经网络中梯度消失问题带来的限制很重要。快捷连接是非常大的模型(如LLMs)的核心构建块，它们将有助于通过在跨层的训练期间确保一致的梯度流来促进更有效的训练，我们将在下一章训练GPT-2模型时看到这一点。

在本节介绍了快捷连接之后，我们将在下一节中将之前介绍的所有概念(层归一化、GELU激活、前馈模块和快捷连接)联系在一个transformer块中，这是我们需要编写的最后一个构建块，以组装GPT-2架构。

下篇博客继续；