3、分布式Checkpoint框架：基于PyTorch Distributed的State Dict保存、Megatron-LM与DeepSpeed的Checkpoint机制

分布式训练里，Checkpoint 从来不是「把模型参数存个盘」那么简单。

你想想看，当你的模型分布在几十甚至几百张 GPU 上，每张卡只持有模型的一部分——要么是数据并行的一份副本，要么是模型并行中的一个切片，要么是流水线并行中的一个 stage。这时候，怎么保证存下来的东西能完整恢复？

我最早踩这个坑是在一次 64 卡的多机训练里。训练跑了三天，突然节点掉线，我心想没事，有 Checkpoint。结果恢复的时候发现，因为数据并行的 rank 映射没对齐，加载的参数全乱了。嗯，从那以后，我对分布式 Checkpoint 的敬畏心就上来了。

3.1 PyTorch Distributed 的 State Dict 保存：最基础的方案

PyTorch 原生的 torch.save(model.state_dict(), path) 在单卡场景下没问题。但到了分布式环境，你得自己处理 rank 的协调。

最朴素的做法是这样：

import torch
import torch.distributed as dist

def save_checkpoint(model, optimizer, epoch, path):
    if dist.get_rank() == 0:
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
        }, path)
    dist.barrier()

这段代码的问题很明显——它假设所有 rank 的 state dict 是一样的。这在纯 DDP（Distributed Data Parallel）下成立，因为 DDP 会同步梯度，模型参数在各 rank 间保持一致。

但一旦你用了模型并行或流水线并行，每个 rank 的 state dict 就不同了。这时候，你需要每个 rank 各自保存自己的那份：

def save_distributed_checkpoint(model, optimizer, epoch, path_prefix):
    rank = dist.get_rank()
    local_path = f"{path_prefix}_rank_{rank}.pt"
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'rank': rank,
    }, local_path)
    dist.barrier()
注意： 每个 rank 保存独立文件，恢复时也必须按 rank 加载。我曾经见过有人把所有 rank 的文件合并成一个，结果加载时 rank 映射全乱套了。千万别这么干。

恢复的逻辑也简单：

def load_distributed_checkpoint(model, optimizer, path_prefix):
    rank = dist.get_rank()
    local_path = f"{path_prefix}_rank_{rank}.pt"
    checkpoint = torch.load(local_path, map_location=f'cuda:{rank}')
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    return checkpoint['epoch']

这个方案的好处是简单直接，坏处是文件数量随 GPU 数量线性增长。你训练 1024 张卡，就得管理 1024 个文件。说实话，这在实际运维中挺头疼的。

3.2 Megatron-LM 的 Checkpoint 机制：模型并行的救星

Megatron-LM 是我个人非常推崇的一个框架。它在 Checkpoint 设计上做了很多工程化的考量。

Megatron-LM 的核心思路是：每个模型并行组内的 rank 共享一个 Checkpoint 文件。它通过 mp_rank_XX 的命名规则来区分不同模型并行分片的文件。

举个例子，假设你有 4 张卡，模型并行度（tensor parallelism, TP）为 2，数据并行度（DP）为 2。那么：

  • TP group 0（rank 0, 1）保存为 mp_rank_00.pt
  • TP group 1（rank 2, 3）保存为 mp_rank_01.pt

每个文件里存的是该 TP group 内所有 rank 的 state dict。恢复时，每个 TP group 加载自己的文件，然后按 rank 切分。

Megatron-LM 的 Checkpoint 代码里有一个关键函数 save_checkpoint，它的流程大致是：

  1. 只在每个 TP group 的 master rank（通常是 rank 0）上执行保存操作。
  2. 收集该 group 内所有 rank 的 state dict。
  3. 打包成一个字典，写入文件。
  4. 其他 rank 等待 barrier。

关键设计： Megatron-LM 把 optimizer state 也按模型并行分片存储。这意味着恢复时 optimizer 的 momentum、variance 等状态也能精确还原。这一点对长训练任务至关重要。

我自己的经验是，Megatron-LM 的 Checkpoint 在千卡规模下表现稳定。但有一个坑：如果你改了模型并行度（比如从 TP=2 改成 TP=4），旧的 Checkpoint 就没法直接用了。你得写一个转换脚本，把旧的分片重新切分。这个转换过程我做过一次，说实话，挺折腾的。

3.3 DeepSpeed 的 Checkpoint 机制：ZeRO 优化的最佳拍档

DeepSpeed 的 Checkpoint 设计跟它的 ZeRO 优化器深度绑定。ZeRO 把 optimizer state、gradient 和 parameter 分散到各个 rank 上，所以 Checkpoint 也必须感知这种分布。

DeepSpeed 提供了两个层次的 Checkpoint：

  • Engine Checkpoint：保存整个训练引擎的状态，包括模型、优化器、学习率调度器等。
  • Client Checkpoint：只保存用户自定义的状态，比如数据加载器的进度。

Engine Checkpoint 的保存接口很简单：

model_engine.save_checkpoint(save_dir, tag="latest")

但背后的逻辑相当复杂。DeepSpeed 会根据 ZeRO 的 stage 来决定如何保存：

ZeRO Stage 保存策略 文件数量
Stage 1 optimizer state 分片，每 rank 保存自己的部分 等于 rank 数
Stage 2 optimizer state + gradient 分片 等于 rank 数
Stage 3 optimizer state + gradient + parameter 分片 等于 rank 数

恢复时，DeepSpeed 会自动从所有 rank 的文件中重建完整的模型状态。这个过程对用户是透明的。

个人建议： 使用 DeepSpeed 时，一定要设置 save_latest 参数为 True。这样每次保存都会更新一个符号链接指向最新的 Checkpoint。训练意外中断后，你可以直接用 load_checkpoint 恢复，不用手动指定 tag。这个小习惯帮我省了不少时间。

DeepSpeed 还有一个很实用的功能：异步保存。它会把 Checkpoint 的写入操作放到后台线程执行，不阻塞训练。这在 IO 成为瓶颈的大规模训练中特别有用。

不过要注意，异步保存时如果训练进程崩溃，正在写的 Checkpoint 可能会损坏。我一般会在关键节点（比如每个 epoch 结束）用同步保存，中间步骤用异步保存。

3.4 三种方案的对比与选型

说了这么多，我们来个直观的对比：

特性 PyTorch 原生 Megatron-LM DeepSpeed
适用并行策略 纯 DDP TP + PP + DP ZeRO + DP
文件数量 等于 rank 数 等于 TP group 数 等于 rank 数
恢复速度 快（每卡独立加载） 中等（需 group 内聚合） 慢（需全量重建）
存储开销 高（每卡存完整副本） 中（分片存储） 低（ZeRO 分片）
易用性 简单 中等 简单（接口封装好）

怎么选？我的经验是：

  • 如果你只是做数据并行，用 PyTorch 原生就够了，别引入额外依赖。
  • 如果你用 Megatron-LM 做模型并行，就用它自带的 Checkpoint 机制，别自己造轮子。
  • 如果你用 DeepSpeed 的 ZeRO，那必须用 DeepSpeed 的 Checkpoint，因为只有它知道怎么正确保存和恢复分片状态。

3.5 一张图看懂分布式 Checkpoint 框架

下面这张图展示了三种框架在分布式 Checkpoint 中的核心差异：

分布式 Checkpoint 框架对比 PyTorch 原生 每 rank 独立保存 state_dict 文件数 = rank 数 仅支持 DDP Megatron-LM TP group 内聚合保存 文件数 = TP group 数 支持 TP + PP + DP DeepSpeed ZeRO 分片感知保存 文件数 = rank 数 支持 ZeRO 1/2/3 选型建议 • 纯 DDP 训练 → PyTorch 原生，简单够用 • 模型并行（TP/PP）→ Megatron-LM，分片管理成熟 • ZeRO 优化 → DeepSpeed，与优化器深度集成 • 混合并行 → Megatron-DeepSpeed 组合，各取所长 • 恢复速度优先 → PyTorch 原生（每卡独立加载最快） • 存储空间优先 → DeepSpeed ZeRO-3（分片最彻底）

3.6 实战中的几个坑

最后，分享几个我实际踩过的坑：

  • 文件系统压力：千卡规模下，所有 rank 同时写文件，文件系统可能扛不住。我建议用 dist.barrier() 错峰写入，或者用 DeepSpeed 的异步保存。
  • Checkpoint 校验：保存后一定要做一次加载验证。我有一次因为磁盘空间不足，文件写了一半，训练中断后恢复才发现 Checkpoint 损坏。后来我加了一个校验步骤：保存后立即加载，对比 loss 是否一致。
  • 版本兼容：PyTorch、Megatron-LM、DeepSpeed 的版本升级可能导致 Checkpoint 不兼容。我习惯在 Checkpoint 文件名里带上框架版本号，比如 checkpoint_v1.2.3.pt。
  • 多节点训练：节点间网络抖动可能导致 Checkpoint 保存超时。我一般会设置 torch.distributed.distributed_c10d::_DEFAULT_TIMEOUT 为 30 分钟，给足缓冲时间。

分布式 Checkpoint 没有银弹。你得根据你的并行策略、硬件环境和恢复需求来选。但不管选哪个，一定要在训练开始前就把 Checkpoint 的保存和恢复流程跑通。别等到训练跑了一周才发现恢复不了——那种感觉，我不想再体验第二次了。

← 上一章 📖 返回目录 下一章 →