Fault-Tolerant Hybrid-Parallel Training at Scale with Reliable and Efficient In-memory Checkpointing

文章标题：大规模可靠且高效的内存检查点混合并行训练容错
作者/机构：
- Yuxin Wang (香港浸会大学)
- Xueze Kang (香港科技大学（广州）)
- Shaohuai Shi (哈尔滨工业大学（深圳）)
- Xin He (香港浸会大学)
- Zhenheng Tang (香港浸会大学)
- Xinglin Pan (香港科技大学（广州）)
- Yang Zheng (华为技术有限公司)
- Xiaoyu Wu (华为技术有限公司)
- Amelie Chi Zhou (香港浸会大学)
- Bingsheng He (新加坡国立大学)
- Xiaowen Chu (香港科技大学（广州）)

A1 主要贡献

核心问题：为了高效扩展大型模型（LM）的训练，研究人员从数据并行（DP）转向GPU集群上的混合并行（HP）。然而，这些集群频繁遭遇硬件和软件故障。现有的内存中检查点优化方法通过将参数快照到设备内存来实现快速故障恢复，但在资源密集型的HP训练下，这些方法会在检查点操作和训练之间引发严重的资源竞争，难以扩展。如图1所示，现有的异步快照方法在DP设置下开销很小（1%），但在HP设置下开销激增至22%；而现有的内存中检查点保护方法，其节点间通信时间甚至可能超过训练迭代时间（高达6.63倍）。

研究目标：本文旨在设计一个分布式内存检查点系统，为混合并行训练提供低开销的检查点，并实现快速的故障恢复。

创新点：本文提出了REFT，一个可靠且高效利用易失性主机内存来保护快照以实现快速恢复的内存检查点系统。REFT包含REFT-save（用于内存检查点保存）和REFT-load（用于内存检查点加载）两个核心组件。其主要创新如下：
- 高效快照（Efficient Snapshotting）：引入了分层异步快照（Hierarchical Asynchronous Snapshotting, HAS），该方法有效利用设备空闲时间，减少混合并行训练中的资源竞争，从而在进行大规模训练的同时以接近零的开销进行快照。
- 高效可靠的内存保护（Efficient and Reliable In-Memory protecting）：REFT最小化了冗余保护的开销，并能根据不同的混合并行设置自动调整方案以生成不同数量的冗余。它引入了三种 intra-node（节点内）冗余方法来增强分布式检查点的完整性，避免了会阻塞训练的 inter-node（节点间）通信：
- 异步冗余复制 (Asynchronous Redundant Copying, ARC)
- 异步纠删码 (Asynchronous Erasure Coding, AEC)
- 异步优化器重计算 (Asynchronous Optimizer Recomputing, AOR)
- 这些方法由HAS提供支持，在冗余快照期间几乎不增加额外开销。
- 快速恢复：通过分布式内存检查点加载（Distributed In-memory Checkpoint Loading），REFT能够为失败的HP训练提供快速重启，绕过了低效的NFS读取。

在Frontier超级计算机上，使用256个MI250X设备（512个GPU）训练Llama-2-34B模型时，REFT实现了零内存检查点保存开销。

A3 背景知识/关键观察/设计原则

2.1 针对大型模型的并行训练技术

混合并行训练。混合并行（HP）是分布式大型模型（LM）训练的常用方法【索引24，Efficient large-scale language model training on gpu clusters using Megatron-LM，2021，SC】，它结合了三种主要的并行技术：数据并行（DP）、张量模型并行（TP）和流水线并行（PP）。DP【索引3, Revisiting distributed synchronous sgd, 2016, ICLR Workshop Track; 索引6, GeePS: Scalable deep learning on distributed GPUs with a gpu-specialized parameter server, 2016, EuroSys; 索引16, Imagenet classification with deep convolutional neural networks, 2017, Commun. ACM; 索引39, Communication-efficient distributed deep learning: A comprehensive survey, 2020, arXiv; 索引40, Gossipfl: A decentralized federated learning framework with sparsified and adaptive communication, 2022, IEEE Transactions on Parallel and Distributed Systems; 索引47, Poseidon: An efficient communication architecture for distributed deep learning on GPU clusters, 2017, USENIX ATC】将模型复制到多个设备上进行并行训练，可能导致参数冗余。为了解决这个问题，ZeRO-DP被引入【索引33，Zero: Memory optimizations toward training trillion parameter models，2020，SC20】，包括ZeRO-1、ZeRO-2和ZeRO-3等变体，它们将参数分布在不同的DP路径上。其中，ZeRO-1引入的通信开销较少，是扩展HP的首选。TP【索引7, Large scale distributed deep networks, 2012, NIPS; 索引15, One weird trick for parallelizing convolutional neural networks, 2014, CoRR; 索引26, Resource elasticity in distributed deep learning, 2020, MLSys; 索引37, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019, CoRR】将大张量的计算分布到多个设备上，使得能够处理超过单个设备内存容量的模型。每个设备独立进行计算，并根据需要进行通信。PP【索引22, PipeDream: Generalized pipeline parallelism for DNN training, 2019, SOSP; 索引34, ZeRO: Memory optimizations toward training trillion parameter models, 2020, SC; 索引41, Fusionai: Decentralized training and deploying llms with massive consumer-level gpus, 2023, arXiv; 索引44, Dorylus: Affordable, scalable, and accurate GNN training with distributed CPU servers and serverless threads, 2021, OSDI】将模型层划分为分配给不同设备的阶段。每个阶段处理前一个设备的中间结果，并将其输出转发到后续阶段。

大规模分布式训练的实践。在实践中，大规模分布式训练通常利用原生DP或ZeRO-1与TP和PP相结合，以确保HP的可扩展性。本研究遵循这种集成这些并行技术的通用方法，以实现有效的LM训练。

混合并行中的故障。研究人员观察到，在LM训练期间，系统故障的发生率很高，特别是当系统规模扩展到数千个GPU时【索引13, Megascale: Scaling large language model training to more than 10,000 gpus, 2024; 索引48, Opt: Open pre-trained transformer language models, 2022】。尽管这些故障可能发生在小范围内【索引13】，但它们需要从最新的检查点进行完整的系统重启。这个重启过程会产生巨大的重计算开销，特别是当最新的检查点是在数小时前创建的时候。

故障的两种主要类别。在LM训练期间遇到的故障主要有两类：硬件故障和软件故障【索引10, Understanding and mitigating hardware failures in deep learning training systems, 2023, ISCA; 索引12, Analysis of Large-Scale Multi-Tenant GPU clusters for DNN training workloads, 2019, USENIX ATC; 索引48】。由温度波动、电力中断和存储不一致等引发的硬件功能故障，在LM预训练中很普遍【索引10, 48】。另一方面，由于频繁的设备间通信和高内存需求，软件故障也很常见，这给MPI【索引9，Mpi: A message-passing interface standard，1994】和PyTorch内存管理【索引27，Pytorch: An imperative style, high-performance deep learning library，2019，NeurIPS】等软件组件带来了压力。

2.2 针对LM训练的内存中检查点

内存检查点的实现方法。LM训练的内存中检查点可以通过两种主要方法实现：异步参数快照和内存中检查点保护。

异步参数快照的资源竞争问题。异步参数快照可能与训练过程争夺内存带宽，从而可能减慢训练速度。现有研究主要引入两种方式来协调快照和训练以减少干扰。如图2所示，主机协调（Host Coordination）【索引20, CheckFreq: Frequent, Fine-Grained DNN Checkpointing; 索引42, TorchSnapshot: A performant, memory-efficient checkpointing library for PyTorch applications, 2022】使用一个单独的进程或线程来指导异步快照，而不考虑设备调度。相反，Gemini【索引46，Gemini: Fast failure recovery in distributed training with in-memory checkpoints，2023，SOSP】和TorchSnapshot【索引13，Megascale: Scaling large language model training to more than 10,000 gpus, 2024】引入了主-设备协调（Host-Device Coordination），利用PyTorch【索引27】的异步能力，允许设备上的计算（主要用于训练）和内存传输（主要用于快照）同时进行。内存中检查点保护由Gemini【索引46】在ZeRO-3中提出，利用节点间流量将本地参数分片分发到对等节点上。Gemini的设计允许训练和节点间检查点通信并发进行。

检查点的关键性能指标。在进行检查点操作时，我们关注两个指标，即内存中检查点开销 $T_{overhead-inmem}$ 和故障恢复时间 $T_{recovery}$。如图3所示，$T_{overhead-inmem}$ 被测量为检查点期间延迟的训练时间，计算公式为 $T_{i+1_{ckpt}} - T_{i_{train}}$，其中 $T_{train}$ 表示每次迭代的训练时间。$T_{recovery}$ 衡量训练从故障中恢复的速度，计算为故障发生时间与最新快照时间之间的时间差。当异步快照效率低下时，累积的 $T_{overhead-inmem}$ 会很大。因此，必须降低快照频率以减少对训练的阻塞。另一方面，降低快照频率会导致更高的 $T_{recovery}$。本文的目标是减少 $T_{overhead-inmem}$，以允许频繁的内存中检查点，从而减少故障恢复时间（即 $T_{recovery}$）。

2.3 局限性与动机

现有技术的开销问题。如图1所示，将现有的内存中检查点技术应用于HP下的LM训练会引入显著的开销，主要有两个原因。

现有快照方法的局限性。当检查点大小较小或GPU利用率较低时，现有的内存中快照方法可能是高效的【索引20, CheckFreq: Frequent, Fine-Grained DNN Checkpointing; 索引42, TorchSnapshot: A performant, memory-efficient checkpointing library for PyTorch applications, 2022】。例如，主-设备协调在ZeRO-3【索引46，Gemini: Fast failure recovery in distributed training with in-memory checkpoints，2023，SOSP】中被评估为高效，因为它在本地检查点大小较小、硬件FLOPs利用率较低且迭代时间较长时放宽了对异步快照的约束【索引35，Zero: Memory optimizations toward training trillion parameter models，2020】。然而，当使用HP训练LM时，FLOPs利用率变高，每个设备的快照大小也变大【索引13，Megascale: Scaling large language model training to more than 10,000 gpus, 2024】。更大的检查点大小需要更多的快照操作，这要求设备（即GPU）执行更多来自主机（即CPU）的异步复制指令，并消耗大量的PCIe带宽。

现有快照保护方法的局限性。内存中的检查点在硬件故障期间是不可靠的。Gemini【索引46】中当前的内存保护方法将检查点分片分发到对等节点以提高可靠性，但这另一方面也引入了通信开销。我们发现这种方法仅在ZeRO-3 DP训练模式下有效。在HP训练中，参数在ZeRO-1和原生DP【索引13】的所有DP组成员中是冗余或部分冗余的，这是HP中的常见设置。当每个节点上的检查点大小增加时，为了分发检查点分片而在对等节点之间进行的通信大小急剧增加。此外，HP训练在流水线和数据并行上下文中需要密集的节点间通信，这进一步加剧了对通信带宽的竞争，导致训练性能下降。

本文的动机。上述局限性促使我们设计新的内存中快照和保护优化方法，以减轻内存中检查点与训练之间的资源竞争，从而在混合并行下实现快速可靠的LM训练。

A2 方法细节

3 设计概览

REFT系统设计目标。为了设计一个能够缓解资源竞争的高效混合并行训练内存检查点系统，我们提出了REFT。如图4所示，REFT主要包含两个组件：用于检查点保存的REFT-save和用于检查点加载的REFT-load。我们主要在REFT-save组件中解决资源竞争的局限性，通过回答以下问题。

❶ 主机与设备如何高效协作以避免快照开销？。通过理解异步快照和训练操作如何相互干扰，我们提出根据干扰程度将快照过程与训练并行分配。如图4(1)所示，在参数分片后，当调用REFT-save时，分布式训练设备执行分层异步快照（Hierarchical Asynchronous Snapshotting, HAS）。HAS在调度快照过程时包括三个层次。首先，快照被分配到设备“气泡”（即设备空闲时间）。如果气泡时间不足，REFT使用第二层，将快照过程分离到不同的计算周期中以减少资源竞争。最后，快照被分配到不使用PCIe的高速节点内通信期间。第4节的性能分析表明这可能会引入开销，因此我们尽量避免此选项。通过HAS协调，REFT充分利用设备的空闲时间来减少竞争。

❷ 如何在增强分布式检查点完整性的同时，提供最小的冗余保护开销？。如图1所示，为内存检查点保护而进行的频繁节点间通信会阻塞训练。为了解决这个问题，REFT-save中的内存保护通过在对等节点间引入冗余来寻求优化机会。需要注意的是，HAS协调通过捕获对等节点负责保存的参数快照，在节点内级别创建了冗余。这种方法使快照大小增加了一倍，但可以从硬件故障中快速恢复。然而，如果这种冗余导致过多的开销，系统可以切换到异步纠删码（Asynchronous Erasure Coding, AEC），它将本地快照段转换为奇偶校验数据，从而减轻PCIe的负载。虽然使用AEC恢复速度较慢，因为它需要进行奇偶校验解码，但它有效地减轻了频繁快照保存的负担。

针对优化器参数的保护。对于在分片DP优化（例如ZeRO-1 DP）中不冗余的优化器参数，每个节点使用异步优化器重计算（Asynchronous Optimizer Recomputing, AOR）来捕获分配给对等节点的梯度分片进行快照，并在本地重构对等节点的优化器参数。这种方法确保了即使对于本身不冗余的参数，在主机的集体内存中也存在参数冗余。系统允许动态组合和切换这些分布式内存保护方法，这些方法可以通过用户设置或在训练开始时进行自动性能分析来配置。如第7节所示，集成这些保护方法以增强内存检查点的可靠性，只会产生极小甚至没有开销。

系统恢复流程。在接下来的部分，我们将详细介绍REFT的设计，重点是利用混合并行中可用的硬件资源进行高效的内存检查点。有了受保护的内存检查点，我们可以在节点故障时让节点通过以太网与对等节点通信其本地参数，并在不访问缓慢的NFS的情况下重启系统。我们将在第4节详细说明这个过程。

4 REFT：可靠且高效的内存容错

REFT系统概述。REFT旨在通过高效的保存和加载功能，为容错语言模型（LM）训练高效、可靠地扩展混合并行性，如图4所示。在包括数据并行（DP）、张量并行（TP）和流水线并行（PP）组合的混合并行部署中，REFT致力于最小化内存检查点开销，允许训练在发生故障时从最近的迭代中恢复，且进度损失最小。为实现这些目标，REFT提供了两个重要模块：REFT-save通过分布式快照和分布式内存保护来优化内存检查点保存；REFT-load通过分布式内存加载来提出内存检查点加载方案。

4.1 使用REFT-save进行分布式快照

分布式快照引擎的目标。REFT-save的分布式快照引擎旨在最大限度地利用整个系统的分布式并行化设备到主机的通信和存储I/O，并以对LLM预训练的最小干扰实现异步快照。

全局参数分片。在混合并行训练中，我们为并行检查点保存定义了标准的分片过程。如图5所示，在混合并行中，每个流水线并行阶段被指定为一个分片组 $SG_i$。在第 $i$ 个 $SG_i$ 中要划分的参数表示为 $P_i$。在同一PP阶段内共有 $D$ 个DP节点，每个DP上的每个阶段拥有 $|P_i|/D$ 的参数，从而将每个节点在检查点期间的数据传输负载减少了 $D$ 倍。在这项工作中，对于混合并行中的DP，我们专注于原生DP和ZeRO-1，它们能高效地扩展训练，并且常用于大规模LM训练。在ZeRO-1【索引33】中，分片的范围仅限于模型参数，因为优化器参数最初就在DP路径间进行了分区。只有模型参数在每个DP路径上是冗余的，需要进行分片以进行并行检查点。REFT将检查所有设备的冗余。如果发现冗余，我们将这些设备分组到一个分片组中。然后REFT在该组内对参数进行分片以保存。

分层异步快照（HAS）。为解决第2节中讨论的异步快照性能问题，该引擎通过策略性地最小化快照和训练操作之间的干扰。在全局参数分片过程将快照任务分片到所有设备以进行并行快照之后，引擎将快照构建为三层结构，如算法1所示。这些层次决定了每个快照操作的确切时间和大小：

Algorithm 1 Hierarchical Asynchronous Snapshotting

1: Input:   
2: P: Pipeline stages, W: Whole model parameters;   
3: $BW_{IO}$: IO bandwidth;   
4: $T_{fw,bw}$: Forward and backward time of one micro batch;   
5: # Scheduling snapshot based on runtime estimation:   
6: for pipeline stage p ∈ P in parallel do   
7:   $T_{snapshot}$ ← EstimateSnapshotTime(p, W, $BW_{IO}$ );   
8:   $T_{bubble}$ ← EstimateBubbleTime(p, $T_{fw,bw}$ );   
9:   if $T_{snapshot}$ ≥ $T_{bubble}$ then   
10:     $W_{bubble}$, $W_{p*}$ ← SplitParameter($W_p$, $T_{snapshot}$, $T_{bubble}$)   
11:   else   
12:     $W_{bubble}$, $W_{p*}$ ← $W_p$, $p_{empty}$
13:   end if   
14:   Launch async snapshot $W_{bubble}$ in bubble time;   
15:   Launch async snapshot $W_{p*}$ in computation time.   
16: end for   
17:   
18: def SplitParameter(W, $T_{snapshot}$, $T_{bubble}$):   
19:   N ← number vectorized W   
20:   $W_{bubble}$ ← W[: ⌊N × $T_{bubble}$ / $T_{snapshot}$⌋]   
21:   $W^*$ ← W[ ⌊N × $T_{bubble}$ / $T_{snapshot}$⌋ :]   
22:   return $W_{bubble}$, $W^*$   
23: def EstimateSnapshotTime(stage p, W, $BW_{IO}$):   
24:   return size($W_p$) / $BW_{IO}$   
25: def EstimateBubbleTime(stage p, $T_{fw,bw}$):   
26:   return (0.8P + 2|P| − p − 2) × $T_{fw,bw}$

❶ 第1层：在设备“气泡”（bubbles）期间异步执行快照以最小化开销。气泡定义为设备既不进行计算也不进行通信的时间。这种方法在流水线并行中特别有效，因为设备气泡很明显。它通过分析每个气泡的可用时间和检查点参数来实现这一点。

❷ 第2层：如果流水线气泡不足，则与计算流并行进行快照。由于过滤到这一层的快照操作减少了，计算和快照之间的干扰可能性很小，从而避免了训练中断。

❸ 第3层：与通信流异步操作。这一层仅在训练和快照使用独立的互连（例如，NVLink用于训练，PCIe用于快照）时才起作用，以减少可能的冲突。注意，除非前几层不足以进行快照操作，否则不会使用这一层。

算法1描述。算法1演示了调度异步快照与训练过程的过程。对于每个阶段 $p$，它首先估计两个关键时间：快照时间 $T_{snapshot}$ 和气泡时间 $T_{bubble}$。快照时间 $T_{snapshot}$ 是根据IO带宽 $BW_{IO}$ 和模型参数的通信成本计算的。气泡时间 $T_{bubble}$ 的估算是考虑到前向和后向传播的计算成本 $T_{fw,bw}$ 和GPU计算FLOPs $FLOPs_{GPU}$。需要注意的是，这里的气泡时间估算是专为1F1B流水线并行结构定制的。我们还为其他流水线并行配置提供了一个在线气泡分析器；我们可以通过运行时分析准确估计计算时间，这有助于快照的调度。如果快照时间 $T_{snapshot}$ 大于或等于气泡时间 $T_{bubble}$，算法将参数分为两部分：$W_{bubble}$ 和 $W^*$。算法优先利用气泡时间来快照 $W_{bubble}$。然后，模型的剩余部分 $W^*$ 在计算时间内进行快照。在我们在第7.1节的测试平台上，HAS在大规模混合并行训练设置中对训练引入的开销极小甚至为零。

4.2 使用REFT-save进行分布式内存保护

节点内冗余的必要性。由单个节点故障导致的内存检查点损坏可能导致全局检查点不完整，从而无法快速重启。因此，内存检查点方法必须为相同的检查点生成冗余。如第1节所讨论，为生成此类冗余而进行的节点间通信不适用于混合并行，会导致巨大的训练开销。因此，REFT选择通过节点内操作来生成冗余。对于不同的混合并行，分布式内存保护引擎确保在分布式节点间对内存检查点进行冗余保护，利用可用的设备和主机资源来优化参数安全性，而不增加训练开销。

异步冗余复制 (ARC)。在ARC中，当一个分片组 $SG_i$ 中的参数存在固有冗余时，每个成员不仅将自己的分片保存到快照中，还会保存来自 $SG_i$ 中对等成员的分片，以在全局主机内存中生成冗余。然后，在发生故障时，失败节点上丢失的参数可以从其他节点上的冗余副本中恢复。图6中的红色框给出了一个ARC的例子。给定PP大小为 $P$ 和DP大小为 $D$，ARC将每个 $SG_i$ 组成员的快照量增加一倍，达到 $2|P_i|/D$。这个负担可能会给其他异步快照方法带来巨大开销，但通过HAS协调，开销可以减少到接近零。使用这种方法，只要在 $SG_i$ 内没有超过一个节点发生故障，训练过程就可以恢复。

异步纠删码 (AEC)。为了在冗余引入快照开销时减小其大小，我们提出使用异步纠删码（AEC）作为保护引擎的第二种方法。AEC依赖于GPU设备的逻辑计算能力，并嵌入到设备气泡中。在我们的微基准测试中，设备上的纠删码在V100和A100 GPU上实现了比快照高12倍至15倍的吞吐量。图6中的蓝色框展示了一个在有四个节点的系统中使用AEC的例子。与ARC相同，为防止节点间通信中断PP通信，AEC对一个设备上的 $p0, p1, p2$ 进行快照，并使用以下编码方法冗余地保护 $p0, p0$ 和 $p0$：

i. 编码。在设备0的训练气泡中，REFT通过使用XOR计算对奇偶校验单元进行编码来计算 $p_0, p_0$ 和 $p_0$ 的奇偶校验 $parity_0$：

ii. 解码。当一个SG中的节点发生离线故障时，REFT可以通过XOR计算恢复该节点上的参数。假设 $DP_0$ 离线，系统使用减法解码器恢复 $p_0$：

异步优化器重计算 (AOR)。在优化器分片的数据并行（DP）中，例如ZeRO-1，优化器参数被均匀分布在DP组成员之间，失去了固有的冗余。然而，模型参数和梯度在每个成员上保持完整。为了在没有节点间通信开销或额外设备计算的情况下提供冗余保护，我们引入了一个名为异步优化器重计算（AOR）的主机上优化器保护模块。具体来说，在启用ZeRO-1的情况下，系统维护来自DP组中对等成员的优化器参数副本。图7给出了AOR的一个例子。具体来说，混合并行中的每个节点（标识为 $(p,d)$，属于流水线阶段 $p$）负责保护与本地分片大小相同的对等分片。

AOR的参数管理。给定总模型参数 $W_{model}$ 和总优化器参数 $W_{optimizer}$，分布在 $P$ 个流水线阶段，每个阶段有 $D$ 个数据并行成员，每个节点管理的模型和优化器参数分片表示为：

协同冗余保护。如果同时启用 $k$ 个冗余保护策略，或者使用不同参数的同一冗余保护策略 $k$ 次，我们的冗余保护系统的保护能力可以无任何开销地增强 $k-1$ 倍。如果在同一个 $SG_i$ 中发生 $k$ 个或更少的节点故障，训练过程可以被恢复。我们还在第7节评估了协同保护的效率。

4.3 使用REFT-load进行分布式内存加载

REFT-load的工作流程。如图4所示，REFT-load由分布式内存加载引擎驱动，该引擎通过三个步骤管理故障后从全局主机内存加载参数的同步。它通过以下方式避免了外部网络文件系统（NFS）存储读取的低效：❶ 指导每个节点从本地主机内存加载其检查点分片。❷ 促进节点间的all-gather操作以在网络上同步参数。❸ 每当节点经历参数丢失时，解码或传输丢失的参数，确保所有节点都能有效地恢复和同步其训练状态。这些组件协同工作，确保REFT在混合并行中提供可扩展且可靠的容错训练环境，最大化内存检查点效率。然而，如果失败的节点数量超过了内存保护的能力，REFT允许训练从NFS中的全局检查点重启。我们在第7节的分析表明，在大多数情况下，REFT中内存检查点的可靠性对于LM训练是足够的。

4.4 工作流

弹性功能。REFT在TorchElastic【索引28，TorchElastic，2021】的rendezvous功能上弹性运行。训练启动时，一个HEALTHY信号触发训练开始。当它接收到SNAP信号以开始快照时，REFTsave启动，主机内存异步保存参数。当所有张量都快照完成后，分片被标记为COMPLETED。发生故障时，一个UNHEALTHY（软件故障）或OFFLINE（节点故障）信号从故障节点广播，训练弹性重启。如果参数是可恢复的，REFT-load将在重启期间尝试恢复它们。

快照管理。在REFT中，快照保存在冗余的主机内存中，并根据需要刷新到tmpfs【索引38，tmpfs: A virtual memory file system，1990】。通常，它至少维护一个已完成的副本和一个正在进行的快照。已完成副本的数量受分配的主机内存限制，以避免内存不足问题。已完成和正在进行的快照在以下条件下进行管理以保持参数一致性：❶ 保存：正在进行的快照接受从设备内存刷新的参数。保存完成后，已完成的快照被正在进行快照的新副本替换。这个循环防止了可能损害LLM收敛的参数不一致性。注意，REFT以较低的频率将内存检查点保存到NFS，以防范大规模硬件故障。❷ 加载：当预训练过程突然失败或停止时，REFT-save将最新的已完成快照及其冗余持久化到tmpfs。如果故障是硬件级别的，导致节点重启或关闭，参数在所有节点进行all-gather后从冗余备份中重构。如果重构失败，训练从NFS中现有的检查点重启。

A7 补充细节

5 系统可靠性分析

假设1。在多节点GPU系统中，我们假设所有节点的故障概率是独立的。这些系统的故障前时间（Time-to-Failure, TTF）——从系统开始运行到第一次故障的时间——遵循威布尔分布。威布尔分布是一种广泛用于描述系统寿命和故障模式的统计模型，常用于故障建模场景【索引43, Reliability model of a system of k nodes with simultaneous failures for high-performance computing applications, 2013, Int. J. High Perform. Comput. Appl.; 索引36, A large-scale study of failures in high-performance computing systems, 2006, DSN'06】。

故障前时间（TTF）。TTF指的是系统或组件可靠运行的持续时间。我们将系统的故障率定义为 $\lambda_{system}$。基于我们的假设以及系统已运行时间 $t$ 的事实，累积的生存概率由 $P = e^{-\lambda t^c}$ 给出，其中 $c$ 是威布尔分布的形状参数，影响分布的形状，从而影响故障概率模式。假设在一个有 $N$ 个节点的系统中存在 $N/D$ 个 $SG$。我们将最新快照中捕获的参数称为当前参数。

使用REFT的生存概率。REFT提供了最大的内存保存效率和接近零的训练开销，因此允许迭代级别的保存。然后，使用AOC或AEC的单层冗余保护在REFT中的生存概率，表示为 $P_{re-survive}$（re代表REFT），表示REFT成功保护内存检查点不丢失的可能性，假设所有子模块正常工作且每个 $SG$ 中最多有一个节点发生故障。这个概率与训练过程无关。在多节点系统故障中，参数能够存活并通过REFT从内存检查点高效重启的总体概率表示为：

不使用REFT的生存概率。不使用REFT，预训练无法在节点故障中存活，必须经历NFS检查点加载。不使用REFT的生存概率取决于所有节点都健康：
$P_{ck-survive} = P_{sw}^N$
其中 $P_{ck-sw}$ 表示单个节点在软件故障中存活的累积概率。$P_{ck-survive}$ 和 $P_{re-survive}$ 随时间 $t$ 的变化如图8所示。我们可以得出结论，REFT通过保护参数显著增强了生存概率，并且使用REFT可以大大降低NFS检查点的频率。

7.6 讨论

REFT的性能表现。REFT在效率和可靠性方面表现出色。凭借快速的快照功能，它可以实现最佳的保存间隔和最小的容错开销。值得一提的是，重启时减少的重计算开销与保存效率同样重要，这一点被一些论文所忽略。尽管如此，这项工作仍存在一些局限性。

洞察1：合理的资源利用。REFT在CPU内存中最多使用优化器和模型参数存储量的3倍，分配给1) 快照缓冲区。举例来说，在我们的Llama-2.7B预训练测试平台上，峰值CPU内存使用量仅为30.45GB，其中还包括数据加载器的缓存。结合参数分片有效地减轻了单个节点上的CPU内存占用。

洞察2：意外的训练加速。我们观察到，当我们通过异步快照或进行优化器参数重构来提高CPU利用率时，迭代时间出人意料地减少了。这个有趣的观察是由于更快的CPU操作（例如，在第一个GPT层中用于输入数据加载的aten::ones和aten::tril）瓶颈了训练。

局限性。由于GPU资源有限，我们无法进行更大规模的预训练。尽管如此，我们已经用部署级别的GPU资源（即Frontier中多达512个GPU）证明了REFT的效率和可靠性。此外，异步容错对训练的低级别干扰可以通过低级别编程进一步减轻。由于混合并行训练需要密集的节点内和节点间GPU通信，在利用优化的GPU到GPU通信拓扑的训练系统上，这种干扰可以被显著减少。

8 相关工作

8.1 同步和异步流水线并行

流水线并行中的参数更新。在预训练期间，流水线并行可以同步或异步地更新参数。虽然异步流水线并行减少了流水线中的气泡大小，但这是以牺牲准确性为代价的。例如，PipeDream【索引22，PipeDream: Generalized pipeline parallelism for DNN training，2019，SOSP】使用来自流水线内不同迭代的梯度来更新参数。相比之下，同步流水线并行，如在Megatron【索引24，Efficient large-scale language model training on gpu clusters using Megatron-LM，2021，SC】和OPT【索引48，Opt: Open pre-trained transformer language models, 2022】中所见，在DP路径间执行同步参数更新。REFT在同步流水线并行的同一迭代的前向和后向传播期间，促进了一致的参数快照。

8.2 容错方向

检查点加速技术。有许多工作使用分层或异步方法来加速检查点【索引2, Distributed snapshots: Determining global states of distributed systems, 1985, ACM Transactions on Computer Systems (TOCS); 索引14, Checkpointing and rollback-recovery for distributed systems, 1987, IEEE Transactions on software Engineering; 索引21, Design, modeling, and evaluation of a scalable multi-level checkpointing system, 2010, SC'10】。

相关领域的容错研究。此外，之前关于推荐模型训练的工作【索引18，Ecrm: Efficient fault tolerance for recommendation model training via erasure coding】探索了基于快照的容错的可能性，并取得了有希望的结果。研究人员也正在努力优化分布式检查点的性能【索引4, Project adam: Building an efficient and scalable deep learning training system, 2014, OSDI; 索引5, Solving the straggler problem with bounded staleness, 2013, HotOS; 索引23, Pipedream: generalized pipeline parallelism for dnn training, 2019, SOSP; 索引25, Deepfreeze: Towards scalable asynchronous checkpointing of deep learning models, 2020, CCGrid; 索引29, Fault tolerance in iterative-convergent machine learning, 2019, ICML】。一个与REFT正交的方向是有损检查点【索引19, CPR: Understanding and Improving Failure Tolerant Training for Deep Learning Recommendation with Partial Recovery, 2020, arXiv; 索引30, Fault tolerance in iterative-convergent machine learning, 2021】，其中【索引30】将故障表述为扰动的概念。部分恢复是基于参数的一部分位于一个服务器上的事实。优先检查点保存自上次保存以来变化最大的参数。

A4 实验环境

测试平台：
- NVIDIA集群：
- GPU：512 x V100-SXM2-16GB
- 互连：每个GPU通过300GB/s的P2P NVLink带宽连接到三个其他GPU。节点间通过100Gbps以太网连接，NFS连接速度为10Gbps。
- Frontier集群：
- GPU：256 x MI250x设备 (共512个GPU)，每个设备包含2个GPU。
- 互连：节点内每个MI250X通过50-100GB/s的P2P Infinity Fabric带宽连接到七个对等设备。节点间使用200Gbps的以太网NIC。

模型与数据集：
- 模型：使用不同规模的Llama-2模型进行预训练，包括Llama-2-1.3B, 2.7B, 6.7B, 34B。
- 数据集：所有评估均使用Wikipedia数据集。

软件配置：
- 框架：基于DeepSpeed v0.12.4和PyTorch v2.2.1实现，包含超过4000行Python代码。
- 核心功能：REFT的核心功能是可插拔的，可用于其他基于PyTorch的深度学习训练框架。
- 实现细节：
- 主机端使用多线程进行异步协调，并为缓冲区启用固定内存（pin memory）以加速设备到主机的传输。
- 设备端使用多个CUDA流和ROCm进程进行协调。
- GPU设备间通信通过CUDA中的NCCL和ROCm中的RCCL处理。
- 纠删码计算通过对FP32或FP16参数进行位异或（bitwise XOR）实现。

基线方法：
- 异步快照基线：
1. 主机协调（Host Coordination）：使用单独的进程或线程指导异步快照，不考虑设备调度。
2. 主-设备协调（Host-device Coordination）：在主机线程指导下，设备使用PyTorch的异步能力调度快照和训练。
- 内存保护基线：
1. 节点间检查点通信：来自Gemini【索引46】的方法。
2. 无快照或保护：作为性能对比的基准。

评估指标：
- 迭代时间（秒）
- 每秒样本数（Samples per second）
- 每秒万亿次浮点运算数（TFLOPs）

A4 实验结果

7.2 分层异步快照（HAS）的效率
- 实验内容：在多达512个V100 GPU上预训练Llama-2模型，评估HAS的效率。
- 实验结果：如图9和图10所示，与现有的异步快照方法相比，HAS的迭代时间和TFLOPs性能提升超过17%，开销接近于零。在弱扩展性实验中（表3），HAS引入的开销也低于3%。这极大地提高了内存检查点的可能频率。

7.3 内存检查点保护的效率
- 实验内容：在512个V100 GPU上预训练Llama-2-34B，评估ARC、AEC以及两者组合的性能。
- 实验结果：如图11所示，使用ARC仅引入1%的开销。ARC和AEC的组合也只对训练引入了有限的开销。

7.4 REFT-Load的效率
- 实验内容：在使用TP4, PP4, DP8的混合并行设置下预训练Llama-1.3B，评估REFT-load在不同冗余保护方法下的加载性能。
- 实验结果：如图12所示，从NFS加载检查点到设备耗时9.27秒。而使用ARC和AEC从内存加载分别仅需0.84秒和0.75秒，比NFS加载快高达12.36倍。这验证了REFT-load驱动的内存检查点加载的效率。REFT是首个允许混合并行训练进行内存检查点加载的系统。

7.5 案例研究：REFT在Frontier上的泛化能力
- 实验内容：在Frontier上训练Llama-34B模型，评估ARC和AOR方法的开销。
- 实验结果（图13）：
- 实验一（DP16, PP4, TP8 + ARC）：ARC的平均迭代时间为6.32秒，相比无检查点的6.26秒，开销仅为1%。
- 实验二（DP16, PP4, TP8, ZeRO-1 + AOR）：AOR的平均迭代时间为6.63秒，相比无检查点的6.59秒，开销仅为0.6%。
- 结论分析：两种方法都引入了极低的开销，有时甚至会因为改善了CPU瓶颈而加速训练。

A5 结论

在混合并行（HP）下训练大型模型（LM）需要大规模计算资源，这也凸显了处理容错问题的重要性。本文介绍了REFT，据我们所知，这是第一个用于基于HP的LM训练的内存检查点框架，其检查点开销极低甚至为零。REFT集成了一系列系统优化，以减少异步检查点快照和保护与模型训练之间的资源竞争。全面的评估表明，REFT极大地减少了参数快照和保护的开销，并允许进行迭代级别的快照。在发生故障时，它能从健康节点上的冗余参数或奇偶校验数据中高效地重建参数，从而最大限度地减少GPU工作损失。REFT已被集成到一个工业级LM训练平台中，有效提高了模型训练的容错能力和效率。