第 4 章 模型构建

第 4 章 模型构建

本章目标：理解如何用 HuggingFace GPT2LMHeadModel 搭一个自定义规模的小 GPT，以及如何加载/保存权重。

4.1 设计哲学：不重复造轮子

本项目刻意不手写 Transformer，而是基于 HuggingFace Transformers 提供的先进小型 GPT（GPT2LMHeadModel）实现。GPT2 的实现经过工业级验证，稳定且高效。

工程取舍

本项目选择后者，把精力集中在「工程流水线」上。想学手写 Attention，看第 9 章进阶方向。

4.2 build_model：三行核心代码

def build_model(gpt_config: GPTConfig) -> GPT2LMHeadModel:
    hf_config = GPT2Config(**gpt_config.to_gpt2_kwargs())  # ① 配置映射
    model = GPT2LMHeadModel(hf_config)                     # ② 实例化（随机初始化）
    return model
​

就这么简单。三件事：

① 配置映射

hf_config = GPT2Config(**gpt_config.to_gpt2_kwargs())
​

to_gpt2_kwargs() 在第 2 章讲过：把项目自己的 GPTConfig 字段名翻译成 HuggingFace 的 GPT2Config 字段名。** 把 dict 解包成关键字参数。

② 实例化（从零初始化）

model = GPT2LMHeadModel(hf_config)
​

注意——这里不传 from_pretrained，所以权重是从零随机初始化的，不是加载 OpenAI 预训练权重。对比加载预训练权重的写法：model = GPT2LMHeadModel.from_pretrained("gpt2")。本项目走的是 from scratch（从零训练）路线。

③ 为什么「从零训练」能在 1MB 数据上学到东西

因为模型也小（25M）。模型规模与数据量必须匹配：

小模型 + 小数据 = 教学黄金组合。

4.3 count_parameters：参数量统计

def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters() if p.requires_grad)
​

逐符号解释

• model.parameters()：迭代器，遍历模型所有权重张量。
• p.numel()：number of elements，张量元素总数。比如 shape=(384, 50257) 的张量，numel = 384×50257 ≈ 1900 万。
• if p.requires_grad：只统计需要梯度（会更新）的参数，排除冻结层。

训练前先打印参数量是个好习惯，能发现：模型规模配错了、参数被意外冻结了、某层没接上（参数量比预期少很多）。

参数量估算公式

GPT-2 架构的参数量粗略公式：参数量 ≈ 12 × n_layer × n_embd² + vocab_size × n_embd。

代入本项目（n_layer=6, n_embd=384, vocab_size=50257）：约 30M。实际 25M，因为 lm_head 与 wte 权重共享省了一份词表参数。

4.4 load_model：两种 checkpoint 格式

def load_model(checkpoint_path, gpt_config=None, device=None, map_location="cpu"):
    if gpt_config is None:
        gpt_config = GPTConfig()

    if os.path.isdir(checkpoint_path):
        # 格式 A：transformers 风格目录
        model = GPT2LMHeadModel.from_pretrained(checkpoint_path)
    else:
        # 格式 B：单个 .pt 文件（本项目训练入口保存的格式）
        model = build_model(gpt_config)
        state = torch.load(checkpoint_path, map_location=map_location)
        if isinstance(state, dict) and "model_state_dict" in state:
            state = state["model_state_dict"]
        missing, unexpected = model.load_state_dict(state, strict=False)
        if missing or unexpected:
            print(f"加载权重：缺失键 {len(missing)} 个，多余键 {len(unexpected)} 个")

    if device:
        model.to(device)
    model.eval()
    return model
​

两种 checkpoint 格式

判断方式很简单：os.path.isdir(path) 是目录就走 A，否则走 B。

strict=False 的意义

• strict=True（默认）：权重字典与模型结构必须完全匹配，任何不一致都报错。
• strict=False：允许不匹配，返回两个列表：missing（模型有但 checkpoint 没有）、unexpected（checkpoint 有但模型没有）。

警告：如果关键层（如所有 Transformer Block）都在 missing 里，模型其实是没加载成功的，等于从零开始。必须看打印的数字判断。经验值：missing=0, unexpected=0 完美加载；几百/几百 架构完全不匹配。

model.eval() 不能忘

• 训练模式（model.train()）：dropout 生效、BatchNorm 用 batch 统计
• 推理模式（model.eval()）：dropout 关闭、用全局统计

忘切 eval() 会导致推理结果每次都不一样（dropout 随机置零），调试时超痛苦。

4.5 GPT2LMHeadModel 内部长什么样（科普）

虽然不手写，但理解内部结构有助于调参：

input_ids (B, T)
    │
    ├─► wte: Embedding(vocab_size, n_embd)        词嵌入
    ├─► wpe: Embedding(n_positions, n_embd)       位置嵌入（GPT 用绝对位置）
    │
    ▼  相加
┌─────────────────────────────────────────┐
│  Transformer Block × n_layer             │
│  ┌─────────────────────────────────────┐ │
│  │ LayerNorm                            │ │
│  │ Multi-Head Causal Self-Attention    │ │ ◄── 下三角掩码，看不到未来
│  │   + Residual                         │ │
│  │ LayerNorm                            │ │
│  │ Feed-Forward (n_embd → 4×n_embd → n_embd) │ │
│  │   + Residual                         │ │
│  └─────────────────────────────────────┘ │
└─────────────────────────────────────────┘
    │
    ▼
LayerNorm
    │
    ▼
lm_head: Linear(n_embd, vocab_size, bias=False)   映射回词表
    │
    ▼
logits (B, T, vocab_size)
​

关键组件逐个讲

Token Embedding (wte)：本质是一张大表（50257 行 × 384 列），input_ids 是索引，查表得到每个 token 的 384 维向量。

Position Embedding (wpe)：GPT 用学出来的绝对位置编码（不是 Transformer 原论文的正弦编码）。位置 0 到 127 各有一个 384 维向量。对比：Transformer 原论文用固定的正弦/余弦编码；LLaMA 用旋转位置编码 RoPE。

两个 Embedding 相加：词义 + 位置信息融合，hidden = wte(input_ids) + wpe(position_ids)。

Transformer Block：每个 block 做三件事——多头因果注意力（让每个位置「看」之前所有位置）、前馈网络 FFN（两层 MLP，中间扩展到 4×n_embd）、残差连接 + LayerNorm（稳定训练）。

因果掩码（causal mask）

标准注意力：每个位置看所有位置
  位置 0 → 看 [0,1,2,3]   ← 但生成时位置 0 不能看未来！

因果掩码（下三角）：
  位置 0 → 只看 [0]
  位置 1 → 只看 [0,1]
  位置 2 → 只看 [0,1,2]
  位置 3 → 只看 [0,1,2,3]
​

用矩阵表示就是一个下三角矩阵，加到注意力分数上，softmax 后未来位置的权重就是 0。这是 GPT 能做自回归生成的根本保证。

lm_head 与权重共享

lm_head 把 384 维 hidden state 映射回 50257 维 logits。关键：GPT-2 默认让 lm_head.weight = wte.weight（权重共享，tie weights）——推理方向 hidden → logits，训练方向 token id → hidden，同一张表两个方向用。好处：参数量省 vocab_size × n_embd（本项目省 19M），且词义空间一致。

4.6 forward 怎么算 loss

训练循环里这一行：

outputs = model(input_ids=x, labels=y)
loss = outputs.loss
​

为什么传 labels 就能拿 loss？看 GPT2LMHeadModel.forward 内部简化逻辑：

def forward(self, input_ids, labels=None):
    hidden = self.transformer(input_ids)         # (B, T, n_embd)
    logits = self.lm_head(hidden)                # (B, T, vocab_size)

    if labels is not None:
        # 内部算交叉熵
        shift_logits = logits[..., :-1, :].contiguous()   # 去掉最后一个位置
        shift_labels = labels[..., 1:].contiguous()       # 去掉第一个位置
        loss = CrossEntropyLoss()(shift_logits.view(-1, V), shift_labels.view(-1))

    return outputs(loss=loss, logits=logits)
​

注意它内部也做了「logits 去尾、labels 去头」的对齐（与数据层的 y=chunk[1:] 一致），所以传进去的 x 和 y 长度都是 block_size 即可。

💡 这种「传 labels 自动算 loss」的 API 很方便，但要知道它内部干了什么。想自定义 loss（如 label smoothing）就要传 labels=None 自己算。

4.7 动手实验

1. 构建并打印模型：build_model(GPTConfig())，打印 count_parameters(m)/1e6 和完整结构，数一数实际参数量与 25M 的差距。
2. 规模扫描：跑这些配置，记录参数量：2×64、6×384（默认）、12×768（GPT-2 small）。
3. 检查权重共享：m.lm_head.weight is m.transformer.wte.weight 返回 True 表示共享。
4. 体验 strict=True vs False：构造一个少一个键的 state_dict，分别用 strict=True（报错）和 strict=False（只警告）加载。
5. 思考题：为什么 GPT 用 LayerNorm 而不是 BatchNorm？（提示：序列长度可变、batch 内不同位置语义独立）

4.8 小结

• build_model 三步：to_gpt2_kwargs() 映射配置 → GPT2Config → GPT2LMHeadModel（从零初始化）。
• count_parameters 用 p.numel() for p in parameters() if p.requires_grad 统计可训练参数。
• load_model 支持两种格式：HF 目录（from_pretrained）和单文件 .pt（build_model + load_state_dict）。
• strict=False 允许权重与结构不严格匹配，但要看打印的 missing/unexpected 数量判断是否真的加载成功。
• model.eval() 推理时不能忘（关 dropout）。
• GPT-2 内部 = 词嵌入 + 位置嵌入 + N 个 Transformer Block + lm_head（与 wte 权重共享）。

4.9 下一章

模型搭好了，去《第 5 章 训练循环》看怎么把它训起来——这是最重的一章。