meta-pytorch/torchft

解决的问题

大规模 AI 模型训练时，单个 GPU 或节点崩溃通常会导致训练失败，需要从最近的检查点重新启动整个训练任务。torchft 提供了每步容错能力，即使出现错误，也能在不中断整个任务的情况下继续训练。

工作原理

该库提供了在训练步级别进行成员变更的原语。它使用一个“灯塔”服务器来协调副本组，并通过每步心跳机制判断工作节点的健康状态。同时实现了容错版的 ProcessGroup 和检查点传输机制，支持在扩容操作期间从健康节点实时恢复状态。

适用人群

使用 PyTorch 训练大型模型（如 Llama 3 70B）的机器学习工程师和研究人员，特别是采用 DDP 或 HSDP 等复制权重策略的用户。

主要亮点

• 容错算法：开箱即用支持容错 DDP、容错 HSDP、LocalSGD 和 DiLoCo。
• 灯塔协调：专用服务器用于管理节点健康状态并协调副本组，避免“世界停顿”式的训练暂停。
• 实时恢复：在扩容操作期间，可通过专用检查点传输机制从健康对等节点恢复状态。
• 集成友好：可轻松集成到现有 PyTorch 训练脚本中，对训练循环的修改极小。