大模型训练中的高效故障恢复方案

在当今大规模机器学习模型（如生成式语言模型或视觉语言模型）的训练过程中，通常需要部署数千甚至数万台 GPU 设备。即使采用如此大规模的并行计算，训练过程仍可能持续数月。面对频繁发生的硬件/软件故障（每天多次），传统检查点方案存在明显缺陷。

现有挑战

传统方案将模型状态定期保存至网络存储服务器，导致：

1. 每次检查点操作耗时 30-40 分钟，通常每 3 小时执行一次

2. 故障发生时可能损失数小时训练进度

3. 检查点恢复需要额外 10-20 分钟

Gemini 系统创新

提出的 Gemini 系统采用 CPU 内存检查点技术：

1. 双层级存储架构：

2. 主检查点保存在本地 CPU 内存的 RAM 驱动区

3. 容灾副本存储在集群其他机器的 CPU 内存中

4. 智能分组策略：

5. 将训练集群划分为 M 台设备的容灾组（M 通常为 2-3）

6. 采用数学证明的最优检查点分布算法

7. 通信优化：

8. 系统分析器动态识别训练流量空闲时段

9. 将检查点传输任务调度至空闲时段执行

关键技术突破

1. GPU 内存管理：

2. 为检查点分配固定大小的 GPU 内存缓冲区

3. 采用双缓冲区分块流水线传输技术

4. 实现检查点传输与训练计算的并行执行

5. 三级检索机制：

6. 优先从本地 CPU 内存恢复

7. 次选同组其他节点内存

8. 最后回退到远程存储

性能表现

在三种主流大语言模型训练测试中：

• 实现每次迭代的实时检查点保存

• 相比最优基线方案减少 92%的故障恢复时间

• 检查点操作对正常训练流量的影响可忽略不计

三种检查点方案的故障恢复时间对比（蓝色：基础远程存储方案，橙色：优化远程存储方案，绿色：Gemini 系统）

该技术已发表于 ACM 操作系统原理研讨会（SOSP 2023），为分布式训练系统提供了新的容错范式。更多精彩内容 请关注我的个人公众号 公众号（办公 AI 智能小助手）公众号二维码

办公AI智能小助手