加载预训练权重

加载预训练权重

本节摘要:从零训练一个 1.24 亿参数的模型是预算决策;加载一个已发布的检查点只是一个普通的周二。本节把预训练 GPT-2 风格权重从 safetensors 文件灌进第 33 节的精确架构,逐张量走参数名映射,并生成续写证明加载成功。发布的权重不是为你的架构打包的——它们带着原始实现的命名,h.0.attn.c_attn.weight 形如 (2304, 768),而你的模型期望 blocks.0.attn.qkv.weight 形如 (768, 2304)(同一个矩阵、不同布局约定),有的甚至整体转置。本节的加载器显式处理这三重身份差异(名字、形状、字节布局),每个赋值都记日志、每个形状都校验,LoadReport 汇总命中、缺失、不匹配,让你读清发生了什么。

对应原课程:Phase 19 · Lesson 37 · loading-pretrained-weights(原英文 phases/19-capstone-projects/37-loading-pretrained-weights/docs/en.md)。本节属「从零构建 GPT」赛道第八节。

学习目标

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

1. 用 safetensors Python 库读取文件,检查张量名与形状。
2. 把每个预训练参数名映射到第 33 节 GPT 模型内部的参数上。
3. 处理发布权重与本赛道模型之间的两种命名约定: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。
4. 在赋值前用清晰报错检测并拒绝形状不匹配。
5. 用加载后的权重生成短续写,确认 token 来自加载的分布而非随机初始化。

一、问题与直觉

发布的权重不是为你的架构打包的。它们带着原始实现用的名字。预训练文件里有 transformer.h.0.attn.c_attn.weight 形状 (2304, 768);你的模型期望 blocks.0.attn.qkv.weight 形状 (2304, 768)(不同布局约定下的同一个矩阵),或者你的模型用 nn.Linear 存的是转置后的矩阵。同一个参数以三种微妙不同的身份出现(名字、形状、字节布局),加载器必须调和三者。

名字映射器只是 string 到 string 的函数。形状检查是一个 if。赋值发生在 torch.no_grad() 内,autograd 不追踪加载。报告记录每个名字的结果。

GPT-2 命名约定

发布的 GPT-2 权重用这样的名字:

两个意外要提前规划。c_attn、c_proj、c_fc 这些线性层存的矩阵相对 nn.Linear.weight 期望的是转置的,加载器在赋值时转置。LM 头根本不在文件里——模型依赖与 wte 的权重绑定,头在 wte 落地后通过别名设置。

本地命名约定

本赛道模型用描述性名字:tok_embed.weight、pos_embed.weight、blocks.N.ln1.scale、blocks.N.attn.qkv.weight、blocks.N.attn.out_proj.weight、blocks.N.mlp.fc1.weight、final_ln.scale 等。映射是固定函数,本节把它做成字典,加载器逐层展开迭代。

桩夹具

真正的 GPT-2 权重 0.5 GB。demo 不下载,而是在首次运行生成一个小 safetensors 夹具,用精确的 GPT-2 命名约定与适合 d_model 192(而非 768)的 12 块模型形状。夹具有正确结构,能跑通加载器每条代码路径。把夹具换成真文件,加载器无需修改即可工作。

二、从零实现

code/main.py 实现:

• 第 33 节 GPTModel 的小副本,使本节自包含。
• make_pretrained_to_local(num_layers):展开逐层条目的名字映射。
• load_safetensors(model, path):迭代名字、映射、检查形状、转置 conv1d 式权重、在 torch.no_grad() 下赋值,返回 LoadReport。
• make_stub_safetensors(path, cfg):用精确预训练命名约定生成夹具文件。
• demo:首次运行创建 outputs/gpt2-stub.safetensors,建新模型、从随机初始化抓一段续写、加载桩、再抓一段续写,打印两者并验证不同(加载真的改了模型)。

映射与赋值骨架:

def make_pretrained_to_local(n):
    m = {"wte.weight": "tok_embed.weight",
         "wpe.weight": "pos_embed.weight",
         "ln_f.weight": "final_ln.scale", "ln_f.bias": "final_ln.shift"}
    for i in range(n):
        p = f"h.{i}"; loc = f"blocks.{i}"
        m[f"{p}.ln_1.weight"] = f"{loc}.ln1.scale"
        m[f"{p}.attn.c_attn.weight"] = f"{loc}.attn.qkv.weight"   # 需转置
        m[f"{p}.attn.c_proj.weight"] = f"{loc}.attn.out_proj.weight"
        m[f"{p}.mlp.c_fc.weight"]   = f"{loc}.mlp.fc1.weight"
        m[f"{p}.mlp.c_proj.weight"] = f"{loc}.mlp.fc2.weight"
    return m

def load_safetensors(model, path):
    report = LoadReport()
    name_map = make_pretrained_to_local(model.num_layers)
    TRANSPOSE = {"qkv", "out_proj", "fc1", "fc2"}                  # conv1d 式布局
    with safe_open(path, framework="pt") as f:
        for src in f.keys():
            if src not in name_map:
                report.unexpected.append(src); continue
            dst = name_map[src]
            tensor = f.get_tensor(src)
            if any(k in dst for k in TRANSPOSE): tensor = tensor.t()   # 转置
            param = get_param(model, dst)
            if param.shape != tensor.shape:
                report.shape_mismatch.append((src, param.shape, tensor.shape))
                continue
            with torch.no_grad(): param.copy_(tensor)
            report.loaded.append((src, dst))
    # weight tying:头不在文件里,别名到嵌入
    model.lm_head.weight = model.tok_embed.weight
    return report
​

设计要点:权重绑定是别名而非拷贝——model.lm_head.weight = model.tok_embed.weight 让头与嵌入共享存储;拷贝嵌入矩阵到新的 lm_head.weight 参数会破坏绑定、悄悄翻倍参数量。

三、框架对比

HuggingFace transformers 的 from_pretrained 把这一切打包成一行:GPT2LMHeadModel.from_pretrained("gpt2") 自动下载、映射、绑权重,但黑盒后你不知道哪步出错。本节的加载器全部显式:每个赋值记日志、每个形状校验、LoadReport 汇总。同样的模式换到 LLaMA、Mistral、Qwen,只需更新名字映射(它们无偏置、用 RMSNorm、融合 qkv 布局不同),形状检查与报告逻辑原样不动。safetensors 格式比 pickle 更安全(无任意代码执行)、加载更快(零拷贝、mmap),是 2023 年后的事实标准。

四、可复用产物

code/main.py 与 outputs/gpt2-stub.safetensors:加载器对任何用预训练命名约定的 safetensors 文件工作。真 GPT-2 文件(small/medium/large/xl)无需改代码,只有模型配置不同。健全性生成是快速闸门——若加载后样本长得像加载前,说明映射悄悄漏了每个张量。

五、生产模式

三模式让加载器扛住非你创建的权重。赋值前先全量校验文件:打开文件、列出每个张量名与 dtype/形状、跑完整映射与形状检查,只在全成功后才开始赋值——半加载的模型是静默失败的源头。每个赋值记源名与目的名日志:出问题时日志告诉你哪个张量落到哪;替代方案是读 hexdump。LoadReport 追踪 loaded/missing/unexpected/shape_mismatch 四个列表,末尾打汇总。LM 头是权重绑定别名而非独立拷贝:加载 tok_embed 后设 model.lm_head.weight = model.tok_embed.weight 是规范模式。

六、练习

1. dtype 转换:给加载器加 dtype 参数,赋值时把每个张量转成目标 dtype(bfloat16/float16/float32),确认 float32 模型能下转 bfloat16 后仍生成。
2. 层数校验:加 expected_layers 参数,拒绝层数索引与模型 num_layers 不符的检查点。
3. 并排样本:把加载器接进第 33 节生成函数,产出随机初始化与加载夹具两段并排样本。
4. 导出路径:写当前模型状态到新 safetensors 文件(用预训练命名约定),往返加载器,确认报告零形状不匹配。
5. LLaMA 映射:扩展 NAME_MAP 处理 LLaMA 命名(无偏置、RMSNorm、融合 qkv 布局),在你生成的桩 LLaMA 夹具上重跑加载器。

本节要点回顾

1. 三重身份:同一参数在名字、形状、字节布局三处微妙不同,加载器须调和三者。
2. 名字映射是字符串函数:预训练名到本地名的固定字典,逐层展开。
3. 转置在加载时:GPT-2 的注意力/MLP 投影相对 nn.Linear 是转置布局,赋值时 .t()。
4. 赋值前校验形状:不匹配就拒绝并记日志,绝不让加载器盲拷。
5. LM 头靠权重绑定:头不在文件里,lm_head.weight = tok_embed.weight 别名是规范。
6. 全量校验再赋值:半加载模型是静默失败机器,先校验全部再开始赋值。
7. 健全性生成即闸门:加载后样本像加载前,说明映射悄悄漏光了。

下一节,我们做「分类微调」——拆掉 LM 头、换上两类分类头,在「冻体/全调」两种策略下训练垃圾短信分类器,用精确率/召回率/F1 评估。