GPT 模型组装:从块到 1.24 亿参数

GPT 模型组装:从块到 1.24 亿参数

本节摘要:12 个块叠起来,加 token 嵌入、学习位置嵌入、最终 LayerNorm、绑定的 LM 头——这就是整个 1.24 亿参数的 GPT 模型。本节把这些件组装成可工作的类,数参数确认匹配参考 124M 形状,并用多项采样、温度、top-k 生成文本。transformer 块单独无用,你得把 token id 变向量、混位置信息、过堆叠、投回词表 logits——忘四步任一,模型要么前向失败、要么位置信息漂移、要么不会说话。形状也要紧:参考 GPT-2 small 在上述配置下正好 1.24 亿参数,数对不上说明接线错了。

对应原课程:Phase 19 · Lesson 35 · gpt-model-assembly(原英文 phases/19-capstone-projects/35-gpt-model-assembly/docs/en.md)。本节属「从零构建 GPT」赛道第六节。

学习目标

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

1. 把第 34 节的 transformer 块组装成完整 GPT:token 嵌入、位置嵌入、N 块、最终 LayerNorm、LM 头。
2. 复现 1.24 亿参数配置:词表 50257、上下文 1024、嵌入 768、12 头、12 层。
3. 把 LM 头权重绑到 token 嵌入,解释为何在此规模省约 3800 万参数。
4. 用多项采样、温度缩放、top-k 截断从提示生成文本,以滑动窗守上下文长。
5. 对照 124M 目标量参数量与前向成本。

一、问题与直觉

参考 GPT-2 small 在上述配置下正好 1.24 亿参数,数字不魔幻:词表 50257 × 嵌入 768 是 token 表;位置 1024 × 768 是位置表;12 块各约 700 万参数 = 8400 万;最终头经 weight tying 复用 token 表。求和落地 1.24 亿。建一个参数量对不上参考的模型,说明你接错了。

token id 变 token 向量,位置 id 变位置向量,两者相加过堆叠,最终 LayerNorm 是块外唯一在每个现代变体存活的部分,LM 头复用 token 嵌入矩阵——这就是 weight tying。

二、从零实现

Weight tying:token 嵌入形 (vocab, d_model),LM 头需从 d_model 投回 vocab——两者互为转置。tying 指字面上同参数张量用两次。词表 50257、d_model 768 时,矩阵 3800 万参数;不绑付两次,绑付一次且得稍干净的梯度信号(嵌入与头一起更新)。

位置嵌入是学习的非正弦:GPT-2 发学习位置嵌入,位置表是 (1024, 768) 的单参数张量,前向查位置 0~T-1 加到 token 嵌入。这是最简位置方案(RoPE/ALiBi/T5 相对偏置是替代),也是 124M 参考用的。

生成:温度、top-k、多项:生成是自回归。每步模型返全词表在每位置的 logits,取末位置,除温度,可选地把除 top k 外掩码为负无穷,softmax 得概率,从分布采一 token。三旋钮三行为:温度近零塌成贪心;温度一匹配自然分布;top-k 一是贪心;top-k 四十滤长尾。组合要紧,下一节训练用生成作定性评估信号。

def generate(model, ids, max_new, temperature=1.0, top_k=None):
    for _ in range(max_new):
        ctx = ids[:, -model.context_length:]            # 滑窗守上下文
        logits = model(ctx)[:, -1, :] / temperature
        if top_k: logits = mask_topk(logits, top_k)
        probs = F.softmax(logits, dim=-1)
        ids = torch.cat([ids, torch.multinomial(probs, 1)], dim=1)
    return ids
​

三、可复用产物

code/main.py 实现:GPTConfig 数据类(124M 默认:vocab_size=50257/context_length=1024/d_model=768/num_heads=12/num_layers=12/mlp_expansion=4/dropout=0.1/use_bias=True/weight_tying=True)、GPTModel(token 嵌入、位置嵌入、嵌入 dropout、12 个 TransformerBlock、最终 LayerNorm、标志开时绑到 token 嵌入的 lm_head)、count_parameters(返唯一参数数, honoring weight tying)、generate(温度/top-k/多项/滑窗)、demo(建模型、打参数量对照 124M、从固定提示生成短序列)。为求快,demo 还端到端跑微型配置(d_model=64、num_layers=2)内联打生成 token 序;124M 配置只建、只跑参数量与一次前向。

四、生产模式

三模式区分能跑与能发。残差投影小初始化:注意力输出投影与 MLP 第二线性都直喂残差加,用与每个其他线性同 std 初始化,会给随深度增长的残差流、把最终 LayerNorm 推进热区——这两投影的 std 缩 1/sqrt(2*num_layers),残差流 12 层保 sane 范围。缓存位置 id 张量:torch.arange(T) 每前向分新内存,在 __init__ 为最大上下文分一次,每调用切前 T 项,省分配往返。参数级绑权重:lm_head.weight = token_embedding.weight 共享张量,拷贝不算;优化器需更一参数、autograd 图需一累积,拷贝会让头漂离嵌入,weight tying 白搭。

五、框架对比

本节的「学习位置嵌入 + GELU + LayerNorm + weight tying」是 GPT-2 small 的 1.24 亿参考。换 RoPE 得 LLaMA 族(不动块或头),换 SiLU 与 RMSNorm 得 LLaMA 族其余改动——同骨架。生成函数与任意 logits 源工作(不只本模型),第 37 节从预训练 GPT-2 文件拉 logits 可复用同生成循环。

六、练习

1. 解绑头:把 LM 头从 token 嵌入解绑,重数参数,确认增量 50257×768=3800 万。
2. 换正弦位置:把学习位置嵌入换构造时算的正弦表,确认仍前向、参数量降 786432。
3. 贪心标志:给生成加 greedy=True 跳采样取 argmax,确认跨运行确定。
4. 重复惩罚:加 repetition_penalty 旋钮,softmax 前把提示或历史中任何 token 的 logit 除常数,固定提示上 >1 降重复。
5. top_p:在 top_k 旁加 top_p(核采样),两行检查保留 token 概率和超 top_p。

本节要点回顾

1. 124M 配置:词表 50257、上下文 1024、嵌入 768、12 头、12 层,数对不上说明接错。
2. weight tying:LM 头与 token 嵌入同张量,省 vocab×d_model 参数,嵌入与头一起更新。
3. 学习位置嵌入:(1024, 768) 参数张量,最简位置方案,124M 参考用。
4. 生成三旋钮:温度(近零贪心)、top-k(一贪心、四十滤尾)、多项采样。
5. 残差投影小初始化:std×1/sqrt(2*num_layers),残差流深 12 保 sane。
6. 参数级绑权重:lm_head.weight = token_embedding.weight 共享张量非拷贝。

下一节,我们建「训练循环与评估」——交叉熵损失、AdamW、生成作定性信号,把这个模型训起来。