【解剖ChatGPT-3】GPT的核心：如何用代码实现Transformer Block

heidsoft

发布于 2026-07-02 10:55:30

130

文章被收录于专栏：云计算与大数据云计算与大数据

作者按：这是《从零理解ChatGPT》系列的第三篇。前两期我们讲了GPT演进史和Attention机制，今天我们进入实践环节——用PyTorch从零实现一个GPT模型。

前言

前面两期我们讲完了理论：GPT是什么、Attention是怎么工作的。

今天我们动手写代码。

我们参考Sebastian Raschka的《Build a Large Language Model from Scratch》，用PyTorch实现一个GPT-2（124M参数版本）。

GPT模型整体架构

先上一个全局图：

┌──────────────────────────────────────────────────────┐
│  Input: Token IDs [batch, seq_len]                  │
├──────────────────────────────────────────────────────┤
│  Token Embedding: [vocab_size] → [emb_dim]          │
│  + Positional Embedding: [context_length] → [emb_dim]│
│  = Input Embeddings [batch, seq_len, emb_dim]       │
├──────────────────────────────────────────────────────┤
│  Dropout                                              │
├──────────────────────────────────────────────────────┤
│  ×12 Transformer Blocks                              │
│    (每个包含: LayerNorm + MHA + FeedForward + 残差) │
├──────────────────────────────────────────────────────┤
│  Final LayerNorm                                     │
├──────────────────────────────────────────────────────┤
│  Output Head: [emb_dim] → [vocab_size]              │
│  = Logits [batch, seq_len, vocab_size]             │
└──────────────────────────────────────────────────────┘

GPT-2配置参数

我们用最小的GPT-2 Small（124M参数）：

GPT_CONFIG_124M = {
    "vocab_size": 50257,   # BPE词表大小（GPT-2使用）
    "context_length": 1024,  # 最大输入长度
    "emb_dim": 768,         # Embedding维度
    "n_heads": 12,          # 注意力头数
    "n_layers": 12,         # Transformer块数
    "drop_rate": 0.1,       # Dropout率
    "qkv_bias": False       # QKV偏置（现代LLM通常False）
}

1. Layer Normalization

作用：稳定训练，让每层的输出均值=0，方差=1

class LayerNorm(nn.Module):
    def __init__(self, emb_dim, eps=1e-5):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(emb_dim))  # 可学习的缩放
        self.shift = nn.Parameter(torch.zeros(emb_dim))  # 可学习的平移
        self.eps = eps
    
    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x_norm = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * x_norm + self.shift

为什么LayerNorm而不是BatchNorm？

• LLM的batch size经常变化（GPU显存限制）
• LayerNorm对每个样本独立归一化，更灵活
• 分布式训练更稳定

2. GELU激活函数

ReLU vs GELU：

传统NLP用ReLU（负值为0），GPT用GELU：

GELU(x) = 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x³)))

为什么GELU更好？

• ReLU是分段线性，负值直接变0
• GELU是平滑的，允许小负值输出
• 梯度更平滑，优化更稳定

class GELU(nn.Module):
    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))
        ))

3. Feed Forward Network

作用：在每个位置独立做非线性变换，增加模型容量

结构：扩展 → 激活 → 收缩

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

为什么扩展4倍？

这是GPT-2的经验值。扩展后的空间让模型能学习更复杂的表示，然后再压缩回来。

4. Transformer Block

把上面的组件组装起来：

class TransformerBlock(nn.Module):
    def__init__(self, cfg):
        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"])
    
    defforward(self, x):
        # 第一个残差块：Self-Attention
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)       # Self-Attention
        x = self.drop_shortcut(x)
        x = x + shortcut      # 残差连接
        
        # 第二个残差块：Feed Forward
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)       # 前馈网络
        x = self.drop_shortcut(x)
        x = x + shortcut      # 残差连接
        
        return x

5. 残差连接（Skip Connection）

为什么需要？

深层网络中，梯度回传时越来越小（梯度消失），导致前面的层学不到东西。

残差连接的作用：

x_out = x_in + Block(x_in)

梯度反向传播：
∂L/∂x_in = ∂L/∂x_out * (1 + ∂x_out/∂Block输入)
         = ∂L/∂x_out + ∂L/∂x_out * ∂Block输出/∂...

"1"这个通路让梯度直接回传，彻底解决梯度消失。

6. 完整的GPT Model

class GPTModel(nn.Module):
    def__init__(self, cfg):
        super().__init__()
        # Embedding层
        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"])
        
        # Transformer块堆叠
        self.trf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ inrange(cfg["n_layers"])])
        
        # 输出部分
        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
    
    defforward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        
        # Token Embedding + Positional Embedding
        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)
        
        # 通过Transformer块
        x = self.trf_blocks(x)
        
        # 最终归一化 + 输出投影
        x = self.final_norm(x)
        logits = self.out_head(x)
        
        return logits

7. 计算参数量

model = GPTModel(GPT_CONFIG_124M)

total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total_params:,}")
# Total parameters: 163,009,536

# 如果算上weight tying（复用embedding）
total_params_gpt2 = total_params - sum(p.numel() for p in model.out_head.parameters())
print(f"GPT-2 124M parameters: {total_params_gpt2:,}")
# GPT-2 124M parameters: 124,412,160

参数量分解：

组件	参数量	占比
Token Embedding	50257 × 768 ≈ 38.6M	24%
Positional Embedding	1024 × 768 ≈ 0.8M	0.5%
12层Transformer	每层约12M × 12 ≈ 144M	88%
Output Head	50257 × 768 ≈ 38.6M (复用)	-

内存占用：

163M parameters × 4 bytes (float32) ≈ 621 MB

8. 测试：生成文字

未训练的模型输出是乱码：

import tiktoken

tokenizer = tiktoken.get_encoding("gpt2")
model = GPTModel(GPT_CONFIG_124M)

# 输入
start_context = "Hello, I am"
input_ids = tokenizer.encode(start_context)
input_tensor = torch.tensor(input_ids).unsqueeze(0)

# 前向传播
with torch.no_grad():
    logits = model(input_tensor)

print("Output shape:", logits.shape)
# Output shape: torch.Size([1, 4, 50257])
# 4个词，每个词是50257维的logits（vocab_size）

为什么是50257维？

这是GPT-2的词表大小。模型对每个位置输出一个50257维向量，表示下一个词的概率分布。

总结

今天我们用PyTorch实现了一个完整的GPT模型：

组件	代码行数	作用
LayerNorm	10行	归一化，稳定训练
GELU	5行	平滑激活函数
FeedForward	10行	非线性变换，4倍扩展
TransformerBlock	20行	核心块：Attn + FF + 残差
GPTModel	25行	完整模型组装