# 加载预训练权重

> 从零训练一个 1.24 亿参数模型，是预算问题；加载一份已经公开的 checkpoint，只该是周二下午的工作。这节课会把 GPT-2 风格的预训练权重从 safetensors 文件装进第 35 课那套完全一致的架构里，一步一步走通参数名映射，再用一次 sanity generation 证明加载真的生效。没有网络、没有第三方 loader、也没有黑盒魔法。

**类型：** Build
**语言：** Python
**前置要求：** 第 19 阶段第 30-36 课
**预计时间：** ~90 分钟

## 学习目标

- 用 `safetensors` Python 库读 safetensors 文件，并检查里面每个 tensor 的名字与形状。
- 把每一个预训练参数名映射到第 35 课 GPT 模型里的本地参数。
- 正确处理两套命名约定之间的差异：发布版 GPT-2 的 `wte/wpe/h.N.attn.c_attn/c_proj`、`mlp.c_fc/c_proj`，以及本课程本地模型的 `tok_embed/pos_embed/blocks.N.attn.qkv/out_proj`、`mlp.fc1/fc2`。
- 在任何权重赋值发生前，先检测并拒绝 shape mismatch，并给出清晰错误。
- 用加载后的权重生成一小段 continuation，确认 token 分布真的来自加载好的参数，而不是随机初始化。

## 问题背景

公开 checkpoint 从来不是按你的架构命名打包的。它们带的是原始实现自己的参数名。比如预训练文件里有：

`transformer.h.0.attn.c_attn.weight`，形状 `(2304, 768)`

而你的模型可能期待的是：

`blocks.0.attn.qkv.weight`，形状同样是 `(2304, 768)`，或者你的 `nn.Linear` 内部权重还把矩阵按另一种方向存了。于是同一块参数会同时有三种身份：名字、形状、字节布局。loader 的职责，就是把三者重新对齐。

一个瞎 copy 的 loader，会把对的 tensor 装进错的槽里，最后模型吐垃圾。一个遇到 shape mismatch 就偷偷跳过、但什么也不记的 loader，会让你完全不知道哪块权重没落下。本课的 loader 走的是硬路径：每次赋值都记日志，每个 shape 都先验证，最后给你一份 `LoadReport`，明白告诉你命中了哪些、漏了哪些、炸了哪些。

## 核心概念

```mermaid
flowchart LR
  SF[safetensors 文件<br/>gpt2-stub.safetensors] --> R[读取器<br/>safe_open]
  R --> N[参数名迭代器]
  N --> M[名称映射器<br/>预训练 -> 本地]
  M --> S[形状检查]
  S -- 匹配 --> A[赋值张量<br/>在 torch.no_grad 下]
  S -- 不匹配 --> E[记录不匹配<br/>不赋值]
  A --> RP[LoadReport]
  E --> RP
  RP --> G[生成<br/>健全性检查样本]
```

name mapper 本质上就是一个 `str -> str` 的函数。shape check 只是一个 if。真正赋值时放在 `torch.no_grad()` 下，避免 autograd 跟踪加载动作。最后 report 把每个名字的命运都记下来。

### GPT-2 的命名约定

公开 GPT-2 权重常见名字如下：

| 预训练名 | 形状 | 含义 |
|---|---|---|
| `wte.weight` | `(50257, 768)` | Token embedding |
| `wpe.weight` | `(1024, 768)` | Position embedding |
| `h.N.ln_1.weight` | `(768,)` | 第 N 个 block 的 LayerNorm 1 scale |
| `h.N.ln_1.bias` | `(768,)` | 第 N 个 block 的 LayerNorm 1 shift |
| `h.N.attn.c_attn.weight` | `(768, 2304)` | 融合 QKV 线性层权重 |
| `h.N.attn.c_attn.bias` | `(2304,)` | 融合 QKV 线性层 bias |
| `h.N.attn.c_proj.weight` | `(768, 768)` | Attention 输出投影 |
| `h.N.attn.c_proj.bias` | `(768,)` | Attention 输出投影 bias |
| `h.N.ln_2.weight` | `(768,)` | LayerNorm 2 scale |
| `h.N.ln_2.bias` | `(768,)` | LayerNorm 2 shift |
| `h.N.mlp.c_fc.weight` | `(768, 3072)` | MLP fc1 权重 |
| `h.N.mlp.c_fc.bias` | `(3072,)` | MLP fc1 bias |
| `h.N.mlp.c_proj.weight` | `(3072, 768)` | MLP fc2 权重 |
| `h.N.mlp.c_proj.bias` | `(768,)` | MLP fc2 bias |
| `ln_f.weight` | `(768,)` | 最终 LayerNorm scale |
| `ln_f.bias` | `(768,)` | 最终 LayerNorm shift |

这里有两个坑最需要小心：

- `c_attn`、`c_proj`、`c_fc` 这些线性层的矩阵布局，与 `nn.Linear.weight` 默认期待的方向是转置关系，加载时要转置。
- LM head 根本不在文件里，因为它通过 `wte` 做 weight tying。只要 `wte` 装好了，head 本身应该通过 alias 指过去。

### 本地模型的命名约定

这套课程里的模型名字更偏描述式：

| 本地名 | 含义 |
|---|---|
| `tok_embed.weight` | Token embedding |
| `pos_embed.weight` | Position embedding |
| `blocks.N.ln1.scale` | 第 N 个 block 的 LayerNorm 1 scale |
| `blocks.N.ln1.shift` | LayerNorm 1 shift |
| `blocks.N.attn.qkv.weight` | Fused QKV |
| `blocks.N.attn.qkv.bias` | Fused QKV bias |
| `blocks.N.attn.out_proj.weight` | Attention 输出投影 |
| `blocks.N.attn.out_proj.bias` | 输出投影 bias |
| `blocks.N.ln2.scale` | LayerNorm 2 scale |
| `blocks.N.ln2.shift` | LayerNorm 2 shift |
| `blocks.N.mlp.fc1.weight` | MLP fc1 |
| `blocks.N.mlp.fc1.bias` | MLP fc1 bias |
| `blocks.N.mlp.fc2.weight` | MLP fc2 |
| `blocks.N.mlp.fc2.bias` | MLP fc2 bias |
| `final_ln.scale` | 最终 LayerNorm scale |
| `final_ln.shift` | 最终 LayerNorm shift |

映射关系是固定函数，本课会直接把它做成 dict，并按层展开。

### Stub Fixture

真实 GPT-2 权重大约 0.5 GB。本课不会下载它，而是在首次运行时生成一个小型 safetensors fixture：命名完全遵守 GPT-2 习惯，但层数和宽度会缩成便于本地验证的版本（比如 12 层、`d_model=192`）。这样 fixture 的结构足以打通 loader 的所有路径。换成真实文件时，loader 本身不用改。

```figure
cc-weight-remap
```

## 动手构建

`code/main.py` 会实现：

- 一份本地版的第 35 课 `GPTModel`，保证本课自洽
- `make_pretrained_to_local(num_layers)`：按层展开 name map
- `load_safetensors(model, path)`：遍历名字、映射、检查 shape、对 conv1d 风格权重做转置，在 `torch.no_grad()` 下赋值，并返回 `LoadReport`
- `make_stub_safetensors(path, cfg)`：按预训练命名习惯生成 fixture 文件
- 一个 demo：首次运行时生成 `outputs/gpt2-stub.safetensors`，构造一个 fresh model，先从随机初始化下生成一段，再加载 stub，再生成一段，并验证两者不同，说明加载真的改变了模型

运行方式：

```bash
python3 code/main.py
```

输出包括：fixture 路径、逐参数加载日志、`LoadReport` 摘要、加载前的 continuation、加载后的 continuation，以及一条故意注入的坏 tensor 所触发的 shape mismatch，用来覆盖失败路径。

## Stack

- `safetensors`：负责磁盘格式与流式读取
- `torch`：负责模型和赋值数学
- 不用 `transformers`，不用 `huggingface_hub`，不做任何网络调用

## 生产里常见的三个模式

**先验证整个文件，再开始任何赋值。** 正确流程是：打开文件、列出所有 tensor 名、dtype、shape，跑完整 name mapping 和 shape check，只有全部通过后才真正写进模型。半加载模型就是静默失败制造机。

**每次赋值都记 source name 和 destination name。** 出问题时，这份 log 会直接告诉你哪块 tensor 被写到了哪里，不然你只能去读 hexdump。`LoadReport` 至少应该跟踪：`loaded`、`missing`、`unexpected`、`shape_mismatch`。

**LM head 是 weight tying alias，不是独立 copy。** 标准写法是：`model.lm_head.weight = model.tok_embed.weight`。若你把 embedding 再 copy 一份给 head，绑定就断了，参数数目悄悄翻倍。

## 实际使用

- 只要 safetensors 文件遵守 GPT-2 命名规范，这个 loader 就能直接用。真实 GPT-2 small / medium / large / xl 只是在 config 上不同，代码不需要变。
- 同样的套路也能扩展到 LLaMA、Mistral、Qwen，只要更新 name map 即可。shape check 和 report 逻辑完全不变。
- 加载后做一次 sanity generation，是最快的门槛测试：若加载后样本和随机初始化时几乎没区别，那就是 loader 根本没真正改到模型。

## 练习

1. 给 loader 加一个 `dtype` 参数，加载时可选强转成 `bfloat16`、`float16` 或 `float32`。验证 `float32` 模型降成 `bfloat16` 后仍能生成。
2. 加一个 `expected_layers` 参数，若 checkpoint 里的 `h.N` 索引与模型 `num_layers` 不符，就直接拒绝加载。
3. 把 loader 接回第 35 课的 generation 函数，打印“随机初始化样本 vs 加载权重样本”的并排对比。
4. 再补一条导出路径：把当前模型 state 按预训练命名规范写回新 safetensors 文件。再 round-trip 一次，验证 report 中 shape mismatch 为 0。
5. 扩展 `NAME_MAP` 去适配 LLaMA 风格命名（无 bias、RMSNorm、fused qkv），并用你自己生成的 stub LLaMA fixture 复跑 loader。

## 关键术语

| 英文 | 大家嘴上怎么说 | 它实际指什么 |
|---|---|---|
| Name map | “Key remapping” | 从预训练 tensor 名到本地参数名的映射函数，通常是一张按层展开的字典 |
| Shape mismatch | “Bad shape” | tensor 名能映射上，但预训练张量形状与本地参数不一致；loader 会拒绝赋值并记录 |
| Transpose-on-load | “Conv1d layout” | GPT-2 发布权重里 attention / MLP projection 的矩阵布局与 `nn.Linear` 期待方向相反，因此加载时要转置 |
| Weight tying alias | “Shared LM head” | 让 `model.lm_head.weight = model.tok_embed.weight`，head 与 embedding 共用同一块存储 |
| Load report | “Coverage summary” | 一个 dataclass，汇总 loaded、missing、unexpected、shape_mismatch；它就是你判断加载是否成功的证据 |

## 延伸阅读

- 第 19 阶段第 35 课：接收这些权重的模型架构
- 第 19 阶段第 36 课：产出同形 checkpoint 的训练 loop
- 第 10 阶段第 11 课：内存紧张时对已加载权重做量化
- 第 10 阶段第 13 课：围绕加载与推理的一整条 LLM pipeline