第 15 章　权重格式、参数映射与低内存加载

下载一个权重文件和得到正确模型之间，还隔着结构、名字、形状、布局、dtype、设备及共享关系。只检查 load_state_dict 没有报错，不能证明 tokenizer 与模型兼容；只检查每层形状相同，也不能发现 Q、K 被交换或方阵忘记转置。

本章覆盖 ch05 主代码的保存恢复、TensorFlow GPT-2 导入、三种替代加载 notebook，以及低内存 state_dict notebook。所有实际运行只读写本书自己创建的小检查点，不下载预训练权重。替代加载

15.1 先区分推理恢复与训练恢复

推理需要结构配置、参数与必要 buffer、tokenizer 及模板约定。恢复训练还需要优化器状态、已完成更新步数、学习率调度状态、随机数状态，以及数据进度。AdamW 的一阶、二阶动量丢失后，下一步更新通常不同，即使权重完全一样。

上游示例分别保存 model.state_dict() 和包含 model_state_dict/optimizer_state_dict 的字典。恢复顺序是建立模型、载入模型状态、用新模型参数建立优化器、再载入优化器状态。不能把旧优化器继续指向一个已经被替换掉的模型对象。

本书 checkpoint_checks.py 在小模型训练 3 步后保存模型、优化器、CPU RNG、配置和 step，继续第 4 步作为基线；另一条路径加载后做相同第 4 步，全部浮点参数最大差为 0。这里数据顺序固定、dropout 为零，没有验证多进程 DataLoader 的精确续跑。记录哪些状态尚未覆盖，比笼统宣称“完全可复现”更有用。

15.2 官方 GPT-2 的名字怎样映射到教学模型

上游 TensorFlow checkpoint 读取器列举变量名、读取数组、去掉多余单维，把 model/ 前缀后的路径拆开。形如 h0/... 的变量放进第 0 个 block，其他项进入全局参数字典。setdefault 逐级创建嵌套字典，最后存实际数组。读取与映射代码

对应关系如下：

源参数	教学模型目标	必要变换
wte	token embedding、输出头初值	[V,D] 直接复制
wpe	position embedding	[C,D] 直接复制
attn/c_attn/w	Q、K、V 权重	最后一轴拆三份，再各自转置
attn/c_attn/b	Q、K、V 偏置	最后一轴拆三份
attn/c_proj/w	注意力输出投影	转置
mlp/c_fc/w	FFN 第一层	转置
mlp/c_proj/w	FFN 第二层	转置
ln_1/ln_2 的 g,b	两层 LayerNorm scale、shift	直接复制
全局 g,b	final LayerNorm	直接复制

源合并 QKV 权重为 [D,3D]，拆分后每份 [D,D]。因为方阵转置后形状不变，形状检查无法发现漏转置。FFN [D,4D] 的转置更容易被形状检查抓住。应使用有不同元素的测试矩阵，而不是全零或单位矩阵来验证布局。

载入 GPT-2 参数前，上游把上下文容量恢复为 1024，并启用 QKV 偏置。训练示例的 256 上下文和无 QKV 偏置配置不能直接接收全部原权重。

15.3 assign 会替换参数对象

上游 assign(left,right) 先检查 shape，再构造新的 nn.Parameter 返回给模块属性。TensorFlow 数组通过 torch.tensor 转成张量；Hugging Face 变体使用 clone().detach()；safetensors 变体使用 detach()，可能保留底层存储共享。这些选择影响复制量与别名关系。

替换 Parameter 后，之前创建的优化器可能仍持有旧对象。配套测试明确演示了这一点：把 token embedding 属性换成同数值的新 Parameter，旧优化器仍引用旧权重，且不包含新权重。安全顺序是完成所有替换后再建优化器，或在已有参数对象上用 no-grad 的 copy_ 保持身份。

输出头与 token embedding 都从 wte 加载，不一定意味着同一 Parameter。复制数值、共享存储、共享参数对象是三个层次。训练中的真正权重绑定应显式核对对象关系和梯度路径，而非只比较初始化数值。

15.4 三种替代文件入口

weight-loading-pytorch.ipynb 使用已经转成教学模型命名的 .pth，因此可以直接 load_state_dict；同一 notebook 也演示同命名的 safetensors 文件。它省去了外部参数名转换，但仍需要正确模型配置。

weight-loading-hf-transformers.ipynb 通过 GPT2Model.from_pretrained 获得 Hugging Face 模型，再从其 state_dict 按 h.{b}.attn.c_attn.weight 等名字映射。源模型没有独立语言建模头时，目标输出头使用 token embedding 权重。该版本在 block 循环内重复设置 final norm 和输出头，最终结果可以相同，但重复赋值和复制不是必要工作，搬到循环外更清楚。

weight-loading-hf-safetensors.ipynb 直接读取 Hugging Face model.safetensors，不先实例化完整源模型，随后进行相同语义映射。文件格式避免了通用 Python 对象反序列化，但不能保证权重来源可靠、架构正确或模型行为安全。固定 notebook 对其中一个尺寸附有生成问题注释，本书没有下载核验，因此不把所有尺寸写成已通过。

TensorFlow 路径还依赖 checkpoint 元数据和相关库；直接 PyTorch/safetensors 路径减少了格式桥接步骤。它们解决的是不同输入文件的适配，不是三种不同 GPT 结构。

15.5 下载器的成功条件不能只看文件存在

上游下载器验证允许的模型尺寸，创建目录，下载 checkpoint、词表、超参数和数据文件。流式读取、超时和 HTTP 状态检查能避免部分网络错误。某些版本以 Content-Length 与本地文件大小相同作为跳过条件；同大小不证明内容相同，正式归档应记录来源 revision 与内容哈希。

独立 gpt_download.py 还提供备用镜像：内部 _attempt_download 先试主地址，网络异常时再试备用地址；全部失败时打印错误，但调用方仍应显式检查文件齐全，不能继续把缺失文件交给 TensorFlow。gpt_generate.py 内嵌的下载版本没有同样的备用参数，这是两个重复实现的实际差异。

直接覆盖目标文件可能在中断后留下不完整文件。更稳健的流程是写临时文件、校验、再原子替换，并核对所有分片齐全。本书未实际下载任何大文件，相关函数只做源代码解释。不能把远端 main 文件路径当成永远不变的版本标识。

15.6 加载峰值为何可能接近两份权重

常见路径先建立随机模型并移到 GPU，再把整个 state_dict 加载到 GPU，最后逐项复制。加载期间，目标参数与源字典可能同时存在，相当于额外一整份权重；初始化还消耗时间和随机数。

上游内存 notebook 比较了多种策略：先在 CPU 建模和加载再整体迁移；GPU 模型配 CPU state_dict 逐张量复制；meta 构建后 to_empty；meta 加 mmap 与 assign=True；以及每个参数独立文件、逐个读取。不同方法把成本转移到 CPU、GPU、文件映射或小张量临时空间，不应只看最终常驻内存。

meta 张量只有形状、dtype 等元信息，没有可用数值；to_empty 分配未初始化存储，不会自动重建原本的三角 mask。不能对未初始化模型直接推理。

15.7 低内存加载不能漏掉 buffer

上游部分“逐参数复制”示例只遍历 named_parameters()，而保存时 state_dict 包含持久 buffer。若目标模型正常初始化，因果 mask 可能恰好已正确生成，遗漏暂时不暴露；若模型在 meta 上构造再 to_empty，mask 也变成未初始化存储，只复制参数就不够了。

本书测试用确定性的全 False mask 代表“没有正确恢复的 mask”，避免读取随机未初始化内存。参数全部正确时，这个错误 mask 仍使 logits 最大偏差达到 0.433731。它证明前向所需的 buffer 必须一起恢复，不意味着未初始化内存每次恰好产生这个数字。

完整 load_state_dict 能统一处理参数和持久 buffer，并报告缺失或多余键。若确实逐个文件加载，应遍历完整状态、核对每项 dtype/shape/device，并恢复可重建的非持久状态。仅打印 missing warning 后继续使用模型不适合作为严格验收。

15.8 mmap 与 assign 的小型实测

本书在 CPU 上运行：

with torch.device("meta"):
    model = TinyGPT(**config)
state = torch.load(path, map_location="cpu", weights_only=True, mmap=True)
model.load_state_dict(state, assign=True)

mmap=True 让支持格式的本地存储可由文件映射，assign=True 让模块采用 state_dict 的张量属性，避免先创建一套随机参数再逐项复制。所有持久 buffer 同样加载。测试确认没有参数留在 meta，模型输出与原模型逐位相同。

采用 assign 后应在加载完成之后创建优化器，并核对共享权重关系。文件映射也不意味着数据永远不占物理内存；访问页面时操作系统仍需提供实际数据。本书没有报告峰值 RSS 或 GPU 内存节省量，只验证该 API 路径在 PyTorch 2.6.0+cpu 的正确性。

15.9 如何解读内存监控代码

上游 start_memory_tracking/print_memory_usage/cleanup 使用 CUDA 峰值统计、垃圾回收和缓存清理；没有 CUDA 时，部分后续调用仍不适合直接执行。max_memory_allocated 统计的是张量分配峰值，不等于整个进程占据的所有显存。

CPU 监控函数每 0.1 秒读取 RSS，用单独线程收集样本后减去基线。它可能漏掉更短的峰值；若被测函数抛错，还应在 finally 中停止监控线程；采样列表为空也要处理。不同实验在同一进程里的缓存和映射状态会影响基线，不能把一次数字当作普适结论。

逐权重文件方案降低单次源张量大小，却增加文件数量、目录开销和 I/O 管理。还需要清单、哈希、全部状态完整性以及中断恢复。选择方案时应先明确瓶颈是 CPU 内存、GPU 峰值还是读取时间。

15.10 练习与答案

上游续训练习加载模型和优化器再训练一轮；预训练模型评价练习比较不同 GPT-2 尺寸；生成练习继续使用载入的模型。机制见本章和第 13—14 章，真实大模型比较未运行。

补充题一：QKV 权重拆分后都是方阵，为什么 shape 检查还不够？答案：Q/K/V 顺序和转置错误都可能保持 shape，必须用数值与前向对照。

补充题二：只恢复权重、没有恢复 AdamW 动量，下一步一定相同吗？不能保证相同，通常会不同，因为优化器历史影响更新。补充题三：meta 后逐参数复制为何仍可能看见未来？答案：因果 mask 是 buffer，没被参数遍历恢复。

执行 python examples/core_training/checkpoint_checks.py qa/core_training/checkpoint_checks.json 可复查续训、meta/mmap/assign、buffer 反例和优化器旧引用。

参考与覆盖

• 主加载与生成脚本、下载器。
• 三种替代加载 notebook。
• 低内存 state_dict notebook：各加载路径与 GPU/CPU 监控均已解释，未运行大模型内存基准。