大模型分布式训练故障恢复篇
来源：AiGC面试宝典 作者：宁静致远 日期：2023年12月24日
一、为什么大模型分布式训练需要故障恢复？
大规模分布式训练场景由于集群规模过大，芯片设备、主机、网络等均会不定期出现故障。如果此时想要继续训练，就需要从上次存储的 ckpt（检查点）进行恢复，然后 resume training（恢复训练）。这个过程中产生的时间间隔就是集群故障带来的开销，虽然不可避免，但我们可以尽可能地减少故障带来的影响。> 📝 通俗解释：想象一下你正在写一篇很长的论文（训练大模型），每写一个小时才保存一次。如果电脑突然死机了，你之前的工作可能丢失很多。分布式训练也是同理，训练集群有成千上万个"电脑"同时工作，任何一个出问题都可能导致之前的工作白费。所以我们需要定期保存"进度"（ckpt），出了问题从最近保存的进度继续。
二、如何获取最优的 ckpt 存储间隔？
假设我们均匀地同步存储 ckpt，那么此时就需要根据集群环境获取最优的 ckpt interval（存储间隔）。首先，集群时间损失可以做如下定义，故障随机发生在 ckpt interval 区间：
$$ \text{集群时间损失} = \text{ckpt 存储耗时} + \text{故障期望次数} \times (\frac{\text{ckpt interval}}{2} + \text{恢复训练耗时}) $$
通过导数为 0，可以根据集群环境得到对应最优的 ckpt interval。当然，ckpt interval 肯定是远大于 1 的。
📝 通俗解释：这个公式在说：如果保存太频繁（间隔短），保存本身会浪费很多时间；如果保存太稀疏（间隔长），一旦出问题就要从头开始，损失更大。最好的策略是找到一个"平衡点"，让平均损失最小。打个比方：开车去加油站加油，加得太频繁浪费时间路程，加得太少可能没油抛锚，要找到最优的加油间隔。
三、ckpt 存储能否实现异步或者部分掩盖？
异步存储 ckpt 的最大问题：设备内存踩踏
如果在另外一个 stream（流）里做 D2H（Device to Host，设备到主机）数据拷贝，同时模型训练过程继续运行，那么就会有一个非常尴尬的问题：所有参数的 D2H 操作还没有完成，这时候下一个 step（步骤）已经开始更新参数或优化器状态，那么后面没完成的操作就会拷贝错误的数据。
由于 ckpt 存储时间不可控，不能确定是否小于下一个 step 的执行时间，所以内存踩踏的问题不可避免，即完全异步的方案是不可行的。
如果要想做到部分掩盖，本人认为可以有如下两个方案供选择：
- 方案一：在训练脚本侧修改，在下一次更新参数或优化器状态之前，强制等待 ckpt 存储完成，这样可以尽可能地 overlap（重叠）；
- 方案二：在框架侧修改，比如 H2D non-blocking 操作在后续有数据依赖的时候会强制加 sync point（同步点），框架侧也可以新增一个 D2H 拷贝，在后续有数据写操作的时候，强制添加 sync point。
📝 通俗解释：这就像一边搬家一边整理东西。异步存储的问题在于：你在把东西从房间A搬到仓库（保存ckpt），但同时你还在房间里继续收拾（训练），可能把之前还没搬走的东西又弄乱了，导致最后搬到仓库的是一团糟。完全不管不行，必须在"收拾"和"搬家"之间加一些协调机制。
四、断点续训/临终遗言是否真实可行？
绝对可行，但有一点受限。
大模型训练场景多是 DP/TP/PP 多维并行场景，任意一个节点出现故障的可能性都是存在的。如果任何一个 PP stage（流水线并行阶段）都存在一个完整的 TP Group，就是该 rank（进程）对应的节点没发生故障，那么整网参数就是完整的，可以在框架侧捕获分布式 error 做临终参数存储，这样 ckpt interval 就趋近于 0。
如果不满足整网参数完整这个条件，那是做不到临终存储整网参数和优化器状态的。根据经验，框架侧开发并不会很难，需要结合 rank 编排做定制研发。
当然，如果故障发生在参数或优化器状态更新的时候，那也是无法保证整网参数完整性的，这种情况也不能做临终处理。
📝 通俗解释：临终遗言就像"回光返照"——在机器彻底挂掉之前，赶紧把最重要的数据保存下来。理论上可行，但需要满足特定条件：至少要有一个完整的"小组"（TP Group）的节点还活着，这样它们的参数拼起来才是完整的。如果所有节点都在更新参数的瞬间同时挂掉，那就真的没办法了。
对于临终遗言/断点续训，基于训练框架对深度学习框架做深度定制是比较好的出路。
📝 整体总结：大模型分布式训练的故障恢复核心问题就是"多久保存一次进度"和"怎么保存才不打扰训练"。保存太频繁浪费时间，保存太稀疏风险大；完全异步保存会出错，需要加同步机制；临终遗言虽然理想，但需要特定条件支持。实际落地需要对训练框架进行深度定制。