专栏名称: GiantPandaLLM

专注于机器学习、深度学习、计算机视觉、图像处理等多个方向技术分享。团队由一群热爱技术且热衷于分享的小伙伴组成。我们坚持原创，每天一到两篇原创技术分享。希望在传播知识、分享知识的同时能够启发你，大家一起共同进步(･ω<)☆

【PyTorch 奇淫技巧】Async Checkpoint Save

GiantPandaLLM · 公众号 · 3D · 2024-09-09 10:15

正文

请到「今天看啥」查看全文

因此，对故障的鲁棒性与训练进度之间很难做到权衡，但现在有了异步Checkpoint，PyTorch 分布式训练能够显著缓解这种张力，并以最小的影响整体训练时间的方式实现频繁Checkpoint。

大约一年前(https://pytorch.org/blog/performant-distributed-checkpointing/)，我们展示了分布式Checkpoint如何大幅加速Checkpoint时间，从最初的 torch.save() 功能开始。正如 IBM 研究团队指出的那样， torch.save 可能需要长达 30 分钟才能检查一个 11B 模型（PyTorch 1.13）。

随着分布式Checkpoint的进步，对于高达 30B 的模型大小，Checkpoint可以在 4 分钟内完成。使用异步Checkpoint，Checkpoint导致的训练时间损失现在降至 30 秒以下，通常仅需 6 秒。

需要明确的是，异步Checkpoint不会压缩实际的序列化Checkpoint时间，如之前的更新所展示的那样。相反， 它将最终的Checkpoint过程移出关键路径（到 CPU 线程），以允许 GPU 训练在单独的线程下完成Checkpoint的同时继续进行 。