DeadPool：通过零开销检查点实现热插拔的弹性LLM训练

# DeadPool：通过零开销检查点的热替换实现弹性LLM训练

来源：https://arxiv.org/html/2607.01646

###### 摘要

最先进的大语言模型（LLM）训练需要数万个GPU运行数月，并在整个软件和硬件栈中遭遇故障。现有的容错机制要么在无故障执行期间引入显著开销，要么在恢复时面临漫长的延迟，尤其是在少量计算节点发生永久性故障的场景下。我们提出DeadPool，同时解决这两个优化目标。DeadPool引入了一种通过热替换恢复LLM训练的容错机制，即在无需终止整个作业的情况下，用备用节点替换故障节点。DeadPool的热替换基于两个关键思想：首先，它利用非关键路径上的内存检查点机制实现空间冗余。其次，它引入了一个通信器重建协议，在运行时用备用节点替换故障节点。DeadPool有效地将内存检查点与计算重叠，从而在无错误执行期间实现零开销。当发生永久性节点故障时，DeadPool可以借助内存检查点以最小的重计算量重建内存状态。我们在不同规模（高达512个NVIDIA A100 GPU）和不同LLM（高达65B参数）上评估了DeadPool，观察到零检查点开销，热替换恢复在40秒内完成。这些结果表明，DeadPool同时实现了零开销的无错误执行和极低的恢复成本。

## 一、引言

LLM训练在数千到数万个GPU上运行，节点故障不再是罕见事件，而是预期的运行条件[12,26]。来自一个32K-GPU LLM预训练部署的已发布中断统计报告显示，共有678次意外中断，其中GPU HBM内存故障、PCIe设备故障和NCCL看门狗超时共占所有事件的49.9%[33]。传统的检查点重启方案会定期写入持久化检查点，例如每1000步写一次[29,12]，以摊薄检查点开销。发生故障时，作业终止，用户需要启动一个新作业，重新加载最新的检查点。这种方法需要为恢复成本付出两次代价：它在正常执行期间引入备份开销，并且当发生故障时，它既承担恢复延迟，又承担最后检查点之后已完成工作的重计算成本。在过度使用的超级计算机上，如NERSC Perlmutter或TACC Vista，单比特内存错误或NCCL错误可能导致整个作业退出。用户必须重新进入已经拥挤的队列，从而增加额外的等待时间。

图1：现有容错策略在无错误开销与恢复成本权衡空间中的概念图

图1概念性地展示了现有容错LLM训练策略在无错误开销与恢复成本之间的权衡空间。检查点策略，如Datastates-LLM[25]、ByteCheckpoint[38]、TRANSOM[43]和Universal Checkpointing[21]，大幅提升了检查点的吞吐量和灵活性。然而，它们仍然通过从持久化检查点重启训练来进行恢复，这可能导致因丢失训练进度而产生的非平凡重计算成本。先前的内存检查点策略[27]降低了写入延迟，但假设在发生故障时内存检查点仍然可访问，这在永久性故障情况下并不成立。另一方面，弹性容错策略，如Bamboo[37]，通过利用流水线并行中的气泡来容忍成员变更。但这些方法会降低训练吞吐量。

我们在这项工作中的目标是同时减少无错误执行和恢复的开销。对于无错误执行，LLM框架需要异步持久化检查点，以避免与训练进程发生潜在冲突。为了降低恢复成本，LLM框架需要在每次迭代时都设置检查点，以最小化故障时的重计算量。此外，LLM框架需要避免完全重启，因为完全重启会引入全局作业终止、检查点一致性验证、文件系统I/O和内存状态重建的额外开销。这些具有挑战性的需求促使我们设计了DeadPool，一个通过热替换实现LLM训练的容错系统，它利用复杂并行中的零开销异步内存检查点。

DeadPool的设计基于两个独特见解。第一个见解是，在复杂并行（例如3D并行）中，LLM训练期间的内存状态使用不同的策略分布。典型的3D并行设置将切分的数据并行与流水线并行和张量并行叠加，它在数据并行对等节点之间复制模型分片，同时将优化器状态（即一阶和二阶矩）分布在一个数据并行rank（即一组GPU）内的GPU上。当发生永久性节点故障时，模型参数可以从共享相同逻辑分区的健康rank中重建，但优化器状态分片无法恢复。因此，我们设计了一种低开销、跨节点的内存优化器状态检查点策略，有效地将通信与计算重叠，从而使DeadPool能够在每次迭代中复制优化器状态分片，同时避免暴露检查点开销。第二个见解是，DeadPool将节点故障恢复视为一个在线拓扑修复问题，而不是完全作业重启。因此，我们设计了一个分布式通信器重建协议，用备用节点替换故障节点，以避免完全分配重启。

我们在PyTorch[31]和Megatron-LM[29]之上实现了DeadPool，但其设计仅依赖于通用能力：分片模型和优化器状态、主机端备用节点管理、共享恢复协调以及运行时组重建。DeadPool包括一个错误分类层，用于区分瞬时故障和永久性故障，允许DeadPool仅对实际使节点成员失效的故障应用拓扑变更恢复。我们在两个生产级HPC系统（即NERSC Perlmutter和TACC Vista）上评估了DeadPool，分别使用多达512个A100 GPU和64个GH200 GPU。我们向正在运行的训练作业中注入节点故障，并保留备用节点，这是容错和恢复研究中的标准方法[2]。在两个平台上，我们观察到每个步骤的检查点开销在训练关键路径上可以忽略不计。此外，DeadPool一致地在故障后约40秒内完成节点替换并恢复训练。我们的主要贡献如下：

- • 我们提出DeadPool，一种用于大规模LLM训练的在线热替换恢复机制。
- • 我们设计了一种用于优化器状态分片的异步邻居内存复制策略。
- • 我们在多达512个GPU和65B模型规模上评估了DeadPool，展示了无错误执行期间的零开销以及40秒内的一致恢复成本。

DeadPool的零开销检查点依赖于GPU集群配置、模型大小和模型部署策略。尽管DeadPool仅在Perlmutter和Vista上进行了测试，但我们预期DeadPool在许多其他GPU密集的超级计算机上同样有效。这是因为在当前一代机器上，检查点开销比计算低一个数量级，如图6所示。鉴于即将到来的NVIDIA Grace-Blackwell超级计算机上GPU计算能力和互连带宽的增长趋势，DeadPool有望继续实现零开销检查点。

## 二、背景

### 二-A 大语言模型训练

LLM的能力通常随着模型和数据集大小的增加而增强[5,15]。因此，为了克服训练时间和内存限制，研究人员设计了各种分布式训练策略，包括数据并行、张量并行、流水线并行和序列并行[4]。**数据并行**跨处理器复制模型。然后训练迭代五个步骤：1) I/O，2) 前向计算，3) 梯度评估，4) 梯度交换，5) 变量更新。**3D并行**通过在模型并行和流水线并行之上叠加数据并行来跨处理器划分模型。处理器被划分为组（数据并行rank），每个组容纳模型的完整副本。在每个组内，模型按层分布到各个节点，形成流水线阶段。在每个节点内，层（张量）均匀分布在多个处理器上，这被称为模型并行或张量并行。使用数据并行方法，唯一的通信是在训练开始时进行广播，以及在每次迭代中进行梯度交换。相比之下，3D并行在前向和后向计算中都需要节点内和节点间通信。**长序列支持**是LLM训练的第四个维度，使得从更大的上下文中学习成为可能。研究人员设计了序列并行，对查询、键和值矩阵采用不同的划分和通信策略[20,22,23]。在现有方法中，环注意力通过使用查询块（跨GPU）的外循环和键值块（GPU内）的内循环实现了最低的内存开销。这些分布式LLM训练策略使得容错LLM训练系统的设计变得复杂，因为它们容易导致次优的恢复决策和无错误阶段中不必要的开销。

### 二-B 系统中的错误

在计算机系统中，硬件瞬时故障（即软错误）[30,3]可能由宇宙射线的高能粒子[9]、晶体管缩小和低电压操作[10]触发。这些硬件故障表现为比特翻转，范围可以从单个比特到多个比特，甚至到数百行存储单元[3]。我们接下来解释故障/错误/失效的概念[36]：硬件**故障**在系统中传播；一旦它到达软件层面，就变成对程序可见的**错误**。该错误进一步在程序执行中传播，可能导致三种可能的结果[6]：1) 无影响（即正确输出），2) 静默数据损坏，3) **失效**，包括应用程序崩溃和系统无响应。前两种结果对用户不可见，尽管如先前研究所量化[45,14]，静默数据损坏可能会改变训练产生的模型参数。对于可察觉的错误，科学家们设计了一整套用于检测和容错的技术。由于我们的恢复方法选择会因错误类型而有很大差异，因此错误特征描述和分析对于理解，更重要的是，对于高效检测和纠正错误至关重要[32]。准确理解错误，如果可能的话，准确预测即将发生的错误和失效，对于我们及时采取缓解措施以避免数据丢失并提高系统可靠性至关重要[32]。

表一：故障恢复方法的比较。DeadPool维护持续可恢复的状态，并通过拓扑修复实现在线恢复。

## 三、相关工作

现有的容错解决方案可以分为两类。**检查点**方法从最近的检查点恢复模型状态，而**排除并运行**方法排除故障设备，同时允许健康工作者继续不间断运行。早期的神经网络训练方法采用弹性训练方法来处理故障，像Horovod[35]和PyTorch[31]这样的框架提供了故障恢复的API。然而，这些方法主要针对能够适应单个GPU的模型的数据并行。随着模型规模的增长，训练过程中遇到故障的可能性也随之增加[17,42,26]。EasyScale[19]和DLRover[40]考虑了混合并行性，然后提出了可应用于数千GPU的弹性训练流水线。在使用竞价训练[1,37]的集群中，故障率甚至更高。Varuna[1]使用定期检查点来从一个或多个故障中恢复，而Bamboo[37]利用流水线并行性，利用相邻阶段之间的气泡创建冗余。然而，这两种方法都引入了额外的重计算或检查点开销，随着故障率的增加，这会显著影响训练效率。Oobleck[16]通过引入能实现快速故障恢复和充分利用GPU的流水线模板，在吞吐量和容错之间取得了平衡。Parcae[7]利用一种新颖的指标“liveput”，在抢占式环境中动态重新配置并行策略，最大化吞吐量。由于对混合并行性的频繁调整可能导致开销，Slipstream[11]利用健康的GPU接管流水线阶段中故障GPU的计算，而不改变策略。由于基于检查点的恢复开销较高，Checkfreq[28]引入了一种细粒度的自动检查点机制，将计算与保存模型状态重叠，通过运行时分析优化检查点频率。Gemini[41]...