跳到主要内容

异步 Checkpoint

两阶段 snapshot/persist 如何与训练重叠，以及一致性保证

核心要点：

  • 两阶段流水线：snapshot (GPU → CPU) + persist （CPU → 存储），persist 完全异步与训练并行
  • 同步阻塞代价：GPT-3 175B 10-20 分钟，OPT 175B 21 分钟，Llama 3 405B ~13.5 分钟
  • CheckFreq：自适应频率 $k = \max((T_c + T_s - T_o)/T_i, \lceil T_o / (p T_i) \rceil)$，把开销维持 $p$ （典型 5%） 以内
  • GPU vs CPU 模式：GPU 内存够则走 GPU-to-GPU copy 不阻塞；否则走 D2H 与 host 加载竞争 stream
  • Gemini 三层放置：本地内存 / 远端内存 / 磁盘，利用集群闲置 DRAM 做 in-memory checkpoint
  • PyTorch DCP Cached Plan: checkpoint 结构不变时复用 collective 规划，仅传增量，6× 加速

本文聚焦 checkpoint 与训练计算的时序重叠 — 两阶段 (snapshot + persistence) 设计、调度策略、一致性与正确性。不写跨节点数据布局 / Save-Load 协议 / IO 带宽 (见 10.4 分布式 Checkpoint 通信), 不写弹性训练 (见 10.5 弹性训练), 不写 SDC (见 10.6 Silent Data Corruption (SDC))。

名词定义​

核心问题：异步 Checkpoint 涉及的核心名词（Snapshot/Persist/CheckFreq/一致性/增量）的定义是什么？

名词定义
Async Checkpoint把 checkpoint 持久化操作从训练主循环解耦，与下一 iteration 的 forward/backward 并行执行的机制
Snapshot在某一时刻把训练状态（参数 + 优化器状态 + RNG）冻结成一致内存副本的操作
Persist把 snapshot 副本异步写入持久化存储（disk / 远端节点内存 / 对象存储）的操作
Snapshot Atomicitysnapshot 期间训练若继续，参数被下一 step 修改，需机制保证 snapshot 内容对应单一 training step
Stop-the-worldsnapshot 期间整个训练暂停，保证一致性，代价是 GPU idle
Copy-on-write (COW)写时复制：snapshot 后训练继续，对参数的修改触发副本拷贝，保护 snapshot 内容
Lazy copy-out利用 forward/backward 期间参数 immutable 的窗口，在该窗口后台 GPU→Host 传输，更新前完成
Mixed Placement PolicyGemini 提出的三层 checkpoint 放置策略：本地内存 / 远端内存 / 磁盘
Interleaved Scheduling把 worker 分组错峰写入 checkpoint，避免全局同步阻塞
Adaptive Checkpoint Frequency通过 profiling 训练步耗时与 checkpoint 开销，动态调整 checkpoint 间隔
Cached PlanPyTorch DCP 优化：checkpoint 结构不变时复用 collective 规划，只传增量内容

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

同步 Checkpoint 的训练阻塞问题​

核心问题：同步 Checkpoint 为何会阻塞训练、阻塞时间的主要构成是什么？

传统同步 checkpoint 把以下三步全部串行执行在训练主循环上：

  1. 冻结状态：暂停训练，确保模型参数/优化器状态不被下一 step 修改
  2. GPU→Host 拷贝：通过 PCIe DMA 把 GPU 显存中的状态复制到 CPU 内存
  3. Host→Storage 写入：把 CPU 内存数据写入持久化存储（NVMe / NFS / HDFS / 对象存储）

三步串行的代价（来自工业实测）：

模型规模同步 checkpoint 阻塞时间来源
GPT-3 175B10–20 分钟DataStates-LLM 论文综述
OPT 175B21 分钟（每次 2.1 TB 写 NFS）OPT 训练日志[1]
Llama 3 405B约 13.5 分钟理论估算（4.05 TB 写入 / 单客户端 ~5 GB/s 估算带宽；Tectonic 聚合带宽 2 TB/s 是全集群上限，单 job 实际可用带宽远低于此）Llama 3 paper[2]
97 GB checkpoint 同步写对象存储20 分钟（128×A100）Nebula 文档[3]

@tbl-async-ckpt-sync-cost 同步 checkpoint 在工业 LLM 训练中的阻塞时间

Maeng et al. (CPR, MLSys 2021)[4] 量化业界基准：checkpoint 开销平均占总训练时间 12%，极端场景达 43%。OPT-175B 训练日志[1]：2 个月窗口内约 35 次手动重启（Gemini 论文[5] 引用同期 OPT 训练观察到约 110 次故障），每次磁盘 checkpoint 写盘约 21 分钟（2.1 TB 写 NFS），故障 + 恢复开销构成训练算力主要损失项。

Async checkpoint 的核心思路：将上述三步重新组织 — Step 1 + Step 2（snapshot）依然要做但尽量短，Step 3（persist）完全异步化与下一 iteration 计算重叠。剩余的问题就是"如何把 snapshot 时间压缩到与单步训练相当"以及"snapshot 阶段如何保证一致性"。

两阶段抽象：Snapshot + Persist​

核心问题：Snapshot 和 Persist 如何分离解耦、各自的时间和资源开销是什么？

业界 async checkpoint 方案高度收敛到同一架构 — 两阶段流水线：

迭代 i:  forward → backward → weight_update → [snapshot_i 启动]
                                                    ↓ (重叠)
迭代 i+1: forward → backward → [snapshot_i 完成] → weight_update
                                                    ↓
                                              persist_i 后台异步落盘
                                                    ↓ (与训练继续并行)
                                              persist_i 完成
阶段操作关键路径数据位置
SnapshotGPU 状态冻结一致副本 → CPU pinned memory（或同 GPU 内存）部分关键路径（取决于实现）GPU HBM → CPU DRAM
PersistCPU 内存 → 持久化存储（disk / 远端节点 / 对象存储）完全异步，不阻塞训练CPU DRAM → Storage

@tbl-async-ckpt-two-stage Async checkpoint 两阶段抽象

一致性约束：snapshot 必须完成于下一次 weight update 开始之前，否则 snapshot 会捕获到被部分更新的不一致状态。这是 async checkpoint 设计的核心约束，下文 Snapshot 一致性 节展开。

关键工程问题：snapshot 阶段是否真的能压到亚秒？答案取决于：

  • GPU 内存是否够放 snapshot 副本（决定走 GPU-to-GPU copy 还是 GPU-to-CPU D2H）
  • 是否使用 pinned memory（决定 PCIe DMA 还是 staging copy）
  • 框架是否支持 collective 协调的 cross-rank snapshot 同步
  • snapshot 是否能与 forward/backward 期间的参数 immutability 窗口重叠

下文按"经典设计 → 学术创新路线 → 工业实测 → 框架实现"四段展开。

CheckFreq：自适应频率与两阶段流水线​

核心问题：CheckFreq 如何根据训练速度和 Persist 耗时自适应调整、两阶段流水线如何组织？

CheckFreq (USENIX FAST 2021, Mohan et al., UT Austin + Microsoft Research)[6] 是第一个系统性解决"高频 checkpoint + 低开销"矛盾的框架，奠定了 async checkpoint 的两阶段范式。

两阶段时序设计​

CheckFreq 把传统同步 checkpoint 拆分为：

  • snapshot()：GPU → CPU 一致内存快照
  • persist()：CPU → disk 异步写入（后台进程，永不阻塞训练）

时序关键决策：snapshot() 在迭代 $i$ 的 weight update 结束立即触发，与迭代 $i+1$ 的 forward + backward 同时进行。weight update 步骤被 semaphore 包装，迭代 $i+1$ 的 weight update 必须等 snapshot() 完成才能开始，确保一致性。

GPU 是否阻塞 — 两种模式​

模式触发条件snapshot 路径GPU 阻塞对应公式中 $T_o$
GPU 模式（推荐）$M_\text{max} - M > m$（GPU 空闲内存 > checkpoint 大小）GPU 内存内复制（GPU-to-GPU copy）基本不阻塞，代价比 D2H 低一个数量级$T_o = T_{og}$（见 \eqref{eq:async-ckpt-overlap}）
CPU 模式（降级）GPU 内存不足GPU → CPU D2H 传输（PCIe DMA）与下一迭代 host-to-device 数据加载竞争 CUDA 流，造成可见 stall$T_o = T_{oc}$

@tbl-async-ckpt-checkfreq-modes CheckFreq snapshot 的两种模式

模式选择即下文 Algorithm 1 中 $T_o = \min(T_{og}, T_{oc})$ 的依据：GPU 内存够时优先选 GPU 模式（$T_{og}$ 远小于 $T_{oc}$），否则降级 CPU 模式并通过自适应频率算法弥补开销。CPU 模式下 CheckFreq 自动降低 checkpoint 频率避免 stall 在关键路径累积。

Adaptive Frequency Algorithm​

CheckFreq 训练开始时自动 profile 前 1% 迭代（或前 50 次，取小值），收集 8 个指标：

符号含义
$T_i$单次迭代时间
$T_w$weight update 时间
$T_g$GPU 内存内复制时间（GPU 模式 snapshot）
$T_c$GPU→CPU 复制时间（CPU 模式 snapshot）
$T_s$写入存储时间
$m$checkpoint 大小
$M$峰值 GPU 内存使用量
$M_\text{max}$GPU 总内存

@tbl-async-ckpt-checkfreq-metrics CheckFreq profiling 指标

频率算法（Algorithm 1）核心：

$$\begin{equation} T_{oc} = \max\left(0,\ T_c - (T_i - T_w)\right),\quad T_{og} = T_g \label{eq:async-ckpt-overlap} \end{equation}$$ $$\begin{equation} k = \max\left(\frac{T_c + T_s - T_o}{T_i},\ \left\lceil \frac{T_o}{p \cdot T_i} \right\rceil\right) \label{eq:async-ckpt-frequency} \end{equation}$$

其中：

  • $T_{oc}, T_{og}$：CPU / GPU 模式下 snapshot 超出可重叠窗口（$T_i - T_w$）的残余开销
  • $T_o = \min(T_{og}, T_{oc})$：实际选用模式下的暴露开销（GPU 内存够则用 GPU 模式）
  • $k$：checkpoint 间隔（每 $k$ 个迭代做一次 checkpoint）
  • $p$：用户指定的关键路径开销阈值（典型 5%）
  • 第一项 $(T_c + T_s - T_o) / T_i$：摊销完整 checkpoint I/O 代价所需的迭代数
  • 第二项 $\lceil T_o / (p \cdot T_i) \rceil$：保证关键路径残余开销不超 $p$ 的最小间隔

典型例子：checkpoint 代价与迭代时间均为 1 个时间单位、$p = 5\%$，则 $k = 20$。

Adaptive Rate Tuning：静态 profiling 在干扰变化时失效（如存储带宽被其他作业抢占）。CheckFreq 引入反馈驱动动态调频：实测开销超 $p$ 时用实测值重算 $k$ 降频；干扰消退后恢复较高频率。实测 co-location 干扰下，CheckFreq 自动降频把开销维持 5% 以内。

恢复时间界限​

CheckFreq 保证系统中同一时刻最多一个 checkpoint 操作进行。中断时最多回退一个 checkpoint：

$$\begin{equation} R_\text{max} = 2 \cdot k \cdot t_i,\quad R_\text{avg} = k \cdot t_i \label{eq:async-ckpt-rollback} \end{equation}$$

其中 $k$ 为 checkpoint 间隔（迭代数），$t_i$ 为单迭代时间。实验中 $k$ 约为总迭代数的 1/100–1/300，恢复时间从数小时降到秒级。

与早期 HPC 异步 checkpoint 的差异​

DeepFreeze (CCGrid 2020, Nicolae et al.)[7] 是早期面向 CPU HPC 集群的 async checkpoint 系统，通过异步序列化 + I/O 流水线 + 分布式 checkpoint 分散 IO 负载。与 CheckFreq 关键差异：

  • DeepFreeze 不处理 GPU snapshot 代价（GPU 时代最大开销项）
  • checkpoint 频率需手动配置，无自动 profiling 与自适应调节
  • 静态频率无法适应运行时干扰

HPC 经典理论 Daly (Future Generation Computer Systems 2006)[8] 给出基于 MTBF 的最优 checkpoint 间隔：

$$\begin{equation} k^* = \sqrt{2 \cdot C \cdot \text{MTBF}} \label{eq:async-ckpt-daly} \end{equation}$$

其中 $C$ 为单次 checkpoint 代价，$\text{MTBF}$ 为平均故障间隔。CheckFreq 改用 profiling 驱动，不依赖故障分布假设，更适合 DNN 训练的交互式抢占场景。两者形成对照：HPC 路径偏理论最优（假设泊松故障）、DNN 路径偏实测自适应（应对运行时干扰）。

Gemini：In-Memory Checkpoint 与 Mixed Placement Policy​

核心问题：Gemini 的 In-Memory Checkpoint 和 Mixed Placement Policy 如何减少 Persist 瓶颈？

Gemini (SOSP 2023, Wang Z. et al., Rice University + Amazon AWS)[5] 把 checkpoint 目标从磁盘切换到 CPU 内存，并用 RDMA 跨节点复制做容错。核心主张：恢复延迟从数十分钟压缩到秒级。

In-Memory Checkpoint 的时序优势​

每个 GPU worker 把模型状态拷贝到本节点 CPU 内存（同步阶段），再通过 RDMA 异步复制到远端节点 CPU 内存。写入与下一轮 forward/backward 重叠，对训练吞吐影响 < 3%。

跟磁盘 checkpoint 比较的时序维度差异（IO 带宽与传输路径绝对值见 10.4 分布式 Checkpoint 通信 §IO 带宽与传输路径）：

时序维度磁盘 checkpointIn-memory checkpoint
Stage 1 阻塞窗口数分钟（写入串行进存储）亚秒至数秒（GPU→CPU PCIe 拷贝）
恢复路径同步窗口写入串行（DP broadcast 可缓解）RDMA 拉取并行
持久性语义永久（断电安全）仅存活节点有效，需 Tertiary 落盘兜底
与训练 overlap 上限受存储写带宽决定受 RDMA + PCIe 决定

@tbl-async-ckpt-mem-vs-disk In-memory 与磁盘 checkpoint 的时序维度对比

Mixed Placement Policy​

Gemini 的核心设计是三层放置策略，解决单纯本地副本或随机远端副本无法同时满足容错与性能的问题：

层级副本位置适用故障恢复延迟
Primary本节点 CPU 内存（PCIe 直写）单 GPU 进程崩溃（节点存活）30–90 秒
Secondary跨故障域的远端节点 CPU 内存（RDMA 异步复制）整节点宕机1–3 分钟
Tertiary磁盘 checkpoint（低频，与版本里程碑对齐）多节点同时故障15–30 分钟

@tbl-async-ckpt-mixed-placement Gemini Mixed Placement Policy 三层副本策略

远端节点选取原则：跨越潜在共同故障域（同电源单元、同交换机端口的节点视为同一故障域），避免双副本同时丢失。

内存开销：每节点存储本地 Primary + 远端 Secondary，总占用约为：

$$\begin{equation} \text{Mem}_\text{ckpt} = 2 \times M_\text{rank} \times n_\text{gpu} \label{eq:async-ckpt-mem-overhead} \end{equation}$$

其中：

  • $M_\text{rank}$：单 rank 的 checkpoint 大小（含 ZeRO 分片后的参数 + optimizer state）
  • $n_\text{gpu}$：单节点 GPU 数
  • 系数 2 = 1 份本地 Primary + 1 份远端 Secondary

Llama 3 405B FSDP 128 rank 下 $M_\text{rank}$ ≈ 31.6 GB（详见 04-章 工业实测案例），8 GPU/节点配置下每节点约 506 GB CPU 内存，刚好处于主流训练节点 512 GB–1 TB CPU 内存上限。中等规模模型（如 70B FSDP 64 rank）$M_\text{rank}$ ≈ 11 GB，每节点约 176 GB，宽裕。

Interleaved Scheduling​

朴素方案在训练过程中定期暂停所有 worker 同步写 checkpoint，造成全局停顿。Gemini 用交错调度：

  • 所有 worker 分若干逻辑分组，各组在不同训练步骤偏移开始写 checkpoint
  • 同一时刻只有一个分组处于写入状态，其余分组正常训练
  • 写入操作与 forward pass 重叠：GPU 做 forward 时，CPU + RDMA 异步传输上一步骤状态

整体 checkpoint 墙时延被分散到多个训练步，单步额外延迟 < 1%，全局"最新一致 checkpoint"在多步内滚动更新。

Gemini 实验数据（256 GPU = 32 节点 × 8 A100）​

跟磁盘 checkpoint 相比的时序差异（绝对 IO 带宽不在此章范围）：

指标Gemini磁盘 checkpoint
单节点故障恢复时序95 秒（含 RDMA pull + warmup）15+ 分钟（约 10× 加速）
训练吞吐量影响（与 forward overlap）2.4%6–12%

@tbl-async-ckpt-gemini-eval Gemini 论文报告的时序维度数据

DataStates-LLM — 邻近路线​

系统设计差异（时序/一致性维度）恢复时序局限
DataStates-LLM (HPDC 2024, Maurya et al.)[9]Lazy copy-out 利用 forward/backward 期间参数 immutability 窗口，无显式 stop-the-world端到端训练时间减少 2.2×不提供多副本容错（与 Gemini 互补，可组合）

@tbl-async-ckpt-inmem-systems In-memory checkpoint 邻近系统的时序/一致性维度对比

DataStates-LLM 与 Gemini 互补 — DataStates-LLM 优化 snapshot 时序（lazy copy-out 利用 immutability 窗口），Gemini 优化恢复延迟（Mixed Placement），可组合使用。其它对等镜像复制（peer-mirror replication）方案在学术界有探索，但具体系统的归属与 venue 信息存在分歧，本文不展开（见 Limitations）。

工业系统实测：MegaScale / Nebula / NeMo / PyTorch DCP​

核心问题：四个工业级 Checkpoint 系统在规模、吞吐、恢复时间上的实测数据如何对比？

MegaScale Stage 1/2（ByteDance, NSDI 2024）​

MegaScale §4.4[10] 在 10K+ GPU 训练集群上实现两阶段 checkpoint：

阶段操作关键路径实测时间
Stage 1GPU → host pinned memory（PCIe DMA + PyTorch 序列化优化）是"数秒"（论文未给精确值）
Stage 2host memory → HDFS（后台进程异步）否异步，不阻塞训练

@tbl-async-ckpt-megascale-stages MegaScale 两阶段 checkpoint 实测

Stage 2 失败的时序后果：Stage 1 写入 host memory 后训练继续；若 Stage 2 写完前节点崩溃，该 host memory 中的 checkpoint 丢失，必须回滚一个完整 checkpoint 周期，最大 rollback 时间为本章 \eqref{eq:async-ckpt-rollback} 中的 $R_\text{max} = 2 \cdot k \cdot t_i$（多版本保留策略与恢复路径细节见 10.4 分布式 Checkpoint 通信）。

Microsoft Nebula 三阶段（已废弃但仍为基准）​

Nebula[3] 实际是三阶段而非两阶段：

阶段操作同步/异步防护目标
Stage 1GPU → CPU DRAM（约 1 秒）同步阻塞—
Stage 2CPU DRAM → 相邻节点本地 SSD异步单节点故障
Stage 3CPU DRAM/SSD → 对象存储（Azure Blob/ADLS）异步多节点故障 + 长期存储

@tbl-async-ckpt-nebula-stages Nebula 三阶段 checkpoint

核心数据点：97 GB checkpoint 在 128 × A100 上从同步 20 分钟压缩到 Stage 1 阻塞 1 秒。GPT2-XL 20.6 GB 节省 96.9% 写入时间。整体 checkpoint 时间减少 95–99.9%。

内存要求：CPU DRAM 容量 ≥ 3× checkpoint 大小（在途 + 当前版本 + 历史版本）。

Stage 2 失败处理：Nebula 维护多个 checkpoint 版本（num_of_version_in_retention 默认 2 个）。Stage 2/3 失败时 Stage 1 已在 CPU DRAM 中的数据仍可用于重启恢复（节点存活前提）。Nebula 已于 2025 年底被 Microsoft 弃用，不随最新 ACPT 镜像提供，设计仍为业界基准参考。

NVIDIA NeMo / Megatron-Core async checkpoint​

NeMo 用 AsyncFinalizableCheckpointIO + AsyncCallsQueue 实现：

  • Stage 1（同步阻塞）：触发 save，数据从 GPU 写到本地 CPU 内存或本地磁盘
  • Stage 2（异步后台）：AsyncCallsQueue 管理后台进程持久化到分布式存储；训练主进程继续
  • 完成后删除 *-unfinished marker 文件；失败检测依赖 marker 扫描

NVIDIA 实测数据[11]（Megatron-Core v0.7）的端到端 checkpoint 开销减少倍数（归因：减少倍数 = Stage 1 staging 时间压缩 × Stage 2 后台 overlap 比例，二者复合效果，不单独归因于某一阶段）：

模型配置checkpoint 开销减少（vs torch.save）
Nemotron-4 340Bwith distributed optimizer42×
Nemotron-4 15Bwithout distributed optimizer50×
Nemotron-4 340Bwithout distributed optimizer26×

@tbl-async-ckpt-nemo-eval NeMo async checkpoint 在 Nemotron-4 上的开销减少倍数

已知限制：async_save 与 EMA（指数移动平均）不兼容。

PyTorch DCP async（Llama3-70B, 1856 H200）​

PyTorch 团队 2025 实测[12]（Llama3-70B，1856 H200 GPU，TorchTitan + HSDP2）：

指标优化前（无缓存计划）优化后（进程模式 + 缓存计划）
Stage 1 阻塞时间（GPU→CPU）~0.78 秒~0.78 秒（不变，PCIe 带宽限制）
Stage 2 后台写入时长~436 秒/checkpoint~67 秒/checkpoint（6.5×）
TPS 抑制窗口~7 分钟~1 分钟
TPS 峰值下降700 → 320700 → 372

@tbl-async-ckpt-pytorch-dcp PyTorch DCP async checkpoint 在 Llama3-70B 上的优化效果

Cached Plan 机制：checkpoint 结构（层名、形状、分片）在训练期间不变，只数据内容变。首次 checkpoint 完整 collective 规划（昂贵），后续复用缓存计划只传更新的增量，显著减少 Stage 2 协调开销。进程模式：使用 multiprocessing 隔离写盘工作进程避开 Python GIL，进一步压缩 Stage 2 时长。

Llama 3 训练中的 async checkpoint​

Llama 3 技术报告 §3.3[2] 中的关键陈述：

  • Tectonic 存储支持 2 TB/s 持续 / 7 TB/s 峰值吞吐（7,500 SSD 服务器，240 PB）
  • 重大挑战：checkpoint 写入是"高度突发"负载，会短时间饱和存储网络
  • 每 GPU 状态 1 MB–4 GB（覆盖各尺寸 Llama 3 模型）
  • 显式目标：最小化 GPU pause time + 提高 checkpoint 频率
  • 54 天训练 466 次中断，>90% 有效训练时长依赖 async checkpoint
  • 论文未给出 async checkpoint 的具体阻塞时间数字

跨系统对比​

系统规模Stage 1 阻塞Stage 2 时长
MegaScale (NSDI 2024)10K+ GPU"数秒"异步后台
Nebula (Microsoft)128 A100~1 秒异步后台
PyTorch DCP async1856 H200~0.78 秒~67 秒（优化后）
Megatron-Core v0.7Nemotron-4 340B未披露42× 减少 vs torch.save（含 staging 压缩 + persist overlap 综合）
Llama 3 405B16K H100未披露（目标最小化）异步后台

@tbl-async-ckpt-industry-comparison 工业系统 async checkpoint Stage 1/Stage 2 时序对比（存储后端 / 写带宽细节见 04-章）

收敛规律（时序维度）：

  • Stage 1 阻塞时间在所有系统中收敛到亚秒至数秒，受 PCIe D2H 带宽主导（PCIe 4.0 x16 单向 ~32 GB/s，假设 80 GB GPU 显存全量做 snapshot 理论下限约 2.5 秒；实际工业系统 ~0.78–1 秒，是因为 checkpoint 只含 model state（参数 + optimizer），不含 activation / KV cache 等运行时状态，实际拷贝量远小于全显存）
  • Stage 2 时长差异主要源于异步 persist 设计本身（与训练 overlap 的程度），具体 IO 带宽与 sharding 格式影响见 10.4 分布式 Checkpoint 通信
  • Cached Plan（PyTorch DCP）通过复用 collective 规划减少 Stage 2 协调开销 6×

框架 Async API 对比​

核心问题：PyTorch DCP/TensorFlow/Megatron/DeepSpeed 四个框架的 Async Checkpoint API 在能力和接口上有何差异？

主流框架的 async checkpoint API 设计在 staging 机制、cross-rank 同步、写盘并行化几个维度有显著差异。

PyTorch DCP async_save​

PyTorch DCP[13] 提供原生 async_save API：

torch.distributed.checkpoint.async_save(
    state_dict,
    *,
    checkpoint_id=None,
    storage_writer=None,
    planner=None,
    async_checkpointer_type=AsyncCheckpointerType.THREAD,
    async_stager=None,
    use_collectives=True,
)  # Returns Future[Metadata]

Staging 隔离机制：分两阶段执行。第一阶段 stage_data() 在主线程同步完成，把 state_dict 复制到 CPU pinned memory（或 shared memory），返回与训练主循环隔离的副本——文档明确要求 staging 完成后对 module data 的任何更新不得反映在 staged state_dict 中。第二阶段 dcp.save 在独立线程/进程中对 staged 副本执行实际写盘。本质上是 先 copy-on-write（staging），再异步落盘。staging 是 stop-the-world 但时间远短于写盘。

AsyncStager 协议：

  • stage(state_dict) → Future | dict：返回隔离副本
  • synchronize_staging() → None：等待 staging 完成
  • should_synchronize_after_execute：控制 async_save 返回前是否等待 staging

DefaultStager：使用 pinned memory + ThreadPoolExecutor + non-blocking device memory copy（CUDA stream synchronization）。

DeepSpeed DataStates-LLM Lazy Snapshot​

DeepSpeed 集成的 DataStates-LLM (HPDC'24)[9] 采用惰性 copy-out而非 stop-the-world staging（DeepSpeed 集成用法见 tutorial[14]）：

核心洞察：model parameters 和 optimizer states 在 forward + backward 阶段保持 immutable，只有 optimizer update 阶段才修改参数。DataStates-LLM 利用此窗口在 forward/backward 期间异步执行 GPU → Host PCIe DMA 传输，在 update 阶段开始前确保传输完成。

实现要点：

  • 三阶段流水：(1) 递归解析 Python 对象，定位大张量指针；(2) 构建文件 offset header；(3) GPU→Host 与 Host→Storage 两路并行
  • ZeRO Stage-1：对多 shard（layer params + optimizer states）合并写入预分配 pinned buffer，避免逐 shard 分配开销
  • 限制：仅 CUDA，仅 ZeRO Stage-1，不支持 elastic checkpoint

性能：端到端训练加速 2.2×，checkpoint 创建速度提升最高 48×。

Megatron-LM --async-save​

Megatron-LM 通过 args.async_save 启用[15]，仅支持分布式格式（torch_dist、torch_dcp、fsdp_dtensor），不支持 legacy 格式。

实现机制：

  • generate_state_dict() 在 async/sync 分叉之前同步执行，用 all_gather_object 收集所有 DP rank 的 RNG states，确保 snapshot 时刻一致
  • schedule_async_save(async_save_request) 调度写盘，add_finalize_fn() 注册完成回调
  • checkpoint 删除通过 multiprocessing.get_context('fork') 后台 daemon process 执行
  • Sync 路径在 save 后立即执行 torch.distributed.barrier()，async 路径推迟 barrier 到回调

FastPersist — Optimizer-边界 Stop-the-world​

FastPersist (arXiv 2406.13768)[16] 设计思路与上述方案均不同 — 不做 volatile memory 中转，optimizer pass 完成后 helper thread 被唤醒，通过 DMA 把 GPU tensor 读入 page-locked CPU memory，直接写 NVMe SSD。写盘与下一 iteration forward/backward 并行。

  • 一致性保证：helper thread 在 optimizer 完成后才开始读 GPU tensors，确保读到的是完整 step T 状态
  • Cross-rank 同步：训练前 setup 阶段按 byte 粒度静态分配 checkpoint 任务给各 DP rank，运行时无跨 rank 通信开销
  • 性能：每 iteration 都做 checkpoint 时开销 < 5%，比 baseline 快最多 116×

框架横向对比​

框架Staging 机制Cross-rank 同步关键路径阻塞窗口
PyTorch DCPstage_data() 主线程 copy-to-pinned-memoryuse_collectives=True 协调各 rank staging 进度staging 时间（~0.78s 实测）
DataStates-LLMLazy copy-out（forward/backward immutability 窗口）两阶段 commit：node-local 验证 → global 验证，与下一 iteration overlap无显式 stop-the-world，依赖 update 前完成
Megatron-LMgenerate_state_dict() 同步收集 + 异步写盘all_gather_object 收集 RNG statessnapshot 收集时间
FastPersistoptimizer 完成后 helper thread DMA 直接写 NVMe训练前静态分区，运行时无通信DMA 时间（在 optimizer 边界）

@tbl-async-ckpt-framework-async 框架 async checkpoint 横向对比

Snapshot 一致性：Atomicity 与 Cross-Rank Sync​

核心问题：Snapshot 的跨 rank 一致性和原子性如何保证？不一致会有什么后果？

Snapshot atomicity 的几种实现​

没有框架实现真正的 zero-overhead snapshot — 所有方案都在某个边界做一次同步，区别在于该同步发生的位置和持续时长：

方案同步发生位置持续时长后续行为
Stop-the-world（CheckFreq CPU 模式）关键路径（snapshot 在 D2H 期间）GPU→CPU 传输时间写盘完全异步
Pinned-memory staging（PyTorch DCP）关键路径（stage_data 主线程）pinned memory copy 时间写盘完全异步
Lazy copy-out（DataStates-LLM）forward/backward immutability 窗口无显式 stop-the-world依赖 update 前完成
Optimizer 边界 stop-the-world（FastPersist）optimizer 完成后DMA 直接到 NVMe写盘与下一 iter 并行
GPU 内复制 + 后台 D2H（CheckFreq GPU 模式）无（GPU 内存内复制几乎瞬时）几乎为零D2H + 写盘均异步

@tbl-async-ckpt-atomicity Snapshot atomicity 实现机制对比

Copy-on-write (COW) 在操作系统 mmap 层面是理论可行的，但主流训练框架均采用显式 copy-out 到独立 buffer的方式，原因：

  • 更可控 — staging 完成时刻明确，避免 COW 缺页时不可预测的延迟尖峰
  • 利于 GPU pinned memory 管理（COW 需要复杂的 pinning + un-pinning 协调）
  • 避免训练过程中参数被频繁修改触发 COW 副本拷贝累计的隐性开销

Cross-rank snapshot 同步​

分布式训练所有 rank 必须 snapshot 同一 training step 的状态，否则 resume 时各 rank 参数不一致导致训练发散。各框架的同步策略：

  • PyTorch DCP：use_collectives=True 通过 collective 协调各 rank 的 staging 和写盘进度
  • Megatron-LM：all_gather_object 收集 RNG states，snapshot 在 barrier 前完成
  • DataStates-LLM：两阶段 commit（node-local 验证 → global 验证），与下一 iteration overlap
  • FastPersist：训练前静态分区，运行时各 rank 独立写，无跨 rank 运行时通信

ZeRO State Race Conditions​

ZeRO 把 optimizer states 分片到各 DP rank，async checkpoint 时需确保所有 rank 完成各自 shard 的 snapshot 后再进入 update 阶段，否则某些 rank 的 shard 包含 step T+1 的部分更新。解法：

  • DataStates-LLM：通过"update 开始前完成 GPU→Host 传输"保证此约束
  • PyTorch DCP：staging 完成后才返回 Future，由调用方控制 update 时序

增量 Checkpoint 与 Async 配合​

核心问题：增量 Checkpoint 如何与 Async 机制配合、进一步减少每轮数据量和 Persist 时间？

主流框架未原生支持增量（delta）checkpoint 与 async 的联合优化。已知工作：

  • DataStates-LLM (HPDC'24)[9] 明确将 differential checkpointing 列为 future work
  • DeltaCheck（SoCC'26 under review）针对 LLM post-training 提出 stability-aware 增量持久化，减少 intra-state 和 inter-state 冗余，兼容 DeepSpeed / Megatron-style 工作流，但 async 机制尚未公开详细设计
  • IncrCP（VLDB'25）研究增量 checkpoint 的分解与调度，指出构建 checkpoint 阻塞训练是核心问题，async 是改进方向

核心难点：增量 checkpoint 需要 diff 前后两个 step 的 state，而 async 时上一个 checkpoint 可能尚未写完，diff 的基准状态不确定。两种解决思路：

  1. 保留 CPU-side 缓存作为 diff 基准（不等待写盘完成）— 内存开销翻倍
  2. 以写盘完成时间点为 diff 基准（放弃部分异步窗口）— 失去部分加速

目前无生产级开源实现。

开放问题​

核心问题：异步 Checkpoint 尚未解决的关键开放问题有哪些？

  • GPU snapshot 的 PCIe 带宽下限：Stage 1 阻塞时间在工业系统中收敛到亚秒至数秒，本质受 PCIe D2H 单向带宽限制（PCIe 4.0 x16 ~32 GB/s）。PCIe 5.0（~64 GB/s 单向）和 NVLink-C2C / CXL 普及后，Stage 1 时间能压到多低？是否会被 CPU 序列化反而成为瓶颈？
  • Cached Plan 在动态拓扑下的失效：PyTorch DCP Cached Plan 假设训练期间 checkpoint 结构不变。弹性训练 / 动态 expert parallel（MoE 自适应路由）场景下，参数布局可能在 step 间变化，Cached Plan 复用机制需要相应扩展。
  • Lazy copy-out 在 PP 训练下的可行性：DataStates-LLM 利用 forward/backward 期间参数 immutability 窗口做异步 copy-out。Pipeline parallel 训练中 forward/backward 在不同 stage 时序交错，immutability 窗口的定义和保证机制需要重新设计。
  • In-memory checkpoint 的内存预算上限：Gemini 假设每节点 CPU 内存 512 GB–1 TB 可容纳 2× 模型状态。下一代 1T+ 参数模型 + ZeRO-3 微调场景下，每节点状态可能超过 CPU 内存容量。Mixed Placement 是否需要引入 GPU HBM 副本作为更高一层？
  • 增量 checkpoint 与 async 的生产级实现：当前 delta / incremental checkpoint 与 async 配合都是研究原型，缺乏与 PyTorch DCP / Megatron / DeepSpeed 三大主流框架的深度集成实现。
  • 跨故障域定义的工程实现：Gemini 的 Mixed Placement Policy 依赖"跨故障域副本"概念，但生产环境如何在调度器层面声明、维护、验证故障域映射，论文未深入。

Takeaway​

知识点核心结论
两阶段流水线snapshot (GPU → CPU) + persist （CPU → 存储），persist 完全异步与训练并行
同步阻塞代价GPT-3 175B 10-20 min; OPT 175B 21 min; Llama 3 405B ~13.5 min
CheckFreq 频率算法$k = \max((T_c + T_s - T_o)/T_i, \lceil T_o / (p T_i) \rceil)$，维持开销 ≤ p （典型 5%）
GPU vs CPU 模式GPU 内存够走 G2G copy 不阻塞；否则走 D2H 与 host 加载竞争 stream，自动降频
Gemini 三层放置本地内存 / 远端内存 / 磁盘，in-memory recovery 95s vs 磁盘 15+ min
DataStates Lazy snapshot利用 forward/backward immutability 窗口后台 D2H，端到端 2.2× 加速
FastPersistOptimizer 边界 stop-the-world + 直接写 NVMe + 静态分区，每 iter 开销 < 5%
MegaScale 两阶段Stage 1 数秒 + Stage 2 异步 HDFS, 10K+ GPU 实测
PyTorch DCP Cached Plancheckpoint 结构不变时复用 collective 规划，仅传增量，6× 加速
ZeRO state raceZeRO-1/2/3 + async 组合下 optimizer 分片状态可能被部分更新，需框架特殊处理

参考资料​

  1. Zhang et al., OPT: Open Pre-trained Transformer Language Models, arXiv:2205.01068, 2022. https://arxiv.org/abs/2205.01068 OPT-175B 训练日志：2 个月窗口约 35 次手动重启（Gemini 论文引用同期观察约 110 次故障），单次磁盘 checkpoint 约 21 分钟。
  2. Dubey et al., The Llama 3 Herd of Models, Meta tech report / arXiv:2407.21783, 2024. https://arxiv.org/abs/2407.21783 Tectonic 2 TB/s 存储下的 async checkpoint 实践，54 天 466 次中断。
  3. 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 三阶段架构、97 GB 1s 实测、已弃用声明。
  4. Maeng et al., CPR: Understanding and Improving Failure Tolerant Training for Deep Learning Recommendation with Partial Recovery, MLSys 2021. https://cs.stanford.edu/people/trippel/pubs/cpr-mlsys-21.pdf 量化 checkpoint 开销占总训练时间平均 12%、极端 43%。
  5. Wang Z. et al., Gemini: Fast Failure Recovery in Distributed Training with In-Memory Checkpoints, SOSP 2023. https://dl.acm.org/doi/10.1145/3600006.3613145 In-memory checkpoint + Mixed Placement Policy + Interleaved Scheduling，恢复 95s vs 磁盘 15+ min。
  6. Mohan et al., CheckFreq: Frequent, Fine-Grained DNN Checkpointing, USENIX FAST 2021. https://www.usenix.org/conference/fast21/presentation/mohan 两阶段 snapshot/persist + profiling-driven 自适应频率，DNN async checkpoint 奠基。
  7. Nicolae et al., DeepFreeze: Towards Scalable Asynchronous Checkpointing, IEEE/ACM CCGrid 2020. https://doi.org/10.1109/CCGrid49817.2020.00-76 CPU HPC 集群 async 序列化 + I/O 流水线，CheckFreq 的前身。
  8. Daly, A higher order estimate of the optimum checkpoint interval, Future Generation Computer Systems 2006. https://doi.org/10.1016/j.future.2004.11.016 HPC 经典理论：k* = √(2·C·MTBF) 最优 checkpoint 间隔。
  9. Maurya et al., DataStates-LLM: Lazy Asynchronous Checkpointing for Large Language Models, HPDC 2024 (arXiv:2406.10707). https://arxiv.org/abs/2406.10707 Lazy non-blocking snapshot 利用 forward/backward immutability 窗口，端到端 2.2× 加速。
  10. 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（Stage 1 数秒 + Stage 2 异步 HDFS），10K+ GPU 实测。
  11. NVIDIA, Train Generative AI Models More Efficiently with New Megatron-Core Functionalities, 2024. https://developer.nvidia.com/blog/train-generative-ai-models-more-efficiently-with-new-nvidia-megatron-core-functionalities/ Nemotron-4 340B 42×、50×、26× checkpoint 开销减少。
  12. PyTorch Team, 6x Faster Async Checkpointing in PyTorch, 2025. https://pytorch.org/blog/6x-faster-async-checkpointing/ Llama3-70B 1856 H200 实测 Cached Plan + 进程模式优化 6×。
  13. PyTorch 官方，Distributed Checkpoint async_save API (PyTorch 2.x). https://docs.pytorch.org/docs/main/distributed.checkpoint.html async_save API、AsyncStager 协议、DefaultStager 实现。
  14. Microsoft, DeepSpeed DataStates Async Checkpointing Tutorial. https://www.deepspeed.ai/tutorials/datastates-async-checkpointing/ DataStates-LLM 集成 DeepSpeed 的使用方式。
  15. NVIDIA, Megatron-LM checkpointing.py async-save 实现. https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/checkpointing.py --async-save 实现、AsyncCallsQueue、schedule_async_save。
  16. Wang G. et al., FastPersist: Accelerating Model Checkpointing in Deep Learning, arXiv:2406.13768, 2024. https://arxiv.org/abs/2406.13768 Optimizer 边界 stop-the-world + 直接写 NVMe + 静态分区，每 iter checkpoint 开销 < 5%。

项目内交叉引用​

Limitations​

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

  • 本章不涵盖跨节点 sharding / Save-Load 协议 / IO 带宽细节（属于 10.4 分布式 Checkpoint 通信）
  • CheckFreq 论文发表于 2021，data iterator 不可恢复的具体框架限制（PyTorch/MXNet/DALI）可能已在新版本部分解决；本章引用论文原始描述
  • Gemini 论文报告基于 256 GPU 实验，10K+ GPU 规模下 Mixed Placement Policy 的副本管理复杂度与故障域定义未公开
  • 工业实测数字多为单次报告，缺少跨厂商同条件对比；MegaScale 的 Stage 1 "数秒"未给精确值
  • 框架 async API 横向对比基于 2025-2026 公开文档，PyTorch DCP / DeepSpeed / Megatron 持续演进，部分 API 可能变化
  • 增量 checkpoint 与 async 配合的部分（DeltaCheck / IncrCP）基于 under review / 早期工作描述，落地情况待验证
  • ZeRO state race condition 仅在 ZeRO-1/2/3 + async 组合下出现，详细 race condition 类别与覆盖范围以各框架最新实现为准
  • 对等镜像复制（peer-mirror replication）的 in-memory checkpoint 在学术界有相关工作（如部分文献中提及的 REFT），但具体论文的 venue 与作者归属在不同 second-hand 引用中存在分歧，本文未能确认权威来源，故不在主文中作具体对比