第 4 章 模型构建（model.py）

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

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

本项目刻意不手写 Transformer，原因写在 model.py 顶部注释里：

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

工程取舍

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

4.2 完整源码

"""
model.py
========
模型构建模块。

本项目不手写 Transformer，而是基于 HuggingFace transformers 提供的
先进小型 GPT (GPT2LMHeadModel) 实现。通过传入自定义 GPT2Config，
我们能够完全控制模型规模（层数、头数、维度等），使其与 config.py 中的
GPTConfig 保持一致。GPT2 的实现经过工业级验证，稳定且高效。
"""

import os
from typing import Optional

import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Config

from config import GPTConfig


def build_model(gpt_config: GPTConfig) -> GPT2LMHeadModel:
    """
    根据项目 GPTConfig 构建一个从零初始化的小型 GPT2 模型。

    参数:
        gpt_config: 项目自定义的 GPT 配置。

    返回:
        transformers.GPT2LMHeadModel 实例（权重随机初始化）。
    """
    # 将项目配置映射为 HuggingFace GPT2Config。
    hf_config = GPT2Config(**gpt_config.to_gpt2_kwargs())
    model = GPT2LMHeadModel(hf_config)
    return model


def count_parameters(model: nn.Module) -> int:
    """统计模型中可训练参数数量。"""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def load_model(
    checkpoint_path: str,
    gpt_config: Optional[GPTConfig] = None,
    device: Optional[str] = None,
    map_location: str = "cpu",
) -> GPT2LMHeadModel:
    """
    从磁盘加载训练好的模型权重。

    支持两种 checkpoint 形式：
        1) 单个权重文件 (.pt/.bin)：可为完整 dict（含 "model_state_dict" 键）
           或纯 state_dict。
        2) transformers 风格的预训练目录。

    参数:
        checkpoint_path: 权重文件路径或目录。
        gpt_config:      模型架构配置；若为 None 则使用默认配置。
        device:          加载后放置到的设备。
        map_location:    torch.load 的 map_location 参数。
    """
    if gpt_config is None:
        gpt_config = GPTConfig()

    if os.path.isdir(checkpoint_path):
        # transformers 风格目录：直接用 from_pretrained 加载。
        model = GPT2LMHeadModel.from_pretrained(checkpoint_path)
    else:
        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"[model] 加载权重：缺失键 {len(missing)} 个，多余键 {len(unexpected)} 个")

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

整个文件只有 3 个函数，但每个都值得细读。

4.3 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 预训练权重。

对比加载预训练权重的写法：

# 加载 OpenAI 预训练的 gpt2（124M）
model = GPT2LMHeadModel.from_pretrained("gpt2")
​

本项目走的是 from scratch（从零训练）路线。

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

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

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

4.4 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：只统计需要梯度（会更新）的参数，排除冻结层。

用法

train.py 里：

n_params = count_parameters(model)
print(f"[train] 模型可训练参数量: {n_params/1e6:.2f}M")
# 输出: 模型可训练参数量: 24.97M
​

训练前先打印参数量是个好习惯，能发现：

• 模型规模配错了（本想要 100M，结果只有 1M）
• 参数被意外冻结了（requires_grad=False）
• 某层没接上（参数量比预期少很多）

参数量估算公式

GPT-2 架构的参数量粗略公式：

参数量 ≈ 12 × n_layer × n_embd² + vocab_size × n_embd
​

代入本项目（n_layer=6, n_embd=384, vocab_size=50257）：

12 × 6 × 384² + 50257 × 384
= 12 × 6 × 147456 + 19298528
≈ 10.6M + 19.3M
≈ 30M  （实际 25M，因为 lm_head 与 wte 权重共享）
​

权重共享（tie weights）省了一份 vocab_size × n_embd 的参数，所以实际比估算少一些。

4.5 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 文件（本项目 train.py 保存的格式）
        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。

单文件格式的兼容处理

state = torch.load(checkpoint_path, map_location=map_location)
if isinstance(state, dict) and "model_state_dict" in state:
    state = state["model_state_dict"]
​

本项目的 checkpoint 是这种结构（见第 5 章 save_checkpoint）：

{
    "step": 500,
    "model_state_dict": {...权重...},
    "optimizer_state_dict": {...},
    "scheduler_state_dict": {...},
    "gpt_config": {...},
    "train_config": {...},
}
​

加载时取 model_state_dict 子表。但也兼容「纯 state_dict」格式（直接保存权重，没包外层 dict）。

strict=False 的意义

missing, unexpected = model.load_state_dict(state, strict=False)
​

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

何时用 strict=False：

• 「模型架构微调后加载旧 checkpoint」（如加了一层、改了维度）
• 容错加载（万一 checkpoint 部分损坏）

警告：如果关键层（如所有 Transformer Block）都在 missing 里，模型其实是没加载成功的，等于从零开始。必须看打印的数字判断。

if missing or unexpected:
    print(f"缺失键 {len(missing)} 个，多余键 {len(unexpected)} 个")
​

经验值：

• missing=0, unexpected=0：完美加载 ✓
• missing=几百, unexpected=0：可能有问题，检查架构是否变了
• missing=几百, unexpected=几百：架构完全不匹配，别用了

model.eval() 不能忘

model.eval()
​

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

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

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

虽然不手写，但理解内部结构有助于调参。一个 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)

wte = nn.Embedding(vocab_size, n_embd)   # (50257, 384)
​

本质是一张大表：50257 行（每个 token 一行）× 384 列（嵌入维度）。input_ids 是索引，查表得到每个 token 的 384 维向量。

Position Embedding (wpe)

wpe = nn.Embedding(n_positions, n_embd)  # (128, 384)
​

GPT 用学出来的绝对位置编码（不是 Transformer 原论文的正弦编码）。位置 0 到 127 各有一个 384 维向量。

对比：Transformer 原论文用固定的正弦/余弦编码；LLaMA 用旋转位置编码 RoPE。各有利弊。

两个 Embedding 相加

hidden = wte(input_ids) + wpe(position_ids)   # (B, T, n_embd)
​

词义 + 位置信息融合。

Transformer Block

每个 block 做三件事：

1. 多头因果注意力：让每个位置「看」之前所有位置，聚合信息。
2. 前馈网络（FFN）：两层 MLP，中间扩展到 4×n_embd 维度。
3. 残差连接 + LayerNorm：稳定训练。

堆 n_layer 个 block 就是完整 Transformer。

因果掩码（causal mask）

注意力的核心：

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

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

用矩阵表示就是一个下三角矩阵：

mask = 
[[1, -inf, -inf, -inf],
 [1,   1,  -inf, -inf],
 [1,   1,    1,  -inf],
 [1,   1,    1,    1  ]]
​

加到注意力分数上，softmax 后未来位置的权重就是 0。这是 GPT 能做自回归生成的根本保证。

lm_head 与权重共享

lm_head = nn.Linear(n_embd, vocab_size, bias=False)   # (384, 50257)
​

把 384 维 hidden state 映射回 50257 维 logits（每个 token 的得分）。

关键：GPT-2 默认让 lm_head.weight = wte.weight（权重共享，tie weights）：

• 推理方向：hidden → logits（lm_head）
• 训练方向：token id → hidden（wte）
• 同一张表两个方向用

好处：

1. 参数量省 vocab_size × n_embd（本项目省 19M，占总量 40%+）。
2. 语义上合理：词义空间一致，能生成好 token 的方向也能被好 token 触发。

4.7 forward 怎么算 loss

train.py 里这一行：

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)
​

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

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

4.8 动手实验

1. 构建并打印模型：

   from config import GPTConfig
   from model import build_model, count_parameters
   
   m = build_model(GPTConfig())
   print(count_parameters(m) / 1e6, "M")
   print(m)   # 打印完整结构
   ​

   数一数实际参数量与 25M 的差距。理解估算公式。

2. 规模扫描：跑这些配置，记录参数量：

   configs = [
       GPTConfig(n_layer=2,  n_embd=64,  n_head=2),
       GPTConfig(n_layer=6,  n_embd=384, n_head=6),   # 默认
       GPTConfig(n_layer=12, n_embd=768, n_head=12),  # GPT-2 small
   ]
   for cfg in configs:
       m = build_model(cfg)
       print(f"{cfg.n_layer}x{cfg.n_embd}: {count_parameters(m)/1e6:.1f}M")
   ​

3. 检查权重共享：

   m = build_model(GPTConfig())
   print(m.lm_head.weight is m.transformer.wte.weight)   # True 表示共享
   ​

4. 体验 strict=True vs False：故意构造一个少一个键的 state_dict：

   m = build_model(GPTConfig())
   state = m.state_dict()
   del state[list(state.keys())[0]]   # 删一个键
   
   m2 = build_model(GPTConfig())
   # strict=True 会报错
   # m2.load_state_dict(state, strict=True)
   # strict=False 只警告
   missing, unexpected = m2.load_state_dict(state, strict=False)
   print(f"missing: {len(missing)}")
   ​

5. 思考题：为什么 GPT 用 LayerNorm 而不是 BatchNorm？（提示：序列长度可变、batch 内不同位置语义独立）

4.9 小结

• 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.10 下一章

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