跳到主要内容

分布式 Checkpoint 通信

跨节点 IO 带宽、Sharding 分片策略与传输路径设计

核心要点：

  • Sharded Checkpoint：每 rank 只写自己分片，所有 rank 并行 IO; Resharding 支持 N→M rank 恢复
  • Save / Load 协议：Coordinator rank 汇总元数据 + 各 rank 并行写；ShardedTensor 携带 global_shape / local_shape / offset
  • 传输路径：传统 GPU→CPU→Storage 走 PCIe + CPU; GPUDirect Storage 绕开 CPU 直接 NVMe ⇄ HBM
  • Llama 3 实测：Tectonic 聚合带宽 2 TB/s 但单 job ~5 GB/s, 405B checkpoint 4.05 TB 写入 ~13.5 分钟
  • Universal Checkpointing: TP × PP × DP × ZeRO 切分维度可变情况下做跨配置恢复
  • Check-N-Run 增量：DLRM embedding 大但稀疏更新，只写 delta 大幅压缩 checkpoint 体积[1]

本文聚焦大规模 LLM 训练中 checkpoint 的跨节点通信代价 — 数据布局、Save/Load 协议、IO 带宽与传输路径。不写 async/overlap 时序细节 (见 10.3 异步 Checkpoint), 不重复集合通信算法本身 （只写在 checkpoint 场景中的用法）。

名词定义​

核心问题：分布式 Checkpoint 涉及的核心名词（Sharding/GDS/NVMe-oF/RDMA/数据布局）的定义是什么？

名词定义
Sharded Checkpoint按 rank 切片保存的 checkpoint：每个 rank 只写自己持有的参数/优化器分片，所有 rank 并行 IO
ReshardingLoad 时把 N 个 save 端 shard 重新切分到 M 个 load 端 rank 的过程，N≠M 时仍可恢复
Stateful 协议DCP / Megatron 用于约定 "什么对象的什么字段进入 checkpoint" 的接口，对象需实现 state_dict() / load_state_dict()
ShardedTensorMegatron-Core 中带 sharding metadata 的张量包装，记录 key / global_shape / local_shape / global_offset / axis_fragmentations / replica_id
Coordinator rankDCP / Megatron 中执行全局元数据汇总（plan 计算、metadata 写入）的特殊 rank，通常是 rank 0
Fully Parallel Save / LoadMegatron 的优化策略：在写入前通过 metadata-only 通信均匀分配写入任务到所有 rank（含 DP 副本 rank），消除 IO 不均衡
Atom CheckpointDeepSpeed Universal Checkpoint 引入的中间表示，将各 rank 分片合并为含完整参数 + mapping 元数据的格式，加载时按目标拓扑重切分
GPUDirect Storage (GDS)NVIDIA 技术：NVMe DMA 引擎直接读写 GPU 显存，绕过 CPU bounce buffer
Universal Checkpoint (UCP)DeepSpeed 的跨拓扑 checkpoint 转换工具，支持改变 TP/PP/DP 度后恢复训练

@tbl-dist-ckpt-terms 名词定义：名词、定义

单点 IO 瓶颈与问题陈述​

核心问题：集中式 Checkpoint 的单点 IO 瓶颈在哪、为什么必须分布式？

传统 torch.save(state_dict, path) 在分布式训练中存在三类根本性瓶颈：

  1. 数据通信：所有 rank 的参数必须 all-gather 到 rank 0，跨节点搬运数据量等于全量 checkpoint
  2. IO 串行：仅 rank 0 写入，写带宽受单节点 NVMe / 网络带宽限制
  3. 训练阻塞：上述两步串行执行期间训练完全暂停

AWS Storage Blog (2025)[2] 量化了代价：4000-GPU 集群每次同步 checkpoint 暂停约 3 分钟，每日损耗约 9600 GPU·小时。对 Llama 3 405B 训练（16K H100）规模，单点 IO 路径完全不可行。

分布式 checkpoint 的目标：把 "all-gather 到 rank 0 → rank 0 写入" 替换为 "每 rank 写自己持有的 shard + 跨 rank 元数据协调"，使 IO 带宽随集群规模线性扩展。

主流框架（PyTorch DCP / Megatron-Core / DeepSpeed）的设计共同点：

  • 数据流不跨 rank：每个 rank 只写自己持有的本地 shard
  • 元数据流集中：coordinator rank 维护全局 layout 元数据（offset、shape、shard 归属）
  • Stateful 协议：训练状态（model / optimizer / dataloader / RNG / iteration）通过统一接口被 checkpoint 系统识别

差异点在 sharding 抽象（PyTorch DTensor vs Megatron ShardedTensor vs DeepSpeed ZeRO partitioned states）、resharding 是否原生支持、与并行框架（FSDP / Megatron TP+PP / ZeRO）的耦合程度。

数据布局与跨节点分片​

核心问题：分布式 Checkpoint 的数据如何跨节点分片布局、分片策略如何影响 Save/Load 性能？

PyTorch DCP — 三层解耦​

PyTorch DCP[3] 将 "保存什么 / 怎么保存 / 保存到哪" 三个关注点分层解耦：

层职责核心类
State Dict 层定义 FQN → Tensor 映射，含 DTensor / ShardedTensorStateful 协议
Planner 层跨 rank 元数据协调、分片策略决策SavePlanner, LoadPlanner
Storage 层物理 IO，可替换StorageWriter, FileSystemWriter

@tbl-dist-ckpt-dcp-layers PyTorch DCP 三层解耦设计

Planner 层的核心数据结构：

  • WriteItem：单次写操作的描述，含 index（MetadataIndex 位置）、type（TENSOR / SHARD / BYTE_IO）、tensor_data（分片 chunk 元数据）
  • SavePlan：本 rank 的全部 WriteItem 列表
  • ReadItem：Load 时的读操作描述，含源/目标 index、offset、hypercube 尺寸（支持任意切片交集）

DefaultSavePlanner 的关键参数：

  • flatten_state_dict=True：嵌套 dict 展平为 FQN 映射
  • flatten_sharded_tensors=True：支持 FSDP 2D 并行（TP + FSDP 同时存在）
  • dedup_replicated_tensors：去重 DDP 各 rank 重复持有的完整参数

Checkpoint 目录结构：每 rank 至少生成一个 .distcp 数据文件，coordinator rank 额外写一个 .metadata 全局元数据文件。

Megatron-Core — ShardedTensor 与 rank-aware sharding​

Megatron-Core 的 dist_checkpointing 包[4] 用 ShardedTensor 显式编码 sharding 信息，每个分布式参数包装为一个对象，关键字段：

字段含义
key全局张量的唯一标识符（跨 rank 共享同一个 key）
local_shape本 rank 持有的张量形状
global_shape完整全局张量形状
global_offset本 rank 分片在全局张量中的起始偏移
axis_fragmentations每个轴被切分成多少份（各维度的并行度）
replica_id标识本 rank 是否为主副本，0 = 主副本（负责实际写入）

@tbl-dist-ckpt-megatron-fields Megatron ShardedTensor 关键字段

构造方法 ShardedTensor.from_rank_offsets() 接受多个 (axis, axis_rank_offset, axis_fragm) 元组，编码不同并行维度：

  • TP 映射：列并行层（Q/K/V 投影、MLP up_proj）沿输出维度切分，行并行层（O 投影、MLP down_proj）沿输入维度切分，切分信息编码在 axis_fragmentations 和 global_offset 中
  • PP 映射：不同 pipeline stage 持有不同层的参数，PP 不产生同一张量的 sharding，而是参数的"按层分区"（partition by layer）
  • DP 映射：模型参数在 DP 维度复制，通过 replica_id = dp_rank 表达，只有 replica_id == 0 的 rank 实际写入（其余跳过）
  • CP / EP：通过额外的 rank_offsets 元组叠加进 ShardedTensor

DeepSpeed — ZeRO 分区状态​

DeepSpeed checkpoint[5] 文件布局直接反映 ZeRO 分区模式：

<save_dir>/global_step<N>/
  latest                                              # 最新 tag 指针
  mp_rank_00_model_states.pt                          # 模型权重（ZeRO-1/2，rank 0 完整保存）
  zero_pp_rank_X_mp_rank_YY_model_states.pt           # 模型权重（ZeRO-3，每 rank 占位符）
  bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt     # rank 0 的优化器分片
  bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt     # rank 1 的优化器分片
  ...

ZeRO 各 stage 的 checkpoint 行为差异：

Stage优化器状态模型权重备注
ZeRO-1按 DP 分片，每 rank 写自己分片rank 0 完整保存最常用生产配置
ZeRO-2按 DP 分片rank 0 完整保存张量 flatten 可能导致文件虚胀，需 clone_tensors_for_torch_save() 缓解
ZeRO-3按 DP 分片每 rank 占位符（需 gather 才能重建）需 stage3_gather_16bit_weights_on_model_save=True 才能保存完整 fp16 权重；zero_to_fp32.py 离线合并 fp32 权重（需 2× checkpoint 内存）

@tbl-dist-ckpt-zero-stages ZeRO 各 stage 下的 checkpoint 行为

调用约束：save_checkpoint() 和 load_checkpoint() 必须由所有进程同步调用，不能只让 rank 0 执行——每个 rank 都持有自己的分片，必须各自落盘。

三框架横向对比​

维度PyTorch DCPMegatron-CoreDeepSpeed
Sharding 抽象DTensor + DTensorSpec（DeviceMesh + placements）ShardedTensor（六元组字段）ZeRO partitioned states（隐式按 DP rank）
Resharding 原生支持✅ load 时自动 reshard✅ load 时按 ShardedTensor 交集匹配❌ 需 ds_to_universal.py 离线转换
主面向并行FSDP / DDP / 2D（TP+FSDP）TP + PP + DP + CP + EPZeRO-1/2/3 + Megatron-DeepSpeed 集成
文件后缀.distcp + .metadata.distcp + common.pt + metadata.json.pt（多文件，命名含 pp/mp/dp rank）
跨框架移植性高（PyTorch 原生）中（依赖 torch_dist 格式，建在 PyTorch DCP 之上）低（深度绑定 DeepSpeed 生态）

@tbl-dist-ckpt-framework-compare 三框架 checkpoint 抽象横向对比

关键关系：Megatron-Core 的 torch_dist 格式不是独立实现，而是 PyTorch DCP 的扩展层——strategies/torch.py::TorchDistSaveShardedStrategy 把 Megatron ShardedTensor 转换为 PyTorch DCP 所需格式，调用 DCP 的 FileSystemWriter 写入 .distcp 文件。common.pt 和 metadata.json 是 Megatron 自己的扩展，承载 ShardedTensor 定义、resharding 逻辑、fully parallel 策略等 PyTorch DCP 不提供的能力。

Save/Load 流程中的通信​

核心问题：Save 和 Load 流程中各 rank 之间需要什么通信原语、通信量和瓶颈在哪？

PyTorch DCP Save 协议​

DCP save 遵循严格的 6 阶段协议，部分阶段仅 coordinator rank 执行：

阶段操作执行范围通信
1set_up_planner所有 rank无
2create_local_plan所有 rank无
3create_global_plancoordinator onlyreduce_scatter("plan") 汇总
4finish_plan所有 rankbroadcast 全局 plan
5resolve_data + write_data所有 rank无（并行 IO）
6finishcoordinator onlyall_reduce("write") 汇总写结果，coordinator 写 .metadata

@tbl-dist-ckpt-dcp-save-protocol PyTorch DCP Save 协议阶段与通信

源码（state_dict_saver.py）的两个关键 collective 调用点：

# 阶段 2→3: plan 汇总
central_plan = distW.reduce_scatter("plan", local_step, global_step)

# 阶段 5→6: 结果汇总
metadata = distW.all_reduce("write", write_data, finish_checkpoint)

单点 IO 瓶颈如何消除：阶段 5 write_data 完全并行，每个 rank 只写自己持有的 shard，无需等待其他 rank，无数据通信。唯一的跨 rank 通信是元数据协调（plan + metadata），数据量远小于参数本身（KB-MB 量级 vs GB-TB 量级）。

Load 协议与 Resharding​

Load 协议比 save 多一个初始元数据 broadcast：

阶段 1: read_metadata       coordinator 读 .metadata → broadcast 给所有 rank
阶段 2: set_up_planner      所有 rank 接收元数据
阶段 3: create_local_plan   所有 rank 算本 rank 需要读哪些 chunk
阶段 4: create_global_plan  coordinator 汇总读取计划
阶段 5: finish_plan         所有 rank 接收最终 plan
阶段 6: read_data           所有 rank 并行读，无通信
阶段 7: commit_tensor       所有 rank 后处理

Resharding 机制：当 save 拓扑（N ranks）与 load 拓扑（M ranks）不同时，DefaultLoadPlanner 通过 ReadItem 的 hypercube 切片描述，计算每个 rank 所需 chunk 与 checkpoint 存储 chunk 的交集，每个 rank 直接读自己需要的切片，DTensor 的 redistribute 机制处理本地重排，无额外跨 rank 通信。

Megatron-Core 的 resharding 由 TorchDistLoadShardedStrategy 实现，调用方提供新拓扑下的 ShardedTensor（含新的 local_shape / global_offset），策略层按 shard 交集读取磁盘片段拼装到本 rank。Resharding 支持配置严格性（strictness）级别：

级别行为
ASSUME_OK_UNEXPECTED（默认）不检查 checkpoint 中多余 key，零开销
LOG_UNEXPECTED记录多余 key
RAISE_ALL任何不匹配都报错
RETURN_ALL返回 missing / unexpected key 集合

Fully Parallel Save / Load 优化​

默认 save 策略每个 rank 只写 replica_id == 0 的 shard，DP 副本中其余 rank 空闲，磁盘写入未充分并行化。Megatron 的 FullyParallelSaveStrategyWrapper 通过 metadata-only 通信解决此问题：

  1. Metadata 交换：all-gather 各 rank 持有的 ShardedTensor 的 key 和 shard id（不传 data）
  2. Greedy 分配：用贪心算法将 save 任务均匀分配到所有 rank
  3. 修改 replica_id：被分配到写入任务的 rank 设为 0，其余设为 1
  4. 调用底层策略：TorchDistSaveShardedStrategy 执行实际写入

特点：

  • 只交换 metadata 不交换 data，通信量极小
  • 支持 do_cache_distribution=True 缓存分配结果，后续 checkpoint 无需重新通信
  • 与 async_save 正交：可叠加 GPU→CPU 同步拷贝 + 后台 FileSystemWriterAsync 异步落盘

FullyParallelLoadStrategyWrapper 对称：交换 load metadata → 均匀分配读取任务 → 各 rank 读取被分配的 shard → 通过 exchange_by_distribution() 在 rank 间互传 shard。exchange 支持三种算法：broadcast（各 rank 逐一广播）、gather_object（all_gather_object 整 state dict）、gather_rounds（分轮 all_gather 单张量）。

DeepSpeed Universal Checkpoint 跨拓扑迁移​

ZeRO checkpoint 按当前并行拓扑（TP×PP×DP）切分，直接用于不同 GPU 数量会 shape 不匹配。Universal Checkpoint[6] 引入 atom checkpoint 中间表示：

Extract → Union → StripPadding → GenUcpMetadata → Load
  • 保存阶段：零额外开销，沿用现有 distributed checkpoint
  • 离线转换：ds_to_universal.py 把分布式碎片合并为 atom checkpoint（含完整参数 + mapping 元数据）[7]
  • 加载阶段：读取 atom checkpoint，按目标拓扑计算分区边界，on-demand 重切分，开销 1.14×–1.37× 标准加载时间

Megatron-DeepSpeed 中 ZeRO stage 的拓扑变更支持矩阵[8]：

Stage支持的拓扑变更
Stage 1TP、PP、SP、DP 均可重配
Stage 2需 --no-pipeline-parallel；DP 变更
Stage 3需 --no-pipeline-parallel；仅 DP 变更

@tbl-dist-ckpt-ucp-matrix DeepSpeed Universal Checkpoint 的 ZeRO stage 拓扑变更支持

跟 PyTorch DCP 的关键差异：DCP 在设计上将 reshard 视为一等公民，load 时只需预分配好 state_dict 即可自动重分配；DeepSpeed UCP 是事后补充的中间层，需显式转换步骤，且对高阶并行（TP + stage3 + PP）支持仍有限制。

IO 带宽与传输路径​

核心问题：GDS/NVMe-oF/RDMA 三种传输路径的带宽、延迟和适用场景差异？

Checkpoint 大小估算​

混合精度训练（BF16/FP16 前向 + FP32 优化器）的内存组成：

组件精度字节/参数
模型参数BF16/FP162
Adam 动量 (m)FP324
Adam 方差 (v)FP324
FP32 主参数副本FP324
梯度BF16/FP162

@tbl-dist-ckpt-mixed-precision-mem 混合精度训练每参数内存组成

合计 16 字节/参数，但 checkpoint 通常只包含参数 + 优化器状态 = 14 字节/参数（梯度可重建不必保存）：

$$\begin{equation} \text{Size}_\text{ckpt} = 14 \cdot N_\text{params} \quad \text{(bytes)} \label{eq:dist-ckpt-size} \end{equation}$$

其中：

  • $N_\text{params}$：模型总参数量
  • 系数 14 = 2 （BF16 参数） + 4 (FP32 Adam $m$) + 4 (FP32 Adam $v$) + 4 （FP32 主参数副本）
  • 假设前提：BF16 前向 + FP32 Adam 优化器，且 checkpoint 包含 FP32 主参数副本

此公式是上界。生产实现可省去 FP32 主参数副本（可从 BF16 模型 + optimizer 重建），从而降到 10 字节/参数；以 Meta Llama 3 405B 为例，实际 checkpoint 总量 4.05 TB（810 GB BF16 参数 + 3240 GB FP32 $m+v$），对应 10 字节/参数，不含 FP32 主参数副本（见 工业实测案例与经验法则 节）。

FastPersist[9] 实测（含 FP32 主参数）验证 14·N 上界：6.7B 模型 ≈ 88 GB，13B ≈ 173 GB，与公式一致。

ZeRO 分片对 per-rank 大小的影响：以 ZeRO-1 为例（最常用生产配置），每 rank checkpoint 大小：

$$\begin{equation} \text{Size}_\text{per\_rank} = \frac{12 \cdot N_\text{params}}{N_\text{DP}} + 2 \cdot N_\text{params} \quad \text{(bytes)} \label{eq:dist-ckpt-zero1-per-rank} \end{equation}$$

其中：

  • $N_\text{params}$：模型总参数量
  • $N_\text{DP}$：数据并行（DP）维度的 rank 数
  • 第一项 $12 \cdot N / N_\text{DP}$：FP32 优化器状态（Adam $m$ 4 + $v$ 4 + 主参数 4 = 12 字节/参数）按 DP 分片
  • 第二项 $2 \cdot N$：BF16 参数全量复制（ZeRO-1 不分片参数）
  • 假设前提：BF16 前向 + FP32 Adam 优化器，含 FP32 主参数副本

ZeRO-3 把第二项也压缩为 $2 \cdot N / N_\text{DP}$（参数也按 DP 分片）。

TP / PP 影响：每 TP rank 只持有模型权重的 1/N_TP 分片，每 PP stage 只保存对应层的参数和优化器状态。每 rank 写入量 = 全量 / 各并行维度的乘积，但总 IO 量等于全量，所有 rank 并发写。

存储层级与带宽​

分布式训练常见的 checkpoint 存储链路：

GPU HBM
  ──[PCIe / NVLink]──> CPU DRAM (pinned buffer)
    ──[PCIe / NVMe]──> 本地 NVMe SSD
      ──[IB / Ethernet]──> 并行文件系统 (Lustre / GPFS / WekaFS) / NAS-NFS
        ──[网络]──> 对象存储 (S3 / GCS / OSS)

各层实测带宽（单节点写）：

存储层带宽备注
CPU DRAM (in-memory ckpt)50–200 GB/s受 PCIe 或 NVLink 限速
本地 NVMe （PCIe Gen4, 1 盘）~7 GB/s 顺序写Samsung 990 Pro 等
本地 NVMe （PCIe Gen5, 1 盘）12–14.9 GB/s 顺序写WD SN8100, GIGABYTE Gen5
本地 NVMe （多盘并联，Gen5）峰值 ~25 GB/s/nodeFastPersist 8 节点聚合
NFS（VAST/WekaFS 现代实现）1.6–8 GB/s/node取决于后端及并发
Lustre/GPFS 并行文件系统数百 GB/s 集群聚合Polaris 650 GB/s
S3/对象存储（直接挂载）100–300 MB/s/node受请求延迟限制

@tbl-dist-ckpt-storage-bw 各存储层 checkpoint 写带宽参考

并行文件系统聚合带宽实例：

  • Argonne Polaris（ALCF）：160 OST + 40 MDT，集群聚合 650 GB/s
  • NVIDIA Eos DGX SuperPOD（DDN EXAScaler Lustre）：4 TB/s
  • Pure Storage FlashBlade//EXA（Lustre）：设计目标 10+ TB/s

实际 per-node 写带宽：DataStates-LLM 在 Polaris 64 节点上达到 ~10 GB/s/node（Lustre 后端）[10]。CoreWeave 在 VAST NFS 64 节点集群上稳定 ~8 GB/s/node（每节点 7.94 GiB/s）。

对象存储不适合做主 checkpoint 路径：单节点 100–300 MB/s 写带宽约为本地 NVMe 的 1/30–1/100，加上数十 ms 请求延迟，适合作为异步归档层（先写本地 NVMe → 再后台同步到 S3）。

传输路径 — CPU 中介 vs GPUDirect Storage vs RDMA​

CPU 中介路径（传统）：

GPU HBM ──[NVLink/PCIe]──> CPU DRAM (bounce buffer) ──[PCIe]──> NVMe

CPU DRAM 作为中转，受 CPU PCIe 带宽限制。HGX/DGX 服务器 CPU 侧 PCIe Gen4 x16 单向上限约 25 GB/s，实践中常为瓶颈。

GPUDirect Storage (GDS) 直通路径：

GPU HBM ──[PCIe P2P DMA]──> NVMe SSD  (bypass CPU DRAM)

GDS 通过 GPU BAR aperture 使 NVMe DMA 引擎直接读写 GPU 显存。NVIDIA 设计指南[11] 实测：

  • 4 盘 RAID-0 配置 GDS：13.3 GB/s
  • A100 HGX 启用 GDS：可用带宽从 CPU 限速的 25 GB/s 提升至 50 GB/s

NVMe-oF over RDMA（远端 NVMe）：

GPU HBM ──[GDS + RDMA]──> 远端 NVMe (InfiniBand / RoCE)

NVIDIA 设计指南实测：NIC 路径（NVMe-oF over RDMA）可达 11 GB/s/node，接近本地 NVMe-oF 性能。

RDMA vs TCP 对 checkpoint 的影响：

指标TCP/IP（标准以太网）RDMA（IB / RoCE）
协议栈延迟µs–ms1–2 µs
CPU 占用高（软件协议栈）极低（zero-copy）
实测带宽（100GbE vs NDR IB）~12.5 GB/s~50 GB/s

@tbl-dist-ckpt-rdma-tcp checkpoint 传输路径中 RDMA 与 TCP 对比

ACM APSys 2025 论文 "Indispensable CPU-centric Checkpointing for GPUs"（DOI 10.1145/3725783.3764394）指出 GDS 在实际系统中受 GPU HBM 出口带宽、PCIe switch 拓扑、多路 P2P 共享等约束，CPU-centric 路径在多 rank 并发竞争 NVMe 时反而有竞争力（CPU 可缓冲并重排 IO）。实践结论：GDS 在单路大块顺序写优势明显；CPU-centric 在多 rank 并发写时有竞争力。

工业实测案例与经验法则​

核心问题：工业界分布式 Checkpoint 的实测数据和经验法则（如 optimal shard size、rank-to-storage ratio）是什么？

关键实测数据汇总​

@tbl-dist-ckpt-industry-cases 工业 checkpoint 实测关键数字汇总

机构模型 / 规模Checkpoint 大小写入时间数据来源
MetaLlama 3 405B (16K H100)4.05 TB（总）= 810 GB BF16 参数 + 3240 GB FP32 m+v；31.6 GB/rank（FSDP 128 假设）；不存 FP32 主参数副本（10 bytes/param）~13.5 min 理论（异步覆盖）Llama 3[12]
字节（MegaScale）175B，12,288 GPU未披露总耗时未披露（采用异步两阶段，时序细节见 10.3 异步 Checkpoint）MegaScale[13]
Microsoft（CheckFreq）DNN 通用—目标 ≤5% 开销；静态方案可达 35%CheckFreq[14]
Microsoft（Nebula，已废弃）GPT2-XL 20.6 GB—节省 96.9% 时间Azure ML 文档[15]
Microsoft（Nebula，已废弃）97 GB on 128×A10097 GB20 min → 1 s同上[15]
NVIDIA（NeMo）Llama 3.1 8B~105 GB~24.9 s（sharded distcp）dgxc-benchmarking[16]
FastPersist8 节点 128 A100—8 节点聚合 ~130 GB/s（单节点峰值 24.8 GB/s）vs torch.save 基线FastPersist[9]
DataStates-LLM64 节点 256 A100 Lustre—10 GB/s/nodeDataStates-LLM[10]
业界基准大模型训练通用—checkpoint 开销平均 12%，最高 43%Maeng et al., MLSys 2021

经验法则​

以下规律来自多个公开工业报告与研究论文的共识：

频率选择

  • 频率 ↔ 开销权衡：checkpoint 越频繁，单次故障丢失工作越少，但 IO 开销越大。CheckFreq 建议把开销约束在 ≤5% compute time。
  • 故障率驱动频率：10K+ GPU 集群实测故障频率约每多周训练 100+ 次，故障间隔约数小时。合理 checkpoint 间隔应小于平均故障间隔（MTBF），但需考虑 IO 成本。最优解形如 P* ∝ sqrt(T_ckpt / λ)（Young/Daly checkpoint 优化，详见 10.1 集群可靠性总览 中的 ETTR 模型）。
  • 大规模工业系统普遍采用异步两阶段 checkpoint（MegaScale / Nebula / NeMo[17]），通过 GPU→CPU 内存暂存解除训练阻塞；时序、阶段划分与重叠策略见 10.3 异步 Checkpoint。

体积规律

  • 参数量决定基础体积：BF16 存储下每 10 亿参数约 2 GB（1× BF16 参数基准）；加上 FP32 Adam 优化器状态（m + v 各 4 bytes/参数，合计 4× BF16 参数），总 checkpoint ≈ 5× BF16 参数大小（10 bytes/参数，不含 FP32 主参数副本）或 7× BF16 参数大小（14 bytes/参数，含主参数副本）。
  • FSDP/ZeRO 线性分片：完全分片并行（FSDP / ZeRO-3）下每 GPU checkpoint 数据量 ≈ 总大小 / GPU 数。Llama 3 405B 在 128 GPU FSDP 下实测 ~31 GB/rank，与公式吻合。

恢复时间

  • MegaScale 目标：从最新 checkpoint 恢复训练进度 ≤15 分钟（含追赶），保证有效训练时长 > 90%
  • 学术前沿（TrainMover, arXiv:2412.12636，OSDI '26 接收）：通过预热备用 GPU，在 32 GPU 实测中把恢复停机时间压至 <10 秒，相比传统 ~300 秒减少 97%（论文引用 Llama 3 1024 GPU 场景但未实测该规模）

开放问题​

核心问题：分布式 Checkpoint 尚未解决的关键开放问题有哪些？

  • GDS 在多 rank 并发竞争场景下的最优策略：ACM APSys 2025 论文指出 CPU-centric 在多路并发竞争 NVMe 时有竞争力，但缺乏系统性 benchmark 量化"何时 GDS 反而不如 CPU-centric"。
  • Resharding 的开销建模：DeepSpeed Universal Checkpoint 给出 1.14×–1.37× 标准加载时间的开销，但 PyTorch DCP / Megatron 在大规模拓扑变更（如 TP=4→16, PP=8→16 同时变）下的开销缺公开数据。
  • Checkpoint 体积压缩技术：当前研究偏写带宽优化（FastPersist / DataStates-LLM），但少有工作探讨 checkpoint 内容本身的压缩（如 quantized checkpoint / delta checkpoint）对 IO 和恢复正确性的权衡。
  • 跨厂商 checkpoint 格式互通：DeepSpeed ZeRO 状态、Megatron torch_dist、PyTorch DCP .distcp 三者之间缺少公开的双向转换工具，影响模型在不同训练栈之间迁移。
  • 百万 GPU 规模下的 checkpoint 频率最优解：现有 Young/Daly 公式基于泊松故障假设，在 100K+ GPU 规模下故障非独立同分布（lemon node、共因故障），最优 P* 需要修正。
  • 推理服务的 KV cache checkpoint：训练侧 checkpoint 研究丰富，推理侧（KV cache 持久化、多副本一致性）的研究和公开实测几乎空白。

参考资料​

  1. Eisenman et al., Check-N-Run: a Checkpointing System for Training DLRM, NSDI 2022. https://www.usenix.org/conference/nsdi22/presentation/eisenman。Facebook 生产环境网络存储瓶颈分析，DLRM 场景增量 delta 压缩。
  2. AWS Storage Blog, Architecting scalable checkpoint storage for large-scale ML training on AWS, 2025. https://aws.amazon.com/blogs/storage/architecting-scalable-checkpoint-storage-for-large-scale-ml-training-on-aws/。量化单点 checkpoint 代价：4000-GPU 集群每日损耗 9600 GPU·小时。
  3. PyTorch 官方，Distributed Checkpoint (DCP) API Reference. https://docs.pytorch.org/docs/main/distributed.checkpoint.html。完整 API、Planner / StorageWriter 协议序列。
  4. NVIDIA，Megatron-LM dist_checkpointing 源码. https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/dist_checkpointing。ShardedTensor、serialization、fully_parallel、optimizer 实现。
  5. Microsoft 官方，DeepSpeed Model Checkpointing API. https://deepspeed.readthedocs.io/en/latest/model-checkpointing.html。save/load API 参数 + ZeRO stage 行为差异。
  6. Lian et al., Universal Checkpointing: Efficient and Flexible Checkpointing for Large Scale Distributed Training, arXiv:2406.18820, 2024. https://arxiv.org/html/2406.18820v2。DeepSpeed UCP 设计：atom checkpoint + declarative operators，跨拓扑迁移开销 1.14×–1.37×。
  7. Microsoft 官方，DeepSpeed Universal Checkpointing Tutorial. https://www.deepspeed.ai/tutorials/universal-checkpointing/。`ds_to_universal.py` 使用流程。
  8. DeepSpeed，Megatron-DeepSpeed Universal Checkpointing README. https://github.com/deepspeedai/Megatron-DeepSpeed/blob/main/examples_deepspeed/universal_checkpointing/README.md。ZeRO stage 拓扑变更支持矩阵。
  9. Wang et al., FastPersist: Accelerating Model Checkpointing in Deep Learning, arXiv:2406.13768, 2024. https://arxiv.org/html/2406.13768v1。NVMe libaio/io_uring 优化 + 多盘并联，8 节点 128 A100 上 130 GB/s 聚合写带宽。
  10. DataStates-LLM: Scalable Checkpointing for Transformer Models, arXiv:2601.16956, 2026. https://arxiv.org/html/2601.16956v1。Polaris 256 A100 Lustre 实测 ~10 GB/s/node，序列化开销分析 22%。
  11. NVIDIA 官方，GPUDirect Storage Design Guide. https://docs.nvidia.com/gpudirect-storage/design-guide/index.html。GDS 直通路径带宽：4 盘 RAID-0 13.3 GB/s、HGX 50 GB/s、NVMe-oF 11 GB/s。
  12. Dubey et al., The Llama 3 Herd of Models, arXiv:2407.21783, 2024. https://arxiv.org/abs/2407.21783。Llama 3 405B 训练 checkpoint 实测：4.05 TB 总量、1MB-4GB/GPU 范围。
  13. Jiang et al., MegaScale: Scaling LLM Training to More Than 10,000 GPUs, NSDI 2024 (arXiv:2402.15627). https://arxiv.org/abs/2402.15627。两阶段异步 checkpoint，10K+ GPU 100+ 次故障，<15 分钟追赶。
  14. Mohan et al., CheckFreq: Frequent, Fine-Grained DNN Checkpointing, FAST 2021. https://www.usenix.org/conference/fast21/presentation/mohan。迭代粒度自适应频率，目标 ≤5% 开销。
  15. Microsoft 官方，Azure ML: Optimize Checkpoint Performance for Large Models (Nebula). https://learn.microsoft.com/en-us/azure/machine-learning/reference-checkpoint-performance-for-large-models。Nebula 实测 + 废弃声明 + 基准 12%/43% 引用。
  16. NVIDIA，DGX Cloud Llama 3.1 Benchmark. https://github.com/NVIDIA/dgxc-benchmarking/blob/main/llama3.1/README.md。NeMo Llama 3.1 8B distcp 实测 ~24.9 s。
  17. NVIDIA 官方，NeMo Framework Distributed Checkpoints User Guide. https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/checkpoints/dist_ckpt.html。ShardedTensor 设计 + save/load 流程 + PyTorch DCP 集成关系。
  18. Wang et al., GEMINI: Fast Failure Recovery in Distributed Training with In-Memory Checkpoints, SOSP 2023. https://dl.acm.org/doi/10.1145/3600006.3613145。CPU 内存 checkpoint，网络存储基线 3.13 GB/s。
  19. DeepSeek-V3 Technical Report, arXiv:2412.19437, 2024. https://arxiv.org/abs/2412.19437。2048 H800 训练 14.8T tokens 无 rollback，FP8 训练但 checkpoint 细节未披露。

项目内交叉引用：

  • 10.1 集群可靠性总览 — checkpoint 频率 P 的 Young/Daly 最优解推导与 ETTR 模型
  • 4.1 总览 — Save/Load 协议中用到的 reduce_scatter / all_reduce / broadcast 算法
  • 1 总览 — ZeRO / FSDP / TP / PP 的分区策略
  • 1.1 总览 — NVLink / PCIe / IB / RoCE 带宽规格

延伸阅读​

正文未直接讨论，作为分布式 checkpoint 的补充实现参考：

Takeaway​

知识点核心结论
Sharded Checkpoint每 rank 只写自己分片，所有 rank 并行 IO；支持 N → M rank resharding
Save / Load 协议Coordinator rank 汇总元数据，各 rank 并行写；ShardedTensor 含 global_shape / local_shape / offset / axis_fragmentations / replica_id
传输路径传统 HBM → CPU → Storage 走 PCIe + CPU; GDS 直通绕开 CPU 直 NVMe ⇄ HBM
GDS 带宽4 盘 RAID-0 13.3 GB/s; HGX 50 GB/s; NVMe-oF 11 GB/s
Llama 3 实测Tectonic 聚合 2 TB/s 但单 job ~5 GB/s, 405B 4.05 TB 写 ~13.5 min
Universal CheckpointingTP × PP × DP × ZeRO 切分维度可变下做跨配置恢复
Check-N-RunDLRM embedding 稀疏更新，只写 delta 大幅压缩
MegaScale 两阶段Stage 1 数秒 + Stage 2 异步 HDFS, 10K+ GPU < 15 分钟追赶
经验法则4× 参数大小 (BF16 + FP32 Adam); FP8 训练 (DeepSeek V3) 显著降低但未量化

Limitations​

核心问题：本文调研在数据来源、基准测试和时效性上存在哪些局限？

  • 本章不涵盖 async / overlap 时序细节（CheckFreq 的两阶段、MegaScale Stage 1/2 的具体调度、Gemini in-memory 写入策略[18]），这部分留给 10.3 异步 Checkpoint
  • 三框架对比基于 2025-2026 公开文档，DeepSpeed Nebula 已废弃但仍列出（其性能数据仍为业界基准引用）；PyTorch DCP 在 PyTorch 2.x 持续演进，部分 API 可能在新版本变化
  • 工业实测数据多为单家厂商单次报告，缺少跨厂商同条件对比；DeepSeek V3、xAI Colossus 未公开 checkpoint 细节
  • ZeRO 各 stage 在不同 framework（Megatron-DeepSpeed vs 纯 DeepSpeed vs DeepSpeed-MII）下的 checkpoint 行为可能有细微差异，本文以纯 DeepSpeed 文档为准
  • GDS 带宽数字基于 A100 HGX 硬件，H100/B200 等新平台数字可能有出入
  • 经验法则中"4× 参数大小"是 BF16+FP32 Adam 假设；FP8 训练（DeepSeek V3[19]）下 checkpoint 大小会显著降低，但公开数据未量化