📚 课程目录第3章模型加载流程

3、模型加载流程：从检查点加载模型的步骤、关键函数、配置参数

当训练任务中断或需要从已有权重继续训练时，DeepSpeed 的模型加载并不是简单地调用一次 torch.load 就结束。它需要经历检查点识别、分片索引解析、并行状态重建、参数映射与填充、优化器状态恢复等多个阶段。理解这条链路，是定位“加载失败”问题的前提。

核心结论：DeepSpeed 加载失败，绝大多数不是权重文件损坏，而是加载时的并行配置与保存时的并行配置不一致，导致分片索引对不上、参数形状不匹配或优化器状态无法映射。

一、加载流程的整体链路

DeepSpeed 的加载入口通常由 deepspeed.initialize() 触发，内部根据配置决定是否调用 load_checkpoint。整体可分为五个阶段：

检查点目录识别：扫描 latest 文件或指定 tag，确定要加载的检查点路径。
分片索引解析：读取 mp_rank_XX_model_states.pt 与 zero_pp_rank_XX_mp_rank_YY_optim_states.pt 等分片文件，建立参数名到分片的映射。
并行状态重建：根据当前 mp_size、pp_size、zero_stage 重建张量并行、流水线并行与 ZeRO 分片视图。
参数映射与填充：将检查点中的分片参数按当前并行布局重新切分或聚合，写入模型。
优化器与调度器恢复：恢复优化器动量、学习率调度器步数、全局步数等训练状态。

下面这张图展示了从检查点目录到模型与优化器状态恢复的关键数据流与函数调用关系：

二、关键函数与调用关系

DeepSpeed 的加载逻辑集中在 deepspeed/runtime/engine.py 与 deepspeed/runtime/zero/stage3.py 中。以下是最常被触发的关键函数：

函数	所在模块	作用
`load_checkpoint()`	DeepSpeedEngine	加载入口，解析 tag 与路径，协调模型与优化器加载
`_load_checkpoint()`	DeepSpeedEngine	实际执行分片读取与状态恢复
`load_model_state_dict()`	DeepSpeedEngine	将模型分片参数映射到当前模型
`load_optimizer_state_dict()`	DeepSpeedEngine	恢复优化器动量、方差与 ZeRO 分片状态
`_get_ckpt_path()`	DeepSpeedEngine	根据 mp/pp rank 拼接分片文件路径
`load_state_dict()`	ZeRO Stage3	按 ZeRO 分片规则加载参数与优化器状态

三、配置参数对加载的影响

加载能否成功，取决于保存与加载时以下配置是否一致。任何一项不匹配，都可能触发 KeyError、RuntimeError: size mismatch 或 AssertionError。

{
  "train_batch_size": 32,
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu" },
    "offload_param": { "device": "cpu" }
  },
  "fp16": { "enabled": true },
  "tensor_parallel": { "tp_size": 2 },
  "pipeline": { "pp_size": 2 }
}

zero_optimization.stage：Stage 1/2/3 的分片粒度不同，跨 stage 加载通常失败。
offload_optimizer / offload_param：CPU offload 与 GPU 驻留的参数布局不同，需保持一致。
fp16 / bf16：混合精度状态中保存了 master weights，精度模式不一致会导致 dtype 冲突。
tensor_parallel.tp_size 与 pipeline.pp_size：决定分片文件命名与切分方式，必须与保存时一致。
train_batch_size：影响梯度累积与优化器状态形状，不一致时可能触发断言失败。

四、典型场景：从 ZeRO-2 检查点恢复到 ZeRO-3

假设保存时使用 ZeRO-2，加载时配置为 ZeRO-3，日志中会出现类似：

AssertionError: Checkpoint was saved with zero stage 2,
but current zero stage is 3. Cannot load optimizer state.

此时有两种处理方式：

对齐配置：将加载配置改回 stage: 2，先恢复训练，再通过转换脚本迁移到 ZeRO-3。
使用转换工具：通过 zero_to_fp32.py 将分片检查点合并为完整权重，再以新配置重新初始化。

如果只是模型权重加载失败，而优化器状态可以丢弃，可设置 load_optimizer_states: false，仅恢复模型参数，从而绕过优化器分片不匹配的问题。

五、常见误区与复盘清单

误区一：认为只要 latest 文件存在就能加载。实际上 latest 只记录 tag，分片文件缺失仍会失败。
误区二：忽略 mp_rank 与 pp_rank 的对应关系，手动移动文件后路径拼接错误。
误区三：在加载时修改了 train_batch_size，导致优化器状态形状不匹配。

复盘时建议按以下清单逐项核对：

检查点目录下是否存在完整的 mp_rank_* 与 zero_pp_rank_* 分片文件。
保存与加载的 zero_optimization.stage 是否一致。
tp_size、pp_size 是否与保存时相同。
是否启用了 CPU offload，且加载时配置一致。
模型结构（层数、隐藏维度）是否与检查点匹配。
日志中是否出现 size mismatch、KeyError 或 AssertionError，并定位到具体参数名。

掌握这条加载链路后，遇到加载失败时就能快速判断问题出在路径识别、分片索引、并行配置还是参数映射，而不是盲目重试或重新训练。