核心

GPT 模型架构从零实现 — 课程A W4

对应课程: course_A W4 (GPT架构实现)

难度: ★★★ (L3 进阶)

知识基础: Build LLM from Scratch Ch4 — 从头实现GPT模型

> 理解模型内部结构 — 每一层都是潜在的审计关注点

0. 复用上一课的多头注意力

先把 Ch3 实现的多头注意力拿过来用。

代码

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout=0.0, n_heads=4):
        super().__init__()
        assert d_out % n_heads == 0
        self.d_out = d_out
        self.n_heads = n_heads
        self.head_dim = d_out // n_heads
        
        self.W_q = nn.Linear(d_in, d_out, bias=False)
        self.W_k = nn.Linear(d_in, d_out, bias=False)
        self.W_v = nn.Linear(d_in, d_out, bias=False)
        self.out_proj = nn.Linear(d_out, d_out)
        self.dropout = nn.Dropout(dropout)
        
        mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
        self.register_buffer('mask', mask.bool())
    
    def forward(self, x):
        b, num_tokens, d_in = x.shape
        queries = self.W_q(x).view(b, num_tokens, self.n_heads, self.head_dim).transpose(1, 2)
        keys = self.W_k(x).view(b, num_tokens, self.n_heads, self.head_dim).transpose(1, 2)
        values = self.W_v(x).view(b, num_tokens, self.n_heads, self.head_dim).transpose(1, 2)
        
        attn = queries @ keys.transpose(-2, -1)
        mask = self.mask.bool()[:num_tokens, :num_tokens].unsqueeze(0).unsqueeze(0)
        attn.masked_fill_(mask, -torch.inf)
        
        out = torch.softmax(attn / self.head_dim**0.5, dim=-1)
        out = self.dropout(out) @ values
        out = out.transpose(1, 2).contiguous().view(b, num_tokens, self.d_out)
        return self.out_proj(out)

print("✓ MultiHeadAttention loaded")

1. LayerNorm (层归一化)

对应书籍 4.2节: 对每个 token 的所有特征做归一化，稳定训练。

公式: $\text{LayerNorm}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$

代码

class LayerNorm(nn.Module):
    """层归一化 (RMS Norm 的简化版之前的标准)"""
    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)
        normed = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * normed + self.shift

# GPT-2 使用 LayerNorm 放在 attention/FFN 之前 (Pre-LayerNorm)
# 现代 LLM 使用 RMSNorm (去掉均值计算，更快)
class RMSNorm(nn.Module):
    """RMS LayerNorm — 现代LLM (Llama, Qwen) 使用的版本"""
    def __init__(self, emb_dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(emb_dim))
    
    def forward(self, x):
        rms = torch.sqrt((x ** 2).mean(dim=-1, keepdim=True) + self.eps)
        return self.weight * (x / rms)

# 两种 Norm 的对比
sample = torch.randn(2, 4, 8)
ln = LayerNorm(8)
rms = RMSNorm(8)
print(f"输入均值: {sample.mean().item():.4f}, 方差: {sample.var().item():.4f}")
print(f"LayerNorm 输出均值: {ln(sample).mean().item():.4f}, 方差: {ln(sample).var().item():.4f}")
print(f"RMSNorm 输出均值: {rms(sample).mean().item():.4f}, 方差: {rms(sample).var().item():.4f}")

2. GELU 激活函数

对应书籍 4.3节: GPT 使用 GELU 而不是 ReLU。

GELU ≈ ReLU 的平滑版本，梯度更好。

代码

class GELU(nn.Module):
    """Gaussian Error Linear Unit (GPT 使用的激活函数)"""
    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 * x**3)
        ))

class FeedForward(nn.Module):
    """前馈神经网络: Linear → GELU → Linear"""
    def __init__(self, emb_dim):
        super().__init__()
        # GPT 在中间层将维度扩展 4 倍
        self.layers = nn.Sequential(
            nn.Linear(emb_dim, 4 * emb_dim),
            GELU(),
            nn.Linear(4 * emb_dim, emb_dim),
        )
    
    def forward(self, x):
        return self.layers(x)

# 测试
ff = FeedForward(emb_dim=8)
x = torch.randn(2, 4, 8)
print(f"FeedForward: {x.shape} → {ff(x).shape}")
print(f"中间维度: {8} → {4*8} → {8}")

3. TransformerBlock — 完整模块

对应书籍 4.5节: 把注意力 + 前馈网络 + 残差连接 + LayerNorm 组装起来。

Pre-LayerNorm 架构 (GPT-2): 先 Norm 再 Attn/FFN

代码

class TransformerBlock(nn.Module):
    """一个完整的 Transformer 层"""
    def __init__(self, cfg):
        super().__init__()
        d_model = cfg["emb_dim"]
        
        self.attn = MultiHeadAttention(
            d_in=d_model,
            d_out=d_model,
            context_length=cfg["context_length"],
            dropout=cfg["dropout"],
            n_heads=cfg["n_heads"]
        )
        self.ff = FeedForward(d_model)
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.dropout = nn.Dropout(cfg["dropout"])
    
    def forward(self, x):
        # Pre-LayerNorm: 先归一化再计算
        shortcut = x
        x = self.norm1(x)
        x = self.attn(x)
        x = self.dropout(x)
        x = x + shortcut  # 残差连接
        
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.dropout(x)
        x = x + shortcut  # 残差连接
        return x

# 测试一个 TransformerBlock
cfg = {"emb_dim": 8, "context_length": 16, "dropout": 0.1, "n_heads": 2}
block = TransformerBlock(cfg)
x = torch.randn(2, 4, 8)
print(f"TransformerBlock: {x.shape} → {block(x).shape}")

残差连接的作用:

让梯度可以直接流回浅层 (解决深层网络梯度消失)
如果注意力层什么都没学到，信息通过 shortcut 直接绕过

4. GPTModel — 完整模型

对应书籍 4.6节: 总装全部组件。

代码

class GPTModel(nn.Module):
    """完整的 GPT 模型"""
    def __init__(self, cfg):
        super().__init__()
        self.config = cfg
        
        # Token + 位置嵌入
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.dropout = nn.Dropout(cfg["dropout"])
        
        # N 层 TransformerBlock
        self.trf_blocks = nn.Sequential(*[
            TransformerBlock(cfg) for _ in range(cfg["n_layers"])
        ])
        
        # 最终 LayerNorm + 输出投影到词表
        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
    
    def forward(self, x):
        # x: (batch, num_tokens) — token IDs
        batch, seq_len = x.shape
        
        # 嵌入 + 位置编码
        token_embeds = self.tok_emb(x)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=x.device))
        x = self.dropout(token_embeds + pos_embeds)
        
        # Transformer 层
        x = self.trf_blocks(x)
        
        # 输出
        x = self.final_norm(x)
        logits = self.out_head(x)  # (batch, seq, vocab_size)
        return logits

# GPT-2 Small (124M) 配置
GPT2_SMALL = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "dropout": 0.1,
}

# 用微型配置测试
mini_cfg = {
    "vocab_size": 1000,
    "context_length": 64,
    "emb_dim": 64,
    "n_heads": 4,
    "n_layers": 4,
    "dropout": 0.1,
}

model = GPTModel(mini_cfg)
x = torch.randint(0, 1000, (2, 16))  # (batch=2, seq=16)
logits = model(x)
total_params = sum(p.numel() for p in model.parameters())
print(f"输入: {x.shape}")
print(f"输出 (logits): {logits.shape}")
print(f"参数量: {total_params:,}")

5. 文本生成

对应书籍 4.7节: 用训练好的 GPT 模型逐个 token 生成文本。

代码

def generate_text_simple(model, idx, max_new_tokens, context_size):
    """逐 token 生成文本 (最简版)"""
    for _ in range(max_new_tokens):
        # 只取最后 context_size 个 token (如果序列太长)
        idx_cond = idx[:, -context_size:]
        
        with torch.no_grad():
            logits = model(idx_cond)
        
        # 只取最后一个位置的 logits
        logits = logits[:, -1, :]  # (batch, vocab_size)
        
        # 取概率最高的 token (贪心解码)
        next_token = torch.argmax(logits, dim=-1, keepdim=True)
        idx = torch.cat([idx, next_token], dim=1)
    
    return idx

# 演示: 用随机初始化的模型生成 (因为没有训练，所以是乱码)
start = torch.tensor([[1, 2, 3, 4, 5]])  # 模拟 token IDs
output = generate_text_simple(model, start, max_new_tokens=10, context_size=64)
print(f"起始: {start[0].tolist()}")
print(f"生成: {output[0, 5:].tolist()}")
print("(因为是随机初始化，输出无意义。训练后才能生成有意义的文本)")

6. 模型配置对比

| 模型 | 层数 | 注意力头 | 嵌入维度 | 参数量 |

|------|------|---------|---------|-------|

| GPT-2 Small | 12 | 12 | 768 | 124M |

| GPT-2 Medium | 24 | 16 | 1024 | 355M |

| GPT-2 Large | 36 | 20 | 1280 | 774M |

| GPT-2 XL | 48 | 25 | 1600 | 1.5B |

| GPT-3 | 96 | 96 | 12288 | 175B |

| Qwen3.5-9B | ~40 | 32 | 4096 | 9B |

代码

# %%ai openai-chat-custom:Qwen3.5-9B-Q4_K_M.gguf
# 解释 GPT 架构中的 "Pre-LayerNorm" 和 "Post-LayerNorm" 有什么区别？为什么现代 LLM 都改用 Pre-LayerNorm？

检查清单

[x] LayerNorm / RMSNorm 实现
[x] GELU + FeedForward
[x] TransformerBlock (残差连接 + Pre-LayerNorm)
[x] GPTModel 总装
[x] 文本生成函数
[ ] 下一步: 08_Pretraining.ipynb — 训练这个模型

🔍 审计角度 (Audit Lens)

1. 每一层都是审计点

GPT 模型的每一层都可能引入或放大问题：

Embedding层 → token表征中是否包含敏感属性（如性别/种族） Attention层 → 关注模式是否公平 FeedForward层 → 激活值是否异常（可能指示数据中毒） Output层 → 输出分布是否偏移（指示模型漂移）

2. 参数规模与合规成本

模型参数越多, 算法备案的审查范围越大:

中国算法备案对"具有舆论属性或社会动员能力"的算法有额外要求
大模型(Qwen3.5-9B)的审计范围包括: 训练数据、模型结构、输出控制、安全评估
可解释性报告需要覆盖到每一层的作用

3. 残差连接的安全含义

残差连接让梯度直接流过 — 这对训练好, 但对审计意味着:

隐患: 如果某一层被攻陷（如通过模型窃取攻击）, 残差连接会让恶意行为绕过后续层的安全检查, "直接"影响输出

4. 输出层温度与风险

GPTModel最后一层是Linear+Softmax
温度=0 → 确定性输出(适合审计/合规场景)
温度>0 → 创造性输出(需要人工审核)

金融合规场景要求温度≤0.3, 因为高温会增加"幻觉"概率。