【深度学习】实验 — 动手实现 GPT【三】：LLM架构、LayerNorm、GELU激活函数

模型定义

编码一个大型语言模型（LLM）架构

• 像 GPT 和 Llama 这样的模型是基于原始 Transformer 架构的解码器部分，按顺序生成词。
• 因此，这些 LLM 通常被称为“类似解码器”的 LLM。
• 与传统的深度学习模型相比，LLM 更大，主要原因在于其庞大的参数数量，而非代码量。
• 我们会看到，在 LLM 架构中许多元素是重复的。

• 我们考虑的嵌入和模型大小类似于小型 GPT-2 模型。

• 我们将具体实现最小的 GPT-2 模型（1.24 亿参数）的架构，参考 Radford 等人发表的 Language Models are Unsupervised Multitask Learners（注意，最初报告中列出该模型参数量为 1.17 亿，但模型权重库后来更正为 1.24 亿）。

• 后续部分将展示如何将预训练权重加载到我们的实现中，以支持 3.45 亿、7.62 亿和 15.42 亿参数的模型大小。

• 1.24亿参数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
}

• 我们使用简短的变量名，以避免代码中出现过长的行。
• "vocab_size" 表示词汇表大小为 50,257，由 BPE 分词器支持。
• "context_length" 表示模型的最大输入词元数量，由位置嵌入实现。
• "emb_dim" 是输入词元的嵌入维度，将每个输入词元转换为 768 维向量。
• "n_heads" 是多头注意力机制中的注意力头数。
• "n_layers" 是模型中的 Transformer 块数量。
• "drop_rate" 是 dropout 机制的强度，在第 3 章中讨论过；0.1 表示在训练过程中丢弃 10% 的隐藏单元，以减轻过拟合。
• "qkv_bias" 决定多头注意力机制中的 Linear 层在计算查询（Q）、键（K）和值（V）张量时是否包含偏置向量；我们将禁用此选项，这是现代 LLM 的标准做法。

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

• 层归一化（LayerNorm），也称为层归一化，Ba 等人，2016 提出，旨在将神经网络层的激活值中心化为 0 均值，并将其方差归一化为 1。
• 这有助于稳定训练过程，并加快有效权重的收敛速度。
• 层归一化在 Transformer 块内的多头注意力模块之前和之后应用，稍后我们会实现；此外，它也应用在最终输出层之前。
  
• 让我们通过一个简单的神经网络层传递一个小的输入样本，来看看层归一化的工作原理：

# create 2 training examples with 5 dimensions (features) each
batch_example = torch.randn(2, 5)

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

输出

tensor([[0.0000, 0.0000, 0.1504, 0.2049, 0.0694, 0.0000],
        [0.0000, 0.0000, 0.1146, 0.3098, 0.0939, 0.5742]],
       grad_fn=<ReluBackward0>)

• 让我们计算上面2个输入中每一个的均值和方差：

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.3448],
        [0.2182]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[0.0791],
        [0.2072]], grad_fn=<VarBackward0>)

• 归一化独立应用于每个输入（行）；使用 dim=-1 会在最后一个维度（此处为特征维度）上执行计算，而不是在行维度上执行。

• 减去均值并除以方差（标准差）的平方根，使输入在列（特征）维度上具有 0 的均值和 1 的方差：

out_norm = (out - mean) / torch.sqrt(var)
print("Normalized layer outputs:\n", out_norm)

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

输出

Normalized layer outputs:
 tensor([[ 1.9920, -0.1307, -0.3069, -0.7573, -0.2769, -0.5201],
        [-0.4793, -0.4793, -0.4793, -0.1003,  2.0176, -0.4793]],
       grad_fn=<DivBackward0>)
Mean:
 tensor([[-9.9341e-09],
        [ 4.5945e-08]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[1.0000],
        [1.0000]], grad_fn=<VarBackward0>)

• 每个输入都以 0 为中心，方差为 1；为了提高可读性，我们可以禁用 PyTorch 的科学计数法：

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

• 上面我们对每个输入的特征进行了归一化。
• 现在，基于相同的思想，我们可以实现一个 LayerNorm 类：

LayerNorm代码实现

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):
        """
        args:
            x: torch.Tensor
                The input tensor
        returns:
            norm_x: torch.Tensor
                The normalized tensor
        Step:
            1. Compute the mean and variance separately
            2. Normalize the tensor
            3. Scale and shift the tensor
            4. Return the normalized tensor
        """
        # complete this section (3/10)
        # 1. 计算每个特征的均值和方差
        mean = x.mean(dim=-1,keepdim=True)
        variance = x.var(dim=-1,keepdim=True,unbiased=False)
        
        # 2. 对张量进行归一化处理
        x_normalized = (x - mean) / torch.sqrt(variance + self.eps)
        
        # 3. 缩放并平移张量
        norm_x = self.scale * x_normalized + self.shift
        
        # 4. 返回归一化后的张量
        return norm_x

scale和shift

• 注意，除了通过减去均值并除以方差来执行归一化外，我们还添加了两个可训练的参数：scale 和 shift。
• 初始的 scale（乘以 1）和 shift（加 0）值不会产生任何效果；但是，scale 和 shift 是可训练的参数，LLM 会在训练期间自动调整它们，以提高模型在训练任务中的表现。
• 这使得模型可以学习适合其处理数据的适当缩放和偏移。
• 另外，在计算方差的平方根之前我们添加了一个较小的值（eps），以避免方差为 0 时的除零错误。

有偏方差

• 在上述方差计算中，设置 unbiased=False 意味着使用公式 $\frac{{\sum }_i(x_i-\bar{x}{)}^2}{n}$ 计算方差，其中 n 为样本大小（在这里为特征或列数）；此公式不包含贝塞尔校正（其分母为 n-1），因此提供了方差的有偏估计。

• 对于嵌入维度 n 很大的 LLM，使用 n 和 n-1 之间的差异可以忽略不计。

• 然而，GPT-2 的归一化层是在有偏方差下训练的，因此为了与我们将在后续章节加载的预训练权重兼容，我们也采用了这种设置。

• 现在让我们实际尝试 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)

输出

Mean:
 tensor([[    -0.0000],
        [    -0.0000]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[0.9999],
        [1.0000]], grad_fn=<VarBackward0>)

实现带有 GELU 激活的前馈网络

• GELU（Hendrycks 和 Gimpel, 2016）可以通过多种方式实现；其精确版本定义为 $\text{GELU}(x)=x\cdot \Phi (x)$，其中 $\Phi (x)$ 是标准高斯分布的累积分布函数。
• 实际中，通常使用计算成本较低的近似实现：$\text{GELU}(x)\approx 0.5\cdot x\cdot \left(1+\tanh \left[\sqrt{\frac{2}{\pi }}\cdot \left(x+0.044715\cdot x^3\right)\right]\right)$（原始的 GPT-2 模型也是在这种近似下训练的）。

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

    def forward(self, x):
        """
        args:
            x: torch.Tensor
                The input tensor
        returns:
            torch.Tensor
                The output tensor
        """
        # Complete this section (4/10)
        # Approximate GELU using the tanh-based formula
        return 0.5 * x * (1 + torch.tanh((torch.sqrt(torch.tensor(2 / 3.1415)) * (x + 0.044715 * torch.pow(x, 3)))))

import matplotlib.pyplot as plt

gelu, relu = GELU(), nn.ReLU()

# Some sample data
x = torch.linspace(-3, 3, 100)
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()

输出

• 接下来，让我们实现一个小型神经网络模块 FeedForward，稍后将在 LLM 的 Transformer 块中使用：

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        """
        implement self.layers as a Sequential model with:
            1. Linear layer with input dimension cfg["emb_dim"] and output dimension 4*cfg["emb_dim"]
            2. GELU activation function
            3. Linear layer with input dimension 4*cfg["emb_dim"] and output dimension cfg["emb_dim"]
        """
        # complete this section (5/10)
        
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),  # 1. 线性层，输入维度 cfg["emb_dim"]，输出 4*cfg["emb_dim"]
            GELU(),                                          # 2. 使用 GELU 激活函数
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"])    # 3. 线性层，输入维度 4*cfg["emb_dim"]，输出 cfg["emb_dim"]
        )
        
    def forward(self, x):
        return self.layers(x)

print(GPT_CONFIG_124M["emb_dim"])

输出

768

测试

ffn = FeedForward(GPT_CONFIG_124M)

# input shape: [batch_size, num_token, emb_size]
x = torch.rand(2, 3, 768)
out = ffn(x)
print(out.shape)

输出

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