知识花园

大模型检查点

7分钟阅读

大模型训练把权重、梯度、优化器状态切散在多个 rank 上，检查点随之从单文件变成一整套分片目录：Megatron-LM、DeepSpeed ZeRO、PyTorch FSDP 各有一套格式。读懂目录结构、保住优化器分片、对齐并行度，是断点续训与格式导出的三条命脉。

检查点版图三分：模型权重保存 讲 PyTorch 原生 torch.save 机制，推理模型格式 讲 HF / safetensors / GGUF 等推理态格式，本篇专讲训练态分片检查点。TP / PP / DP / ZeRO 等并行策略本身的原理见 分布式训练总览，这里只关心它们如何决定落盘形态。三大框架的共性：按 rank 分片落盘、附带 optimizer 与 RNG 状态、与并行拓扑强绑定——因此不能 torch.load 直读，续训受并行度约束。

Megatron-LM：TP/PP/DP 三级分片目录

Megatron-LM（含 Megatron-DeepSpeed、NeMo 等衍生分支）按 TP → PP → DP 组织目录：每个 TP×PP 组合一个目录；DP 组内权重相同，不占目录层级（分布式优化器状态才按 DP 维再切片落盘）。

iter_0068000/                          # 按全局迭代步命名
├── latest_checkpointed_iteration.txt  # 最新迭代标记，加载入口
├── mp_rank_00/                        # TP rank 0
│   ├── pp_rank_00/                    # PP stage 0
│   │   ├── model_optim_rng.pt         # 权重分片 + 优化器超参 + RNG 状态
│   │   └── distrib_optim.pt           # 分布式优化器分片（Adam 的 m/v）
│   └── pp_rank_01/
└── mp_rank_01/                        # TP rank 1
  • model_optim_rng.pt 三合一：权重切片、optimizer 超参、RNG 状态——续训时 dropout 与数据 shuffle 精确对齐靠它
  • distrib_optim.pt 仅在启用 Distributed Optimizer 时存在，Adam 动量按 DP 维分片；丢失后续训等效换优化器，loss 必跳变
  • 目录缺 pp_rank 层即 PP=1；mp_rank 个数即 TP 度
  • 新版 Megatron-LM 另支持基于 PyTorch DCP 的统一分片格式（.distcp + .metadata），可跨并行度 resharding，但旧式目录仍是存量主流

加载约束：MP/PP/DP 必须与保存时严格一致；任一分片缺失即失败；单分片只是切片，必须由 Megatron 按拓扑重建模型后填充；各分支目录细节有差异，不能互相混用。

DeepSpeed ZeRO：stage 决定分什么

Stage训练时分片对象检查点形态
ZeRO-1optimizer states每 rank 文件内权重为完整冗余副本
ZeRO-2+ gradients同上（梯度不落盘）
ZeRO-3+ parameters权重真分片，需全部 rank 文件才能重建全量
checkpoint/
├── latest                     # 文本文件，指向最新 step 目录
└── global_step_1000/
    ├── zero_pp_rank_0_model_states.pt   # 各 rank 的权重（形态随 stage）
    ├── zero_pp_rank_1_model_states.pt
    ├── zero_pp_rank_0_optim_states.pt   # 各 rank 的优化器状态分片
    └── zero_pp_rank_1_optim_states.pt

命名中的 pp 指 ZeRO 参数分区序号（沿 DP 维），与 Megatron 的流水线 pp_rank 无关，勿混淆。

标准 ZeRO 检查点与 world size、stage 强耦合，换卡数、换 stage 都无法直接续训。Universal Checkpoint 解决这一点：用 ds_to_universal.py 把标准检查点改写为统一命名、按参数连续存储的布局，加载侧设 load_universal: true，即可跨 ZeRO 配置、跨 world size 续训——训练中扩缩容、多人接力训练的标准做法。

PyTorch FSDP：三种保存模式与 DCP

StateDictType落盘形态通信开销用途
FULL_STATE_DICTrank 0 gather 合并全量单文件高导出推理、跨框架迁移
LOCAL_STATE_DICT每 rank 写自身分片零同拓扑断点续训
SHARDED_STATE_DICTDCP 目录：.distcp 分片 + .metadata低超大模型训练主力

DCP（torch.distributed.checkpoint）是 PyTorch 官方分布式检查点后端：各 rank 并行写自身分片，元数据记录每个张量在各文件中的字节区间，加载时按新拓扑重新切分——保存与加载的 world size 可以不同（resharding），这是它取代 LOCAL 模式的本质优势，也是 torchtitan 等官方参考实现的标准选择。协调依赖 torch.distributed 通信原语 的进程组；加载必须走 torch.distributed.checkpoint.load，不支持 torch.load。

from torch.distributed.checkpoint import save
from torch.distributed.checkpoint.state_dict import (
    get_model_state_dict, get_optimizer_state_dict)
 
# FSDP2 + DCP：各 rank 并行写分片，无全局 gather，支持异步保存
state = {"model": get_model_state_dict(model),
         "optim": get_optimizer_state_dict(optimizer), "step": step}
save(state, checkpoint_id="ckpt/step_1000")  # 生成 __0_0.distcp... + .metadata

并行度匹配与续训铁律

  • 续训并行度与保存时严格一致：Megatron 旧式目录 MP/PP/DP 全对齐，ZeRO 对齐 stage 与 world size，FSDP LOCAL 对齐分片数。错配表现为维度不匹配或权重缺失，是分布式加载失败的头号原因
  • 例外仅限显式支持 resharding 的格式：DCP 系（FSDP SHARDED、新版 Megatron）与 Universal Checkpoint，且需按文档声明加载配置
  • 分片完整性：手删”看似冗余”的 rank 文件（尤其 optim_states、distrib_optim.pt）会破坏优化器状态；只要权重推理，应走官方转换抽取，而非手动删文件
  • RNG 与调度器状态决定”无缝”续训：只用权重热启动是另一种训练策略，不是续训

格式转换

跨格式转换一律官方工具，不手拼权重：

转换工具
Megatron → HFMegatron-LM tools/checkpoint/convert.py（旧称 checkpoint_util / convert_checkpoint 一族）；Megatron-DeepSpeed 用 megatron_to_transformers.py
ZeRO 分片 → 全量单文件检查点目录内自动生成的 zero_to_fp32.py
ZeRO → Universalds_to_universal.py
FSDP → HFFULL_STATE_DICT 合并后 save_pretrained
# 合并 ZeRO-3 分片为全量 fp32 权重（脚本随检查点自动生成）
cd checkpoint/global_step_1000
python zero_to_fp32.py . pytorch_model.bin

转换的通用路径是「分片目录 → 全量单文件 → 推理态格式」。用 Accelerate 训练时，底层无论 FSDP 还是 DeepSpeed，save_model 已封装 gather 逻辑，通常无需手动转换。

关系图谱

反向链接