跳到主要内容

数据截至 (上游 commit b6c0bfe04c82)

02 · PreTrainedModel 与权重加载

这一章讲什么: AutoModel.from_pretrained("meta-llama/Llama-3.1-8B") 这行代码背后的完整流水线。重点是 v5 的新权重加载内核:checkpoint 键名与模型参数名不再要求一致,中间隔着一套声明式的 WeightTransform 转换规则。

1. 它要解决的小问题​

「把 Hub 上的权重装进模型」听起来是一次 load_state_dict,实际上要同时处理六件事:

  1. 找文件:repo id / 本地路径 / 显式文件名 / GGUF / 分片索引,还要缓存、锁版本。
  2. 建模型:按 config 实例化,且 70B 模型不能真分配内存——先在 meta 设备上建「空壳」。
  3. 改结构:量化要把某些 nn.Linear 换成量化层;张量并行要给参数加切分钩子。
  4. 对键名:checkpoint 是别的框架/老版本存的,键名和当前模型对不上(QKV 可能存成一块,专家可能存成 8 个独立矩阵)。
  5. 摆设备:按 device_map 把每层放到对的 GPU/CPU/磁盘上。
  6. 绑权重:lm_head 和 embed_tokens 常常共享同一张矩阵,加载完要重新绑起来。

PreTrainedModel.from_pretrained(modeling_utils.py:3881)就是把这六件事编成一条固定流水线。

2. 流水线全景​

from_pretrained 函数体约 550 行(3881-4434),按序做这些事(行号均指 modeling_utils.py):

步骤干什么位置
① 解析 config没传 config 就从 repo 读 config.json 实例化:4233-4256
② 解析 dtype_get_dtype 决定计算 dtype,处理 dtype="auto":4308
③ 实例化模型cls(config, ...),量化器先替换量化层:4349-4359
④ 取转换规则get_model_conversion_mapping 收集该模型的 WeightTransform 列表:4366
⑤ 分布式改造maybe_distribute_model 挂张量并行钩子:4369
⑥ 准备 device_map决定每层去哪个设备:4372 起
⑦ 加载权重_load_pretrained_model:读文件 + 转换 + 写入:4393
⑧ 后处理多卡/磁盘 dispatch 钩子、量化后处理、adapter 加载:4412-4425

第 ③ 步值得单独一句:实例化发生在 ContextManagers(model_init_context) 里(:4349-4351),这个大模型通常建在 meta 设备上——只有形状没有数据,所以 70B 模型也能秒建;真正的内存由第 ⑦ 步逐张量物化时占用。

3. 惰性读取:checkpoint 不是一次读进内存的​

_load_pretrained_model(modeling_utils.py:4436)读 safetensors 的方式是这条流水线里最工程化的一段:

# 摘自 modeling_utils.py:4516-4519(精简)
file_pointer = safe_open(file, framework="pt", device=device, backend=backend)
for k in file_pointer.keys():
    merged_state_dict[k] = file_pointer.get_slice(k)  # don't materialize yet

三个要点:

  • safe_open 用 mmap(Apple Silicon 上换 pread 后端)打开文件,此时不占内存。
  • get_slice(k) 拿到的只是「某键在文件里的切片描述」,merged_state_dict 里装的是惰性切片而不是张量——注释里明写 "don't materialize yet"。
  • 张量只在后面被 WeightTransform 真正需要、且目标设备已确定时才物化。这就是「8GB 内存的机器也能加载 70B 模型(device_map 分到磁盘)」的底座。

4. v5 权重转换系统:把格式漂移建成算子库​

4.1 思路​

旧代码里,键名转换是加载函数里一坨针对具体模型的 if-else。v5 把它建成声明式规则:模型(或注册表)声明「源键模式 → 目标键模式 + 一串算子」,加载引擎对 checkpoint 键做模式匹配,匹配上的交给算子链处理。

核心抽象在 core_model_loading.py:

抽象干什么位置
ConversionOps张量操作基类core_model_loading.py:81
Chunk / Concatenate / MergeModulelist / Transpose / PermuteForRope …具体算子:切、拼、合并 ModuleList、转置、RoPE 重排:112 / :150 / :222 / :312 / :428
WeightTransform一条转换规则:源模式 + 目标模式 + 算子:769
WeightRenaming只改名不变形的特化:992
WeightConverter改名 + 变形:1152
convert_and_load_state_dict_in_model总引擎:匹配 → 收集 → 算 → 写入:1511

4.2 两个官方例子(来自引擎 docstring)​

convert_and_load_state_dict_in_model 的 docstring(core_model_loading.py:1518-1596)自己就是最好的教材,两个典型场景:

场景 A:checkpoint 里 QKV 存成一块,模型里是分开的三个投影。

# 示意,非源码(规则形态与 docstring 一致)
WeightConverter(
    source_patterns=["qkv"],
    target_patterns=["q", "k", "v"],
    operations=[Chunk(dim=0, chunks=3)],   # 沿 0 维切成 3 块,分别改名 q/k/v
)

场景 B:MoE checkpoint 里 8 个专家是 8 个独立矩阵,模型里是一块融合权重。

# 示意,非源码
WeightConverter(
    source_patterns=["mlp.experts.*.gate_proj.weight", "mlp.experts.*.up_proj.weight"],
    target_patterns="mlp.experts.gate_up_proj",
    operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
    # 先把每个模式下的 8 个张量 stack 成一块,再沿 1 维把 gate/up 拼起来
)

引擎内部按目标参数名建 param_name_to_load 映射,每个条目带着 collected_tensors(收集到的源张量,惰性)和 layer_targets(哪些目标键共享同一源);收集齐 → 执行算子链 → 量化/反量化(若需要)→ 写入模型。

4.3 规则从哪来:中央注册表​

规则有两个来源:

  1. 库内置:中央注册表 conversion_mapping.py 的 _build_checkpoint_conversion_mapping()(:144)返回一个巨型 dict,按 model_type 存一批 WeightRenaming/WeightConverter。查的时候类名优先、model_type 兜底(extract_weight_conversions_for_model,:1882)——这样同一架构的不同任务头可以有自己的转换。
  2. 用户注册:register_checkpoint_conversion_mapping(model_type_or_class_name, mapping)(:1857),或 from_pretrained(key_mapping={...}) 直接传一次性映射(get_model_conversion_mapping 里可见,conversion_mapping.py:1924-1925)。

get_model_conversion_mapping(:1902)还负责给子模块的规则打 scope_prefix,保证父子模型规则不互相串——docstring 里专门解释了「兄弟子模型同 model_type 时各自拿自己的规则」的去重逻辑(:1909-1917)。

5. 绑权重:同一张矩阵的两个名字​

很多 decoder 模型的 lm_head.weight 与 model.embed_tokens.weight 是同一张矩阵(输入输出共用词表)。checkpoint 里通常只存一份,加载后要把另一个名字指过去。

tie_weights(modeling_utils.py:2606)做的就是这件事,v5 里它维护一个 {target: source} 映射(all_tied_weights_keys)。有一个不显然的设计值得记住,写在它的 docstring 里(:2614-2617):

在 from_pretrained 期间,绑定是对称的:如果映射说「target 绑到 source」,但 checkpoint 里 source 缺、target 在,就交换方向,绑到实际存在的那个参数上。

这个对称性让「checkpoint 只存了 lm_head 没存 embed」和「只存了 embed 没存 lm_head」两种老格式都能正确加载,而不是只容忍一种。

6. 保存是加载的逆运算​

save_pretrained(modeling_utils.py:3300)在本 commit 的形态:

  1. 先还原转换:revert_weight_conversion 把内存里的权重变回「原始格式」再存(:3518)——加载时做的 WeightTransform 在保存时反向执行,保证存出去的文件仍能被老版本/别的框架读。
  2. 去掉 tied 权重:remove_tied_weights_from_state_dict(:3512),因为 safetensors 不支持共享存储,同一张矩阵只存一份。
  3. 永远存 safetensors:文件名直接取 SAFE_WEIGHTS_NAME(:3522-3524),不再提供 .bin 分支(PEFT adapter 除外)。
  4. 超阈值自动分片:split_torch_state_dict_into_shards(state_dict, max_shard_size=max_shard_size)(:3530-3532),默认 50GB 一片,产出 model-00001-of-00005.safetensors + 一份 index.json 记录每个键在哪片。

7. 关键细节与坑​

  • missing/unexpected keys 警告要会读。 "Weights from XXX not initialized" = 模型有但 checkpoint 没有(通常是新加的头,需要微调);"Weights not used" = checkpoint 有但模型不要(from_pretrained docstring,:3903-3908)。v5 里这两条经过转换引擎后语义不变,但键名已经是转换后的。
  • device_map 里有 "disk" 会走另一条路。 _load_pretrained_model 开头先建磁盘 offload 索引(modeling_utils.py:4459-4467),后续参数直接落到磁盘文件而不是内存。
  • 多卡/磁盘混合时加载完还要装 dispatch 钩子(:4412-4414)——模型 forward 时输入会被自动搬到该层所在设备,这是 device_map="auto" 用起来「无感」的原因,也是性能陷阱的来源(跨设备搬运不免费)。
  • DeepSpeed ZeRO-3 是特例中的特例(:4477-4504):参数本来就是分片的,走 _load_state_dict_into_zero3_model,且假定 checkpoint 完全匹配,连 unexpected_keys 都不报(注释自承 "we assume only best-case scenario",:4491)。
  • 转换规则写错不会立刻炸。 模式没匹配上的 checkpoint 键会进 unexpected_keys 警告而不是报错——调试自定义 key_mapping 时,先开 output_loading_info=True 看 info,别靠猜。

8. 代码地图​

主题文件路径符号名
加载总入口src/transformers/modeling_utils.pyPreTrainedModel.from_pretrained
加载内核(读文件+分发)src/transformers/modeling_utils.pyPreTrainedModel._load_pretrained_model
转换引擎src/transformers/core_model_loading.pyconvert_and_load_state_dict_in_model
转换算子src/transformers/core_model_loading.pyConversionOps、Chunk、Concatenate、MergeModulelist、Transpose、PermuteForRope
转换规则src/transformers/core_model_loading.pyWeightTransform、WeightRenaming、WeightConverter
中央转换注册表src/transformers/conversion_mapping.py_build_checkpoint_conversion_mapping、register_checkpoint_conversion_mapping、extract_weight_conversions_for_model、get_model_conversion_mapping
权重绑定src/transformers/modeling_utils.pyPreTrainedModel.tie_weights
保存与分片src/transformers/modeling_utils.pyPreTrainedModel.save_pretrained(split_torch_state_dict_into_shards 来自 huggingface_hub,modeling_utils.py:36 import)
逆转换(保存时)src/transformers/core_model_loading.pyrevert_weight_conversion