预训练迷你 GPT

预训练迷你 GPT

本节摘要:GPT-2 Small 有 1.24 亿参数——12 层 Transformer、12 个注意力头、768 维嵌入。你可以在单卡上几小时从零训出来,但多数人从不这么做,只用预训练权重。可你要没亲手训过一个,就根本不理解你拿来做产品的那个黑盒里发生了什么。本节用 numpy(不是 PyTorch)从零搭建 GPT-2 Small:token 嵌入、位置嵌入、因果掩码多头注意力、残差连接、前馈网络、权重绑定的输出头,每一个矩阵乘法都看得见。你会看到 1.24 亿个数字如何「合谋」预测下一个词,理解 KV 缓存为何让推理变快、prefill 与 decode 为何有截然不同的瓶颈。

学习目标

阅读完本节,你应当能够:

1. 从零实现完整 GPT-2 架构(124M 参数):token 嵌入、位置嵌入、Transformer 块、语言模型头。
2. 在文本语料上用下一 token 预测与交叉熵损失训练 GPT 模型。
3. 实现带温度采样与 top-k/top-p 过滤的自回归文本生成。
4. 监控训练损失曲线,验证模型学到了连贯的语言模式。

一、问题与直觉

你知道 Transformer 是什么。你看过那些图。你能背「Attention is all you need」,能在白板上画「Multi-Head Attention」的框。

这些都不意味着你理解模型生成文本时发生了什么。

GPT-2 Small 有 124,438,272 个参数(含权重绑定)。每一个都是跑训练循环设出来的:前向、算损失、反向、更新权重。12 个 Transformer 块,每块 12 个注意力头,768 维嵌入空间,50,257 个 token 的词表。模型每生成一个 token,全部 1.24 亿参数都参与一条矩阵乘法链,把 token ID 序列变成下一个 token 上的概率分布。你要是从没自己搭过,你就是在跟黑盒打交道——能用 API、能微调,但当模型幻觉、重复、拒不遵循指令时,你对为什么毫无心智模型。

本节就用 numpy 从零搭 GPT-2 Small。每个矩阵乘法可见,每个梯度由你的代码算。你会清楚看到 1.24 亿个数字如何合谋预测下一个词。

GPT 架构

GPT 是自回归语言模型。「自回归」指它一次生成一个 token,每个都以所有先前 token 为条件。架构是一摞 Transformer 解码块。从 token ID 到下一 token 概率的完整计算图:

1. token ID 输入,形状 (batch, seq_len)。
2. token 嵌入查表,每 ID 映成 768 维向量。
3. 位置嵌入查表,每个位置 (0,1,2,...) 映成 768 维向量。
4. 两者相加。
5. 过 12 个 Transformer 块。
6. 最终层归一化。
7. 线性投影到词表大小 (batch, seq_len, vocab_size)。
8. softmax 得概率。

整个模型就这些:没有卷积、没有循环,只有嵌入、注意力、前馈、层归一,堆 12 次。

Transformer 块

12 块都遵循同一模式(GPT-2 用 pre-norm,不是原始 Transformer 的 post-norm):(1) LayerNorm;(2) 多头自注意力;(3) 残差(把输入加回);(4) LayerNorm;(5) 前馈网络(MLP);(6) 残差。残差连接至关重要——没它,反向传播时梯度传到第 1 块就消失了;有它,梯度能沿「跳连」路径从损失直达任意层。这就是为什么能堆 12、32 甚至 96 块(GPT-4 据传用 120 块)。

注意力:核心机制

自注意力让每个 token 看所有先前 token,决定各看多少。每个位置从输入算三个向量:查询 Q(「我在找什么」)、键 K(「我含什么」)、值 V(「我带什么信息」)。

Q = input @ W_q;  K = input @ W_k;  V = input @ W_v
scores = Q @ K^T / sqrt(d_k)
scores = mask(scores)        # 因果掩码:未来位置填 -inf
weights = softmax(scores)
output = weights @ V
​

因果掩码让 GPT 自回归——位置 5 能看 0~5,不能看 6、7、8。多头注意力把 768 维空间切成 12 个 64 维的头,各学不同注意力模式:一个可能追踪主谓一致,一个追踪同义语义,一个追踪位置邻近。12 个头的输出拼接后投影回 768 维。除以 sqrt(d_k)(=8)是缩放——不做它,高维向量点积变大,把 softmax 推到梯度近零的区域。这是「Attention Is All You Need」的关键洞见之一。

KV 缓存:推理为何快

训练时一次处理整条序列。推理时一次生成一个 token。不优化的话,生成第 N 个 token 要为所有 N-1 个先前 token 重算注意力——每 token O(N²),整条 O(N³)。KV 缓存解决:算完每个 token 的 K、V 后存起来,生成第 N+1 时只算新 token 的 Q,从缓存读先前所有 K、V。对 GPT-2(12 层 12 头),每 token 的 KV 缓存存 2×12×12×64 = 18,432 个值;1024-token 序列约 75MB(FP32)。对 Llama 3 405B(128 层),单条序列的 KV 缓存可超 10GB——这就是长上下文推理受内存限制的原因。

Prefill 与 Decode:推理的两阶段

发提示给 LLM,推理分两阶段。Prefill 并行处理整条提示(所有 token 已知,可同时算所有位置注意力)——这阶段受算力限制,A100 上满吞吐做矩阵乘,1000-token 提示约 20~50ms。Decode 一次生成一个 token,每个新 token 依赖所有先前 token——这阶段受内存带宽限制,瓶颈是从显存读模型权重与 KV 缓存,矩阵运算本身微秒级就完成,计算核心干等内存。这个区别对生产系统至关重要:prefill 吞吐随 GPU 算力缩放,decode 吞吐随内存带宽缩放。这就是 NVIDIA H100 相比 A100 重点提升内存带宽的原因——直接加速 token 生成。

训练循环

训练 LLM 就是下一 token 预测:给定 [0,1,...,N-1] 预测 [1,2,...,N]。损失是模型预测分布与真实下一 token 的交叉熵。一步训练:(1) 前向,得每位置 logits;(2) 算交叉熵损失(目标是输入右移一位);(3) 反向,为全部 124M 参数算梯度;(4) 优化器步,更新权重(GPT-2 用 Adam,带学习率预热与余弦衰减)。学习率调度比想象中重要:GPT-2 在前 2000 步从 0 预热到峰值,再按余弦曲线衰减——一起始用高学习率模型会发散,后期保持高学习率会震荡。「预热后衰减」是每个主流 LLM 的标配。

GPT-2 Small 的数字账

输出投影(logits 头)与 token 嵌入矩阵共享权重——这叫权重绑定(Weight Tying),省 38M 参数还提升性能,因为它强制模型对输入与输出用同一表示空间。

二、从零实现

Step 1:嵌入层

import numpy as np
class Embedding:
    def __init__(self, vocab_size, embed_dim, max_seq_len):
        self.token_embed = np.random.randn(vocab_size, embed_dim) * 0.02
        self.pos_embed = np.random.randn(max_seq_len, embed_dim) * 0.02
    def forward(self, token_ids):
        seq_len = token_ids.shape[-1]
        return self.token_embed[token_ids] + self.pos_embed[:seq_len]
​

0.02 的初始化标准差来自 GPT-2 论文——太大则前向输出极端值致训练不稳,太小则各输入初始输出近乎相同、早期梯度信号无用。

Step 2:带因果掩码的自注意力

def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(..., -2, -1) / np.sqrt(d_k)
    if mask is not None: scores = scores + mask
    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))   # 减最大值防溢出
    weights = weights / weights.sum(axis=-1, keepdims=True)
    return weights @ V
​

设计要点:softmax 前减去最大值是数值稳定技巧——不做它 exp(大数) 会溢出。这不改变输出,因为 softmax(x-c)=softmax(x) 对任意常数 c 成立。

Step 3:多头注意力

把 768 维输入切成 12 个 64 维头,各独立算注意力,拼接后投影回 768 维。这段 reshape-transpose-reshape 是多头注意力最绕的部分:(batch, seq, 768) → (batch, seq, 12, 64) → (batch, 12, seq, 64),于是 12 个头各有自己的 (seq, 64) 矩阵;注意力算完再倒回去。

class MultiHeadAttention:
    def __init__(self, embed_dim, num_heads):
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.W_q = np.random.randn(embed_dim, embed_dim) * 0.02
        # W_k, W_v, W_out 同理
    def forward(self, x, mask=None):
        B, S, D = x.shape
        Q = (x @ self.W_q).reshape(B, S, self.num_heads, self.head_dim).transpose(0,2,1,3)
        # K, V 同理
        scores = Q @ K.transpose(0,1,3,2) / np.sqrt(self.head_dim)
        if mask is not None: scores = scores + mask
        weights = softmax(scores)
        out = (weights @ V).transpose(0,2,1,3).reshape(B, S, D)
        return out @ self.W_out
​

Step 4:Transformer 块

完整一块:LayerNorm、带残差的多头注意力、LayerNorm、带残差的前馈。前馈把 768 维扩到 3072 维(4×),加非线性,再投回 768——这种扩张-收缩让模型在每个位置有更「宽」的内部表示。

class TransformerBlock:
    def __init__(self, embed_dim, num_heads, ff_dim):
        self.ln1 = LayerNorm(embed_dim); self.attn = MultiHeadAttention(embed_dim, num_heads)
        self.ln2 = LayerNorm(embed_dim); self.ffn = FeedForward(embed_dim, ff_dim)
    def forward(self, x, mask=None):
        x = x + self.attn.forward(self.ln1.forward(x), mask)   # 残差
        x = x + self.ffn.forward(self.ln2.forward(x))          # 残差
        return x
​

Step 5:完整 GPT 模型

堆 12 块,前置嵌入,后置输出投影。注意权重绑定:logits = x @ self.embedding.token_embed.T——输出投影复用 token 嵌入矩阵(转置)。

class MiniGPT:
    def __init__(self, vocab_size=50257, embed_dim=768, num_heads=12,
                 num_layers=12, max_seq_len=1024, ff_dim=3072):
        self.embedding = Embedding(vocab_size, embed_dim, max_seq_len)
        self.blocks = [TransformerBlock(embed_dim, num_heads, ff_dim) for _ in range(num_layers)]
        self.ln_f = LayerNorm(embed_dim)
    def forward(self, token_ids):
        S = token_ids.shape[-1]
        mask = np.triu(np.full((S, S), -1e9), k=1)             # 因果掩码
        x = self.embedding.forward(token_ids)
        for block in self.blocks: x = block.forward(x, mask)
        x = self.ln_f.forward(x)
        return x @ self.embedding.token_embed.T                 # 权重绑定的输出头
​

Step 6:训练循环

真正在 124M 上训练要 GPU 和 PyTorch。这里用纯 numpy 在小模型(4 层、4 头、128 维)上演示机制。损失从 ln(vocab_size) 起步——256-token 字节级词表即 ln(256)=5.55;随机模型给每 token 等概率。训练推进,损失下降,因为模型学会了预测常见模式:「t」后接「h」、句号后接空格等。

Step 7:文本生成

def generate(model, prompt_tokens, max_new_tokens=100, temperature=0.8):
    tokens = list(prompt_tokens)
    seq_len = model.embedding.pos_embed.shape[0]
    for _ in range(max_new_tokens):
        ctx = np.array(tokens[-seq_len:]).reshape(1, -1)       # 不超上下文窗口
        logits = model.forward(ctx)[0, -1, :] / temperature
        probs = np.exp(logits - logits.max()); probs /= probs.sum()
        tokens.append(int(np.random.choice(len(probs), p=probs)))
    return tokens
​

温度控制随机性:1.0 用原始分布,0.5 锐化(更确定),1.5 抹平(更随机),0.0 是贪心(总取最高概率 token)。tokens[-seq_len:] 窗口是必要的——模型有最大上下文长度(GPT-2 是 1024),超了就得丢最老的 token,这就是人人在说的「上下文窗口」。

三、框架对比

在小语料小模型上,生成文本最多半连贯——它能从训练文本学些字节级模式,但无法像有 40GB 数据和完整 124M 架构的 GPT-2 那样泛化。重点不是输出质量,而是你能追踪每一步:嵌入查表、注意力计算、前馈变换、logit 投影、softmax、采样,每个操作都可见。生产中你会用 PyTorch(自动微分、GPU、FSDP)而非手写 numpy,但前向-损失-反向-更新的循环完全一致,只是优化器更精巧。

四、可复用产物

本节产出 outputs/prompt-gpt-architecture-analyzer.md——一个分析任意 GPT 风格模型架构选择的提示。喂给它模型卡或技术报告,它拆解参数分配、注意力设计、缩放决策。

五、练习

1. (Easy) 把模型改成 24 层 16 头(原 12/12),数参数。深度翻倍与宽度(嵌入维)翻倍,影响有何不同?

2. (Medium) 实现 GELU 激活(GELU(x)=x*0.5*(1+erf(x/sqrt(2))))替换前馈里的 ReLU,各训 500 步对比最终损失。

3. (Medium) 给生成函数加 KV 缓存:首次前向后存各层 K、V,后续 token 复用。测生成 200 token 有无缓存的耗时差。

4. (Hard) 实现 top-k(只考虑最高概率的 k 个)与 top-p(核采样,累计概率超 p 的最小 token 集)采样,在温度 0.8 下对比 top-k=50 与 top-p=0.95 的输出质量。

5. (Hard) 训 1000 步画损失曲线,识别三阶段:快速初降(学常见字节)、缓慢中段(学字节模式)、平台(小语料上过拟合)。无论 128 维还是 GPT-4,曲线形状一样。

本节要点回顾

1. 自回归 = 一次一 token:每个输出 token 以所有先前 token 为条件,预测 P(token_n | token_0..n-1)。
2. 因果掩码让训练不偷看未来:上三角填 -inf,softmax 后未来位置权重为 0。
3. 多头注意力切空间:768 维切成 12 个 64 维头,各学不同关系(语法、语义、位置邻近),拼接投影回 768。
4. 残差连接救梯度:没它,反向传到第 1 块梯度就消失;有它,梯度能沿跳连直达任意层,所以能堆 96+ 块。
5. 权重绑定省 38M 参数:输出投影复用 token 嵌入矩阵(转置),强制输入输出同表示空间。
6. KV 缓存把 O(N³) 降下来:存住先前 token 的 K、V,生成新 token 只算新 Q,每 token 成本从 O(N) 降到 O(1) 的 K/V 计算。
7. Prefill 算力受限,Decode 内存受限:prefill 并行处理全提示满吞吐,decode 一次一 token 干等内存读权重。
8. 缩放除 sqrt(d_k):防高维点积把 softmax 推到梯度近零区,是原始 Transformer 论文的关键洞见。
9. 损失从 ln(vocab) 起步:随机模型给每 token 等概率,训练下降代表学到了「t 后接 h」这类模式。
10. 温度控制随机性:0 抽取/代码,0.8 平衡,1.5 创意;上下文窗口超了就丢最老 token。

下一节,我们解决一个残酷现实:70B 模型在单卡上放不下。分布式训练、FSDP、DeepSpeed 是唯一出路。