第4章：从零实现GPT模型

#从零构建大语言模型 #GPT #Transformer

学习目标

  • 弄清 LayerNorm、GELU、残差连接各自的作用。
  • 实现 Feed-Forward 子层。
  • 把第 3 章的多头注意力 + FFN 拼装为 Transformer Block。
  • 堆叠成完整的 GPT 模型，并能跑通一次前向传播。

4.1 GPT 的总体结构

一个 GPT 模型可以画成下面的”汉堡”:

输入 ids ──► Token Embed ┐
                          ├─► h0
位置 0..T-1 ─► Pos Embed ─┘
              │
              ▼
       Transformer Block × N        (本章重点)
              │
              ▼
         Final LayerNorm
              │
              ▼
          Linear → logits  (V 维)

每个 Transformer Block 的内部:

x ──► LayerNorm ──► MultiHeadAttention ──► Dropout ──► (+ x) ──► h'
h' ─► LayerNorm ──► FeedForward ────────► Dropout ──► (+ h') ──► out

注意 GPT-2 采用 Pre-LN（LayerNorm 放在子层之前），而原始 Transformer 论文是 Post-LN。Pre-LN 的训练稳定性显著更好，是后续几乎所有 LLM 的默认配置。

4.2 LayerNorm

LayerNorm 在最后一维（特征维）上做归一化：

$$\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \varepsilon}}, \quad y = \gamma \odot \hat{x} + \beta$$

其中 $\mu, \sigma$ 在每个样本、每个时间步独立计算；$\gamma, \beta$ 是可学习的逐通道缩放和偏置。

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, d_model, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(d_model))
        self.shift = nn.Parameter(torch.zeros(d_model))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x_hat = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * x_hat + self.shift

和 BatchNorm 的差别：BatchNorm 在 batch 维度统计，依赖 batch 大小且推理时要”冻结统计量”；LayerNorm 在特征维统计，单样本就能算，对序列建模特别友好。

4.3 GELU 激活

GPT 系列默认使用 GELU（Gaussian Error Linear Unit）：

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

直观理解：ReLU 是 0/1 硬开关，GELU 是用标准正态 CDF 做”软开关”，对小负值仍保留少量梯度，训练更平滑。

PyTorch 提供 nn.GELU()，也可以用近似版本：

$$\text{GELU}(x) \approx 0.5 x \left(1 + \tanh!\left[\sqrt{\tfrac{2}{\pi}} (x + 0.044715 x^3)\right]\right)$$

4.4 Feed-Forward 子层

每个 Transformer Block 中都有一个逐位置（point-wise）的两层 MLP：

class FeedForward(nn.Module):
    def __init__(self, d_model, expansion=4, dropout=0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, expansion * d_model),
            nn.GELU(),
            nn.Linear(expansion * d_model, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

“逐位置”意味着对 (B, T, C) 张量来说，时间维 T 完全没有被混合，FFN 只在通道维做非线性变换。注意力负责”位置之间”的混合，FFN 负责”通道之间”的混合——这是 Transformer 的核心结构哲学。

4.5 残差连接

残差连接是 ResNet 的发明，在 Transformer 里同样关键：

$$x_{l+1} = x_l + \text{Sublayer}(\text{LN}(x_l))$$

它带来两个好处：

  1. 梯度直通：反向传播时 $\partial x_{l+1}/\partial x_l$ 至少为 1，缓解深层网络的梯度消失；
  2. 恒等初始化：当 Sublayer 输出接近 0 时，整个 Block 退化为恒等映射，深堆叠不会破坏初始信号。

4.6 Transformer Block

把上面所有零件拼起来：

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.attn  = MultiHeadAttention(
            d_model=cfg["emb_dim"],
            num_heads=cfg["n_heads"],
            context_length=cfg["context_length"],
            dropout=cfg["drop_rate"],
        )
        self.drop1 = nn.Dropout(cfg["drop_rate"])

        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.ff    = FeedForward(cfg["emb_dim"], dropout=cfg["drop_rate"])
        self.drop2 = nn.Dropout(cfg["drop_rate"])

    def forward(self, x):
        x = x + self.drop1(self.attn(self.norm1(x)))
        x = x + self.drop2(self.ff(self.norm2(x)))
        return x

MultiHeadAttention 即第 3 章实现的版本。

4.7 GPT 主类

最后把 N 层 Block 堆起来，前后接好 embedding 和输出头：

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    = nn.Dropout(cfg["drop_rate"])

        self.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, idx):                    # idx: (B, T)
        B, T = idx.shape
        positions = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(positions)
        x = self.drop(x)
        x = self.blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)              # (B, T, vocab_size)
        return logits

GPT-2 Small（124M）的标准配置：

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
}

4.8 跑一次前向 + 参数量统计

import torch

model = GPTModel(GPT_CONFIG_124M)
model.eval()

x = torch.randint(0, 50257, (2, 16))
with torch.no_grad():
    logits = model(x)
print(logits.shape)   # (2, 16, 50257)

n_params = sum(p.numel() for p in model.parameters())
print(f"参数量: {n_params/1e6:.2f}M")
# 实际约 163M（包含位置嵌入），即所谓 "GPT-2 124M" 的口径仅算非嵌入部分

“124M” 的命名沿用了 OpenAI 的口径——只统计 Transformer 主体参数，不含 token/位置嵌入和输出头。

4.9 权重共享（可选优化）

GPT-2 把 token embedding 矩阵 $E \in \mathbb{R}^{V \times C}$ 和输出投影 $W_{out} \in \mathbb{R}^{C \times V}$ 绑定为同一组参数（转置共享）。这能：

  • 减少约 $V \cdot C \approx 38\text{M}$ 参数；
  • 让”词义空间”和”输出空间”对齐，通常略微提升困惑度。

实现：

self.out_head.weight = self.tok_emb.weight  # 形状必须一致：(V, C)

检查清单

  • 我能在白板上画出 Pre-LN Transformer Block 的全部数据流。
  • 我能解释 LayerNorm 与 BatchNorm 在序列建模中的取舍。
  • 我能跑通一次前向，并算出参数量。

练习题

  1. 把 Pre-LN 改成 Post-LN（即先做子层再 LN），用同一份训练数据训一两个 step，比较 loss 曲线的稳定性。
  2. 关掉残差连接（x = self.drop1(self.attn(...)) 而不再加 x +），观察前几个 step 内 loss 是否爆炸。
  3. 试着把 FFN 的 expansion 从 4 改成 2 或 8，记录参数量与 loss 的变化。

📖 第4章补充材料 → — Pre-LN vs Post-LN、GELU/SwiGLU详解、GPT-2四尺寸配置

← 上一章 · 返回目录 · 下一章 · 预训练 →