Checkpoint
3秒看懂
Checkpoint（检查点）在深度学习中有两种紧密相关但目标不同的含义：
- 训练快照（Training Snapshot）：周期性将模型参数、优化器状态、学习率调度器和随机数状态等完整训练状态保存至持久化存储。训练中断或硬件故障后可从最近快照精确恢复，也可用于实验分支回滚。
- 激活检查点（Gradient / Activation Checkpointing）：一种以计算时间换取显存空间的优化技术。前向传播时仅保留部分关键激活，丢弃中间激活，反向传播时根据保留的输入动态重算被丢弃的激活，大幅降低峰值显存占用，使得在有限的 GPU 显存上训练更大模型或使用更大的 micro batch size 成为可能。
3分钟产业解释
在大模型竞争白热化的背景下，Checkpoint 被赋予双重的工程与战略价值：
- 可靠性基石：数千张 GPU 集群连续训练数周至数月，硬件故障、节点掉线、进程崩溃近乎是必然事件。训练快照是唯一有效的“保险”——每 N 步保存一次完整状态，使价值数十万甚至数百万美元的算力投入不至于因一次机架掉电而清零。同时，云上训练大量采用可抢占实例，定期快照让“中断-恢复”成为标准化运营流水线。
- 显存杠杆：GPU 显存容量增速（近年来约 2 年翻倍）远落后于模型参数膨胀速度（约每年 5–10 倍）。激活检查点通过将显存压力转换为可控的额外计算（通常 20%–40% 的单步时间增加），使训练 GPT-4、Llama 3、Gemma 等数百亿至数千亿参数模型成为可能。它已从一个选题性技巧演变为所有主流训练栈的“默认开启”基础能力。
产业实践中，两类 Checkpoint 往往同时配置：训练快照保障训练的物理连续性，激活检查点保障显存边界内的数学可行性。PyTorch、TensorFlow 等框架与 DeepSpeed、Megatron-LM 等分布式训练库已将它们深度融合，用户仅需配置若干参数即可启用。云厂商（AWS、Azure、GCP）则围绕检查点提供高吞吐对象存储、并行文件系统和托管训练服务，进一步降低门槛。
技术原理
训练快照检查点
保存内容
一个能完整恢复训练的快照通常包含：
- 模型参数（权重、偏置，即
state_dict）； - 优化器状态（如 Adam 的一阶动量与二阶动量，其体积可达模型参数的 2 倍）；
- 学习率调度器状态；
- 随机数生成器状态（保证数据加载与 dropout 模式的可复现性）；
- 训练元信息（步数、epoch、数据迭代器位置或数据集的确定性随机种子）。
保存时机与策略
- 按步间隔保存：每 100 或 1000 步保存一次，适配上万个训练步数。
- 按时间间隔保存：适合训练时长高度不确定的场景。
- 滚动保留：维持最近 K 个检查点，平衡存储开销与恢复精度；旧检查点按策略删除或移动到冷存储。
恢复机制
从快照重新构建模型与优化器对象，加载上述全部状态；若采用确定性数据加载（如固定的全局随机种子 + 按步数同步的数据读取器），训练可从断点精确延续，如同从未中断。
分布式场景的一致性保障
- 数据并行：每个 rank 拥有完整模型副本，通常仅 rank 0 执行保存（简单但存在单点写入带宽瓶颈），或各 rank 写入分片状态。
- 模型并行（张量并行、流水线并行）：模型状态在多个设备上切分存放。保存时需并行收集所有分片，并按协商一致的分片方案写入；恢复时反向分发。Megatron-LM 和 DeepSpeed 均实现分布式检查点读写器，通过 NCCL 集合通讯或并行 I/O 完成高效聚合和分发。大集群中常采用多节点并发写入分布式文件系统，以避免单节点 I/O 成为瓶颈。
激活检查点（Gradient Checkpointing）
问题来源
经典反向传播要求前向过程中产生的所有激活张量驻留于显存，直至对应反向计算完成。对于 Transformer 模型，激活显存占用与层数、序列长度、隐藏维度和 batch size 成正比，极易成为瓶颈。
核心原理
将网络分割为若干“检查点段”（如一个 Transformer Block）。前向传播时，仅在段边界保留输入张量，段内所有中间激活计算后立即释放。当反向传播到达该段时，利用保存的段输入重新执行一次段的前向计算，重新获得段内激活后完成梯度计算。这相当于用额外一段完整的前向计算换取段内全部激活的显存释放。
在典型配置下，激活显存可由与层数线性相关降低至与段数线性相关（通常为常数或平方根级别），显存节省显著。额外计算量约为一次完整前向，总训练时间增加大致在 20%–40%，具体取决于前向计算与反向计算的耗时比例以及重算段落的选择（参考：Chen 等 2016 年原论文及后续社区实践经验总结）。
选择性重计算
为控制额外开销，实现中通常只对显存压力最大的部分网络段做检查点（如 Transformer 层本身），而对计算量大但显存消耗小的操作（如注意力矩阵乘法）保留完整激活。PyTorch 的 torch.utils.checkpoint 允许用户单独标记检查点分段；DeepSpeed 的 Activation Checkpointing 支持自动或手选层配置，配合模型并行在流水线气泡中隐藏部分重计算开销。
关键参数
-
检查点体积（Checkpoint Size）
单次保存的存储量，受模型参数量、优化器状态和精度影响。举例：一个 175B 参数的 GPT-3 类模型，若采用 Adam 优化器与 FP16 混合精度训练（参数 FP16，优化器状态 FP32），单一全局检查点体积约为 2.8 TB（参数 350 GB + 优化器约 700 GB × 2 + 其他状态，参考社区估算；具体数字因并行策略和分片方式变化）。万亿参数模型单检查点体积可达数十 TB 级别。 -
保存/恢复吞吐（Save/Restore Throughput）
持续写入/读取检查点的速度，以 GB/s 或 TB/min 计。例如，在 NVIDIA DGX SuperPOD 中，通过多节点并行写入 Lustre 文件系统，实测聚合写入吞吐可超过 100 GB/s（NVIDIA 白皮书，2023 年）。该指标直接决定检查点保存间隔下限和恢复时间。 -
恢复时间目标（Recovery Time Objective, RTO）
从故障发生到训练完全恢复至故障前步数所需墙钟时间，包含读取检查点、重建分布式状态、重放数据等。万卡集群中，单个大模型检查点的加载可能需要数分钟至数十分钟。RTO 越短，有效训练时间占比越高。 -
显存节省率（Memory Saving Ratio）
激活检查点开启后峰值显存的下降比例，通常以“× 倍”表达。例如，开启 Transformer 层粒度的激活检查点可使激活显存降低约 80%–90%（具体取决于隐藏维度与序列长度组合）。PyTorch 文档中记载，开启检查点后模型所需显存可缩减至原来的 1/√L 或更低。 -
额外计算开销（Recomputation Overhead）
激活重计算导致的总训练时间增加比例。典型区间为 15%–40%。通过选择性检查点和与流水线调度协同，DeepSpeed 等框架可将开销控制在 20% 左右（DeepSpeed 文档，2024 年版本）。 -
检查点保真度（Checkpoint Fidelity）
确保保存过程中的状态是确定的、可完全复现的，分布式环境下无竞态导致的状态不一致。通常通过全局barrier同步和确定的集合通信顺序来保障。 -
存储成本占比
大规模训练中，检查点占据的存储空间可能远大于数据集本身。例如，训练一个千亿参数模型期间保留 10 个滚动检查点，总存储需求可达数十 TB 到百 TB 级。对云存储费用影响显著，也成为架构设计需考虑的变量。
上述体积与吞吐数值因并行策略（数据并行、张量并行、流水线并行、ZeRO 分片等级）、硬件（NVMe、InfiniBand 网卡、GPU 显存）和数据类型而异，实际项目应以实测为准。
技术路线
演化脉络
- 2016 年以前：训练快照作为基础容错手段，主要依靠框架提供简单的
save/loadAPI，解决 GPU 集群故障率增高下的训练可恢复性。 - 2016 年：Chen 等人发表《Training Deep Nets with Sublinear Memory Cost》，首次系统化提出激活重计算，论证可将显存消耗由线性降至亚线性，为极深网络的训练打开新可能。
- 2017–2019 年：PyTorch 和 TensorFlow 内置激活检查点接口；Horovod 等推出更高效的分布式保存；分布式训练社区开始形成“按步保存 + 滚动保留”的最佳实践。
- 2020–2022 年：大模型元年。DeepSpeed 推出 ZeRO 状态分区，并与激活检查点深度耦合，大幅降低巨模型训练显存门槛；Megatron-LM 实现张量并行/流水线并行下的分布式检查点读写器；各大云厂商存储网关支持并行写入。
- 2023 年至今：针对 MoE（混合专家）模型、动态网络的选择性检查点策略成为热点；异步检查点写入技术出现（在训练主循环不阻塞的情况下将状态复制到 CPU 内存后异步刷盘），减少保存暂停时间；万亿参数模型推动分层次、分片压缩的检查点方案研究。
主流方案对比
| 维度 | 训练快照检查点 | 激活检查点 |
|---|---|---|
| 主要目标 | 容错恢复、实验回溯、模型交付 | 降低训练峰值显存 |
| 保存内容 | 模型/优化器/调度器/随机状态等全套 | 仅段边界输入张量 |
| 存储/IO 开销 | 大（每检查点数 GB~数十 TB） | 极小（少量张量，通常 KB~MB） |
| 计算开销 | 可忽略（仅序列化和 I/O 耗时） | 显著（额外前向重算，总时间增加 15%–40%） |
| 启用方式 | 按步/按时触发 | 始终启用，作用于指定模块 |
| 能否独立恢复训练 | 可精确恢复 | 不能，仅影响激活管理 |
| 代表性实现 | PyTorch torch.save、Lightning Checkpoint | PyTorch checkpoint_sequential、DeepSpeed Activation Checkpointing |
| 分布式兼容性 | 需要并行保存与恢复逻辑 | 与模型并行协同设计，需考虑重算与通信重叠 |
上游
Checkpoint 的效能高度依赖底层硬件与系统栈：
- GPU 显存与算力：激活检查点的重计算增加了计算总量，对 GPU 浮点算力提出更高要求；同时，检查点段的选择需要精确匹配显存带宽和容量，不同 GPU 微架构（A100/H100/B200）上的最优策略可能不同。
- 存储系统：训练快照要求高突发顺序写入带宽和低延迟元数据操作。常用方案包括高端并行文件系统（如 Lustre、GPFS）、NVMe over Fabrics（NVMe-oF）和云原生的高吞吐对象存储（AWS S3 Express One Zone、Azure Blob Hot）。存储系统的写入缓冲能力、读镜像分发会直接决定 RTO。实践中，多节点并发写入时小文件随机元数据操作可能成为瓶颈，推动面向检查点优化的异步写入聚合中间件的发展（如 Facebook 内部使用的 Tectonic 文件系统定制修改，据 Meta 工程博客）。
- 网络互连：分布式检查点需要多节点并行写入，集合通信（NCCL）和并行 I/O 生成的网络负载非常可观。高带宽低延迟网络（如 InfiniBand HDR/NDR、RoCE v2）是维持检查点吞吐的必备条件。万卡集群中，检查点保存会产生瞬时的全交换数据流，网络拥塞控制方案影响保存完成时间。
- CPU 内存与系统总线：异步检查点方案需先将状态拷贝到 CPU 内存，而后异步写入磁盘，因此 CPU-DRAM 容量和 PCIe/NVLink 带宽成为缓冲瓶颈。极大规模下可能需使用 CPU 内存池或 NVDIMM 作为写缓冲。
下游
- 大规模预训练：千卡/万卡集群上持续数月的训练，完全依赖定期检查点实现“无限长时间运行”。Llama 3 的训练就使用了密集的检查点策略以确保在多重故障下依然可持续（Meta 公开 blog，2024 年）。
- 故障恢复与弹性训练：云上可抢占实例、电网波动、设备维护等场景中，检查点实现分钟级恢复，使训练任务具备云原生的弹性；配合 Kubernetes Job 或 SLURM 的自动重启，实现训练自动化。
- 实验管理：从不同检查点启动多分支实验（如微调、RLHF、数据配比消融），大幅加速研发迭代。Hugging Face Hub 等模型仓库中，多数模型以检查点形态分发，供社区自由续训或微调。
- 模型交付：最终检查点常作为模型发行的标准格式。提供者可能同时发布原始训练快照（含优化器状态）和仅推理权重的轻量检查点，用于下游部署。
- 知识蒸馏与剪枝：蒸馏、量化感知训练、结构化剪枝等工作通常从特定检查点起始，以保证可比性与复现性。
受益公司
- Meta（PyTorch）：PyTorch 核心维护了
torch.utils.checkpoint及相关分布式保存 API，其技术路线极大影响了业界检查点实践。Meta 自身在大规模训练中积累的检查点方案通过开源反哺社区，巩固了框架地位。 - 微软（DeepSpeed / Azure）：DeepSpeed 将 ZeRO 状态分片与激活检查点深度融合，降低了超大模型训练门槛，并天然引导用户使用 Azure 云服务实现协同优化。DeepSpeed 的检查点方案已成为多数 GPT 类模型训练栈的标配组件。
- NVIDIA（Megatron-LM / NeMo / DGX 系统）：NVIDIA 提供从芯片（H100/B200）到系统（DGX SuperPOD）到软件（Megatron-LM、NeMo）的全栈检查点优化，尤其是并行文件系统与网络配置的紧密整合，使其在高端训练集群市场占据主导。
- Hugging Face：通过
transformers库统一了预训练模型的检查点保存/加载接口，并运营业界最大的模型检查点托管平台，掌握重要的分发节点地位。 - 云基础设施厂商（AWS、GCP、CoreWeave 等）：提供高性能存储与网络，决定检查点 I/O 上限；其托管训练服务（如 Amazon SageMaker、GCP Vertex AI）均内建检查点策略，对用户透明。
- 新兴训练框架企业（ColossalAI、MosaicML/Databricks 等）：通过创新的异步检查点、梯度累积感知检查点等差异化特性争夺市场关注，部分已被大厂整合（如 MosaicML 被 Databricks 收购）。
以上仅列举公开信息中与 Checkpoint 技术栈直接相关的公司，不构成任何投资评价。
市场规模
直接统计缺失
截至 2024 年 7 月，尚无第三方机构发布以“Checkpoint 软件/服务”为独立品类的市场规模报告。Checkpoint 属于深度学习训练基础设施中的一项横向能力，其商业价值隐含在框架、云服务和存储硬件中，公开资料未见拆分数据。
参考关联市场
- 整体 AI 基础设施市场：根据 IDC 2024 年初发布的《Worldwide AI Infrastructure Tracker》，2023 年全球 AI 服务器与存储支出约 326 亿美元，预计 2027 年将超过 520 亿美元。与检查点直接相关的高吞吐并行文件系统和对象存储在该市场中占据重要份额，但具体比例未单独披露。
- HPC 存储市场：Hyperion Research 2023 年报告指出，全球 HPC 存储（含 AI 训练）市场约 75 亿美元，特别是支持高速突发写入的并行文件系统与 NVMe-oF 方案增速超过 20%。
- 云存储消费：大规模训练项目中，检查点存储费用可占训练总云成本的 5%–15%（根据各云厂商公开案例估算，如 Azure 博客 2023 年某客户案例）。随着模型参数增长，检查点相关存储消费的绝对额与相对比例预计将继续上升。
以上关联市场数据取自第三方公开报告，口径与具体数字请查阅原报告。
玩家对比
| 解决方案 | 训练快照能力 | 激活检查点策略 | 分布式优化亮点 | 社区/生态 |
|---|---|---|---|---|
| PyTorch 原生 | torch.save、支持 FSDP 分片保存 | torch.utils.checkpoint 选择性包装 | 与 DDP/FSDP 深度整合，但高级并行需手配 | 最广泛的模型实现依托，生态最大 |
| DeepSpeed | 提供 ZeRO 1/2/3 的分片检查点，支持异步写入 | 自动或手动按层配置，可与流水线气泡重叠 | ZeRO 状态下仅保存/加载分片，节省内存与 I/O；CPU/NVMe 卸载 | 大量开源大模型训练采用，由微软维护 |
| Megatron-LM（NVIDIA NeMo） | 张量/流水线并行下分布式检查点读写器 | 提供与模型并行模式绑定的激活重算 | 分布式检查点写入器高度优化，针对 NVIDIA 硬件/网络调优 | NVIDIA 内部预训练首选，社区有部分复用 |
| Hugging Face Transformers | 统一的 save_pretrained/from_pretrained | 依赖 PyTorch 或 DeepSpeed 插件转发 | 主要关注模型交付与微调场景，对分布式写入支持有限 | 最大模型分发平台，下游应用丰富 |
| ColossalAI | 支持多维并行分片检查点，异步写入 | 提供灵活的激活检查点定制 | 将检查点与自动并行化结合，降低用户感知 | 新兴，社区增速较快 |
| JAX/Flax 生态 | 基于 Orbax 提供异步检查点 | jax.checkpoint（原 jax.remat） | 支持多主机多设备检查点，与 TPU 训练紧密配合 | 主要由 Google 及 TPU 用户使用 |
功能状态均基于对应项目截至 2024 年 7 月公开文档与代码仓库信息。
风险
- 技术瓶颈风险：万亿参数级模型的单检查点体积可达数十 TB，当前最先进的并行文件系统在未作专门优化的情况下也可能触达写吞吐天花板；恢复时间可能长至小时级别，影响有效训练时间。虽然研究层面不断探索检查点压缩、增量保存等方法，但尚未成为工业界主流，存在过渡期瓶颈。
- 异构算力适配风险：GPU + AI 芯片（如 TPU、NPU）混合训练或跨集群训练时，不同芯片架构之间的检查点格式、算子确定性重算行为不一致，可能产生状态恢复错误或额外转换开销，增加工程复杂性。
- 安全与合规风险：检查点含完整可复现的训练状态，包括模型权重、优化器状态、数据采样状态等。一旦泄露，攻击者可重建模型甚至推断训练数据分布。同时，某些行业（如金融、医疗）的合规要求可能对检查点的保留和传输施加限制，涉及数据驻留和加密传输等成本。
- 成熟度与创新空间：检查点技术整体已趋于成熟，框架层面的颠覆式创新空间有限，这可能导致部分创业公司难以实现差异化壁垒。然而与 MoE、动态架构、在线学习等新范式的结合仍存在未充分解决的技术点，成为潜在的不确定性来源。
- 存储成本失控风险：滚动保留多个检查点产生的存储费用随模型规模与项目数量指数增长，如果缺乏精细化的生命周期管理策略（如自动冷热分层、增量差异保存），存储成本可能侵蚀训练总预算，尤其对于长期运行的自监督训练项目。
误读纠偏
-
误读：“检查点就是模型最终权重文件”
纠正：仅保存模型权重（如model.pt）不包含优化器状态等，无法实现精确恢复训练。如果从纯权重文件继续训练，优化器将从零初始化，改变训练动力学，导致收敛路径漂移。用于恢复训练的检查点必须保存完整训练状态。 -
误读：“激活检查点一定会加速训练”
纠正：激活检查点本质是显存优化技术，通过释放显存允许增大 batch size，可能间接提高吞吐（更高 GPU 利用率）。但若显存本非瓶颈，重计算只会单纯延长单步时间，降低训练速度。是否“加速”取决于显存约束是否为系统主要瓶颈。 -
误读：“激活检查点适用于所有模型”
纠正：该技术要求模型可分段且重算具有确定性。涉及不可逆随机操作（如某些动态 dropout 或依赖全局状态的层）时，重算可能产生不同结果，导致梯度误差。通常需配合固定随机种子和可控的随机遮挡。 -
误读：“保存越频繁越好”
纠正：过于频繁的检查点保存会消耗大量 I/O 带宽和存储空间，拖累训练主循环。在大规模集群中，一次全量检查点保存的耗时可达到数分钟甚至更高，若保存频率超过合理阈值，有效训练时间占比会显著下降。需根据模型规模、集群故障率（MTBF）和 RTO 综合确定最优间隔。
最新事件
- 2024 年 4 月：Meta 发布 Llama 3 模型，其训练流程中大量运用 PyTorch FSDP + 检查点技术，并在公开工程博客中提及长期训练中的故障恢复策略与检查点频率配置经验（Meta AI Blog，2024 年 4 月）。
- 2024 年 6 月：DeepSpeed 发布 0.14.x 版本，增强异步检查点（Asynchronous Checkpointing）支持，允许在训练步不阻塞的情况下完成状态持久化，据项目发布说明，可将检查点 IO 开销隐藏超过 90%。
- 2024 年上半年：PyTorch 2.3 版本改进了
torch.utils.checkpoint与compile的兼容性，并增加了对选择性激活重计算的更好支持；同时社区提出用于 FSDP 的更高效的分片检查点合并方案。 - 2024 年：NVIDIA 在其 NeMo 框架中集成面向 MoE 架构的优化检查点方案，可减少 expert 参数分片的保存冗余；多家云厂商发布增强型并行文件系统实例，以应对万卡集群突发写入需求（如 AWS 推出 FSx for Lustre 高吞吐实例）。
- 公开报告：部分研究机构开始发布千亿-万亿参数模型训练的检查点 IO 基准测试（如 MLPerf 基准的存储相关指标），推动行业内标准化的比较。
以上事件均来源于相应项目官方公告、技术博客或公开报告，截至 2024 年 7 月。
跟踪指标
若需持续把握 Checkpoint 领域的技术与产业动态，建议关注以下量化指标与信号：
- 最大公开单检查点体积：观测公开报告或论文中使用的检查点大小，反映参数规模增速对存储系统的冲击。
- 主流框架检查点写入吞吐基准：例如在 MLPerf、社区基准中出现的 GB/s 写入速度，反映存储栈与框架协同的进步。
- 典型 RTO 改善曲线：关注云厂商与系统厂商公布的恢复时间目标，衡量弹性训练能力的提升。
- 激活检查点节省率与额外开销的最新数据：不同模型架构下的实测值，可判断技术效益是否饱和。
- 新并行策略（如 MoE、专家并行）下的检查点方案成熟度：是否出现针对新架构的专用检查点库或论文。
- 异步与增量检查点的生产采用度：主流训练框架 Release Note 中是否添加相关特性，以及开源大模型训练代码中是否启用。
- 存储硬件路线图：PCIe 6.0、CXL 内存共享、NVIDIA Spectrum-X 等新型互连与存储介质对检查点 I/O 带宽的潜在推动作用。
- 训练故障率与集群规模关系：Meta、Google 等发布的故障统计数据，可评估检查点频率设置的经济区间。
信源
- PyTorch 官方文档：
torch.save与torch.utils.checkpoint章节 - Chen, Tianqi, et al. “Training Deep Nets with Sublinear Memory Cost.” arXiv 1604.06174, 2016.
- DeepSpeed 文档：Activation Checkpointing 与 ZeRO Checkpoint 部分
- Megatron-LM GitHub 仓库：分布式检查点设计与实现文档
- Hugging Face Transformers 文档：模型保存与加载
- NVIDIA NeMo 框架文档：检查点与恢复
- Meta AI Blog, “How Meta Trains Large Language Models at Scale,” 2024.
- IDC, “Worldwide AI Infrastructure Tracker, 2023–2027,” 2024.
- Hyperion Research, “HPC Storage Market Update,” 2023.
- 各云厂商官方博客与技术白皮书（AWS、Azure、GCP 训练存储最佳实践）
- MLPerf Storage Benchmark 相关白皮书与公开结果
注：所有技术参数、性能数字请以对应框架和硬件的官方文档及最新论文为准；市场数字引自第三方研究机构，仅用于说明行业规模量级，不构成任何投资评价。