检查点分片和可恢复训练
将模型的训练状态保存为片段（分片）的技术，以便可以保存和重新加载大型模型，而不会受到内存或磁盘限制的影响，因此崩溃的运行可以准确地从中断的地方继续。
概述
对于任何在多个 GPU 上运行数天或数周的训练作业来说都是至关重要的。
深入探讨
训练检查点是恢复所需的所有内容的快照：模型权重、优化器状态、学习率计划、数据加载器的位置和随机数生成器种子。对于大型模型，此快照可能有数百 GB，对于单个文件或单个计算机的内存来说太大了。检查点分片将该快照拆分为多个文件和多个等级，因此每个 GPU 仅并行写入自己的切片。然后，可恢复训练会重新加载这些分片并精确恢复完整状态。如果没有它，在第 200 小时崩溃的多周运行将不得不从头开始。 PyTorch 分布式检查点、DeepSpeed 和 Hugging Face Hub 的分片安全张量格式等框架可以实现此例程。
技术洞察
分片之所以有效，是因为分布式训练已经跨等级划分了权重和优化器状态（通过数据、张量或零并行）。每个等级仅序列化其分区，通常序列化为安全张量之类的格式，允许延迟、内存映射加载。索引文件将参数名称映射到分片文件。为了确定性地恢复，系统还保留 RNG 状态、优化器步数和确切的数据加载器偏移量，因此重新运行会重现相同的批次序列。
战略影响
成本与预算
多年来，架构决策决定着性能和运营成本。
更清晰的判决
技术教育帮助团队选择正确的堆栈，而不仅仅是最新的堆栈。
质量控制
更好的工程选择可以减少生产中的可靠性事故。
检查点分片和可恢复训练的未来
检查点正在从周期性的停止世界事件转变为异步且几乎免费的事件。预计会有更多的内存中和重叠检查点，在训练继续的同时在后台写入分片，加上纠删码和复制检查点，可以在千个 GPU 规模上常见的节点故障中幸存下来。云对象存储和更快的本地 NVMe 层将托管分片，安全张量等标准化格式将不断改进训练恢复和推理部署的安全、快速、部分加载。
现实世界的实施
前沿模型运行在数千个 GPU 上，每隔几百步自动保存分片检查点，因此单个失败的节点只需要几分钟，而不是几天。
Hugging Face 将大型开放模型分发为多个 safetensors 分片和一个 index.json，以便用户可以逐个下载和加载它。
研究人员恢复中断的微调，恢复精确的优化器动力、步数和数据加载器位置以无缝继续。
在廉价的可抢占式云 GPU 上进行现货实例训练，其中频繁的分片检查点使作业能够在被驱逐和重新安排的情况下幸存下来。
风险与防护栏
优化一项基准测试可以隐藏更广泛的系统弱点。
基础设施和维护成本常常被低估。
随着系统变得更加复杂，安全性和可观察性差距可能会扩大。
实施路线图
在实施之前定义延迟、质量和成本目标。
在实际负载和数据条件下进行基准测试。
仪器监控错误、漂移和用户影响。
在扩展之前准备回滚和事件响应路径。
不断探索
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the Checkpoint Sharding and Resumable Training quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
常见问题
什么是检查点分片和可恢复训练？
将模型的训练状态保存为片段（分片）的技术，以便可以保存和重新加载大型模型，而不会受到内存或磁盘限制的影响，因此崩溃的运行可以准确地从中断的地方继续。对于任何在多个 GPU 上运行数天或数周的训练作业来说都是至关重要的。
为什么大模型检查点要分成分片？
将巨大的检查点（数百 GB）作为一个文件是不切实际的；分片允许许多队列并行写入其切片并启用分段加载。
除了模型权重之外，可恢复检查点通常还必须保存什么？
为了准确地恢复，检查点存储优化器状态、随机数生成器状态、步数以及数据加载器停止的位置。
长期工作可恢复培训的主要好处是什么？
通过可恢复的检查点，中断的运行会重新加载其上次保存的状态并继续，而不是丢弃数天的计算。
每个 GPU 等级通常如何对分片检查点做出贡献？
由于分布式训练已经对模型进行了跨等级划分，因此每个等级仅写入其切片，从而实现并行保存和内存高效。
索引文件在分片检查点中起什么作用？
索引（例如，index.json）记录哪个分片保存每个参数，以便加载器可以获取并重新组装正确的片段。