Universal Checkpointing: A Flexible and Efficient Distributed Checkpointing System for Large-Scale DNN Training with Reconfigurable Parallelism

文章标题：通用检查点：一种灵活高效的分布式检查点系统，用于具有可重构并行性的大规模DNN训练
作者/机构：Xinyu Lian (UIUC), Sam Ade Jacobs (Microsoft), Lev Kurilenko (Microsoft), Masahiro Tanaka (Microsoft), Stas Bekman (Snowflake), Olatunji Ruwase (Microsoft), Minjia Zhang (UIUC)

---

A1 主要贡献

本文旨在解决大规模深度神经网络（DNN）训练中的一个核心问题：现有训练系统在通过检查点技术重新配置并行策略方面的支持非常有限。随着模型规模、数据量和序列长度的急剧增长，训练任务常常需要跨越多台机器，并采用复杂的分布式并行策略。然而，训练过程漫长且容易因硬件故障、软件错误或资源弹性管理而中断。当中断发生时，现有的检查点系统由于其结构与特定的模型并行策略和硬件配置紧密耦合，导致无法灵活地适应新的训练环境（例如，使用不同数量的GPU或迁移到其他集群）。

研究目标：
本文的研究目标是设计并实现一个新颖的检查点系统，名为通用检查点（Universal Checkpointing, UCP），以实现灵活高效的、具有可重构并行性的DNN训练。该系统旨在让训练任务能够灵活地重新配置其并行策略，同时保持训练的收敛性和准确性，从而提高大规模训练的弹性和效率。

创新点与核心贡献：
为了应对现有系统的挑战，UCP引入了以下核心创新：
1. 原子检查点（Atomic Checkpoint）：提出了一种新的检查点结构，它将检查点文件与任何特定的并行策略和硬件配置解耦。原子检查点为模型状态（权重、优化器状态等）提供了一个统一的、整合的视图，作为不同并行策略之间转换的通用中间表示，从而实现了灵活的重构。
2. 基于模式的重构管道（Pattern-Based Reconfiguration Pipeline）：设计了一套精心定义的并行模式（如Unique, Replicate, Shard-V, Shard-NC等）和基于模式的重构操作（如Extract, Union等）。该管道能够系统性地、自动化地将来自任意源并行策略的分布式检查点转换成原子检查点，并再映射到任意目标并行策略，覆盖了包括3D并行、ZeRO以及稀疏混合专家（MoE）和分组查询注意力（GQA）等复杂模型架构。
3. 高效的系统优化：为降低重构开销，UCP引入了多项优化技术：
* 嵌套并行重构（Nested Parallel Reconfiguration）：将重构过程构建为一个MapReduce问题，并利用多节点和多核CPU进行并行转换，同时通过负载均衡避免“掉队者”问题，相比顺序方法，转换时间缩短了14-257倍。
* 冗余旁路加载（Redundancy-bypassing loading）：在数据并行组内，通过让每个工作节点只加载部分原子检查点，然后通过高速互联（如NVLINK）进行all-gather通信，避免了从存储中重复读取数据，减轻了I/O压力。
* 惰性重构调用（Lazy Reconfiguration Invocation）：重构过程仅在源并行策略与目标策略不同或硬件发生变化时才被触发，对正常的训练检查点保存过程无任何性能影响，使得重构开销在总训练时间中的占比低于0.001%。

本文通过在多种大型语言模型（包括数十亿到万亿参数级别的密集和稀疏模型）上的广泛评估，证明了UCP在支持的并行策略广度、准确性、覆盖范围和效率方面的优越性。UCP已成功应用于多个真实世界的大规模模型训练任务中（如BigScience BLOOM 176B），显著增强了训练过程的灵活性和对动态硬件环境的适应能力，并已在DeepSpeed中开源。

---

A3 背景知识、关键挑战与机遇

2 背景知识

分布式训练的并行策略。高效地进行大规模DNN训练通常涉及以下并行策略。

数据并行（DP）。数据并行是扩展DNN训练最常用的策略之一。它将一个小批量（mini-batch）数据分割成多个微批量（micro-batch）分配给数据并行工作节点，但模型状态是复制的。每个工作节点独立计算一个微批量的梯度，随后通过all-reduce操作同步梯度，以在整个训练过程中保持模型参数的一致性。为了消除DP中复制模型状态的冗余，零冗余优化器（ZeRO）【索引46，ZeRO: Memory Optimization Towards Training A Trillion Parameter Models，arXiv preprint arXiv:1910.02054，2019】被提出，用以逐步消除优化器状态（ZeRO-1）、梯度（ZeRO-2）和权重（ZeRO-3）中的模型冗余。

张量并行（TP）。为了训练超过单个设备内存的模型，张量并行沿非批处理轴将计算和内存密集型的张量操作（例如矩阵乘法）分区到多个设备上。然后，它让每个GPU并行计算分区后的操作，并在这些分区算子的分割点和合并点使用通信集合操作（例如all-reduce、all-gather）来确保状态的一致性【索引48，Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism，arXiv preprint arXiv:1909.08053，2020】。

流水线并行（PP）。与张量并行不同，流水线并行将模型状态划分为多个阶段，分布在不同的GPU上【索引18，GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism，NIPS’19，2019】、【索引37，PipeDream: generalized pipeline parallelism for DNN training，SOSP’19，2019】。同时，流水线调度器将一个小批量分割成多个微批量，并使这些微批量流经不同的流水线阶段。两个连续的流水线阶段通过点对点通信交换中间数据。这种数据依赖性会导致流水线气泡，即计算停顿。为了减少流水线气泡以提高整体流水线效率，已经引入了多种流水线调度策略，如1F1B【索引37，PipeDream: generalized pipeline parallelism for DNN training，SOSP’19，2019】和交错式1F1B【索引38，Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM，arXiv preprint arXiv:2104.04473，2021】。

序列并行（SP）。受训练长上下文LLM的驱动，研究人员还开发了序列并行，它沿着序列输入维度划分模型激活，并利用分布式注意力机制等技术来容纳训练长序列输入时产生的大量激活内存【索引21，DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models，arXiv preprint arXiv:2309.14509，2023】、【索引30，Sequence Parallelism: Long Sequence Training from System Perspective，arXiv preprint arXiv:2105.13120，2022】。

并行策略的组合。这些并行策略通常被结合起来以扩展现代DNN的训练。训练系统需要协调不同的并行策略，以实现高训练效率。例如，像Megatron-LM和DeepSpeed这样的先进训练系统，将这些并行策略结合起来用于扩展基于Transformer的LLM，这被称为3D/4D并行【索引24，MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs，NSDI’24，2024】、【索引38，Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM，arXiv preprint arXiv:2104.04473，22021】、【索引51，Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model，arXiv preprint arXiv:2201.11990，2022】。

3 问题描述

可重构并行问题的形式化定义。本节中，我们对分布式深度学习训练的可重构并行问题进行形式化描述。具体来说，给定以下输入：（1）一个训练任务，其DNN模型为 $\Phi$，在训练迭代 $t$ 时的模型状态为 $\Phi(t)$（包括权重和优化器状态）；（2）一个源并行策略 $P_{src}$，该策略决定了 $\Phi(t)$ 如何在 $N_{src}$ 个并行工作节点（例如64个NVIDIA A100 GPU）上进行分区和并行执行；（3）一个目标并行策略，运行在具有 $N_{tgt}$ 个工作节点（例如48个NVIDIA A100 GPU）的目标硬件配置上，模型训练需要被重构至此。在这两种情况下，并行策略可以是数据并行（DP）【索引29，PyTorch Distributed: Experiences on Accelerating Data Parallel Training，arXiv preprint arXiv:2006.15704，2020】、张量并行（TP）【索引48，Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism，arXiv preprint arXiv:1909.08053，2020】、流水线并行（PP）【索引18，GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism，NIPS’19，2019】、【索引37，PipeDream: generalized pipeline parallelism for DNN training，SOSP’19，2019】、序列并行（SP）【索引30，Sequence Parallelism: Long Sequence Training from System Perspective，arXiv preprint arXiv:2105.13120，2022】、ZeRO【索引46，ZeRO: Memory Optimization Towards Training A Trillion Parameter Models，arXiv preprint arXiv:1910.02054，2019】，或这些策略的组合，如3D/4D并行【索引38，Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM，arXiv preprint arXiv:2104.04473，2021】、【索引51，Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model，arXiv preprint arXiv:2201.11990，2022】。目标是设计一个重构机制，将 $\Phi(t)$ 从源并行策略 $P_{src}$ 转换到目标策略 $P_{tgt}$。

4 挑战与机遇

与模型并行实现耦合的检查点保存/加载。分布式检查点与特定的模型并行策略高度耦合。例如，数据并行（如PyTorch DistributedDataParallel【索引29，PyTorch Distributed: Experiences on Accelerating Data Parallel Training，arXiv preprint arXiv:2006.15704，2020】）作为最常用的并行策略之一，它在GPU之间复制模型权重和优化器状态，因此只需一个rank（如rank 0）保存整个模型状态到一个检查点文件中。当恢复训练时，每个数据并行工作节点在启动训练迭代前加载相同的检查点文件。而其他并行技术，如ZeRO【索引46，ZeRO: Memory Optimization Towards Training A Trillion Parameter Models，arXiv preprint arXiv:1910.02054，2019】和3D并行【索引38，Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM，arXiv preprint arXiv:2104.04473，2021】，则在GPU之间分片模型参数和优化器状态。由于模型状态在GPU之间是分区的，每个GPU会分别保存和加载只包含其所拥有的部分模型状态的检查点文件。图1展示了ZeRO-3检查点文件的结构，它与其他策略的结构截然不同。

耦合设计带来的重构困难。与模型并行耦合的分布式检查点使得保存/加载流程相对简单，因为对每个工作节点而言，检查点加载过程只是保存过程的逆操作。这种设计还带来了性能优势，因为分布式检查点保存需要阻塞所有并行工作节点，而让每个并行工作节点保存其拥有的模型状态不会产生额外的同步开销，也不会对整体训练速度产生负面影响。然而，这种耦合设计使得将检查点重构为不同的并行策略变得痛苦且容易出错，因为开发者需要手写转换脚本来将检查点文件从一种并行策略转换为另一种。因此，现有框架仅支持对极少数并行策略的重构，例如，在分布式数据并行中更改数据并行度【索引10，PyTorch Distributed Checkpoint - torch.distributed.checkpoint，PyTorch Contributors，2023】，或者仅支持分布式检查点的权重转换用于推理评估（例如，不支持优化器状态的转换）【索引41，Megatron-LM: Evaluation and Tasks，NVIDIA，2024】。因此，如果缺乏支持训练任务在新硬件环境下重构其并行策略的能力，训练过程中的任何节点故障或资源重新分配都将严重影响整个训练进程。

自动化重构管道对广泛并行策略的覆盖有限。表1总结了当前最先进的深度学习训练检查点系统。CheckFreq【索引36，CheckFreq: Frequent, Fine-Grained DNN Checkpointing，FAST’21，2021】是一个高效的检查点系统，它将训练计算与检查点保存重叠，并根据I/O分析在运行时自适应地调整检查点频率。Gemini【索引60，GEMINI: Fast Failure Recovery in Distributed Training with InMemory Checkpoints，SOSP’23，2023】引入了一种内存中检查点技术，将GPU状态检查点到本地和远程CPU，并将检查点IO与训练计算交错以减少开销，令人印象深刻地实现了在每个训练迭代上进行检查点。CheckFreq和Gemini都加速了检查点效率，但不支持对高级模型并行（如TP、PP和ZeRO）的重构。一些同期的工作提出了通过检查点改变并行策略的方法。例如，PyTorch Distributed Checkpointing (DCP) 支持通过加载时重分片来改变并行度，例如，在所有分布式检查点文件加载到CPU后，每个形状为[X/n, Y]的分片参数被连接起来得到完整的[X, Y]参数。然后根据当前的DP度重新分片该张量。类似地，Megatron Dist Checkpointing (MCP) 保存形状为[PP * TP, X/PP, Y/TP]的本地检查点。加载时，这些检查点被连接成全局形状[1, X, Y]，然后根据当前的PP和TP度进行分片。这两种方法都有限，因为它们只支持一部分并行策略。鉴于不同的并行策略在不同硬件上具有不同的性能影响，理想的情况是拥有一个可重构的并行机制，它能覆盖广泛使用的并行策略，并且易于扩展到不同的模型架构。
Table 1: 不同并行策略下重构分布式检查点的支持情况比较。三角（▲）表示部分支持且有限制。

高昂的重构开销。现代DNN的规模通常因其庞大的参数数量而变得巨大，例如，Transformer模型【索引26，Scaling Laws for Neural Language Models，arXiv preprint arXiv:2001.08361，2020】、【索引57，Attention is All you Need，NIPS’17，2017】在过去几年中增长了超过1000倍。著名的例子包括拥有5400亿参数的PaLM【索引9，PaLM: Scaling Language Modeling with Pathways，arXiv preprint arXiv:2204.02311，2022】和拥有5300亿参数的MT-NLG【索引51，Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model，arXiv preprint arXiv:2201.11990，2022】。检查点的大小约是参数数量的12倍（4倍用于模型参数，8倍用于优化器状态），因此这些模型的检查点大小达到了数十TB。鉴于巨大的检查点大小，重构时间可能受到持久化存储带宽的限制【索引14，Check-N-Run: a Checkpointing System for Training Deep Learning Recommendation Models，NSDI’22，2022】、【索引60，GEMINI: Fast Failure Recovery in Distributed Training with InMemory Checkpoints，SOSP’23，2023】。例如，当带宽为20Gbps时，将MT-NLG【索引51，Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model，arXiv preprint arXiv:2201.11990，2022】的模型状态检查点到远程持久化存储需要42分钟。考虑到模型规模的庞大且不断增长，重构过程需要具有相对较低的开销，以避免成为主要的性能瓶颈。

---

A2 方法细节

5.1 UCP 设计概览

UCP系统设计。UCP是一个专为处理具有可重构并行性的大规模深度学习训练而设计的分布式检查点系统。它允许从任意源并行策略灵活高效地重构到任意目标并行策略。由于训练准确性是主要要求，我们专注于那些不会导致训练准确性损失或相关指标下降的设计选择。我们首先对UCP进行概述，并在后续章节中更详细地讨论每个组件。

UCP的核心组件。如图2所示，UCP由三个主要组件构成：原子检查点、基于模式的重构管道和嵌套并行转换。在第4节中，我们描述了与并行策略耦合的检查点所面临的挑战，即这些检查点重构起来非常麻烦。我们通过提升抽象层次来解决这个问题，即通过原子检查点（§5.2）将检查点文件格式与其并行策略及硬件配置解耦。原子检查点使得实现灵活重构的机制变得异常简单，覆盖了各种并行策略选择。UCP还提供了一个基于模式的重构管道（§5.3），通过原子检查点，系统地、自动地将分布式检查点从源并行策略重构到目标并行策略。UCP通过精心设计一套并行模式集，并开发了基于模式感知的重构操作，来自动地将源分布式检查点转换为原子检查点，并将原子检查点映射到任何目标并行策略。UCP利用嵌套并行重构、冗余旁路加载和惰性重构调用来减少重构开销，使得重构对训练过程仅增加极小的开销（§5.4）。

5.2 原子检查点

现有检查点机制的问题。在现有的深度学习分布式训练系统中，不同的并行策略以不同的方式处理模型状态。例如，数据并行（DP）在数据并行工作节点间复制并维持模型状态的一致视图，因此单个GPU进行检查点保存就足够了。而其他并行策略，如张量切片并行或流水线并行，则沿着不同的轴和层将模型状态划分到不同的GPU上。因此，每个GPU都基于其本地内存中分区模型状态的快照创建检查点文件。由于这种高度与并行策略耦合的检查点机制，从业者需要手写转换脚本来为不同的并行策略转换分布式检查点，这带来了巨大的工程成本且难以维护。

原子检查点的定义。在UCP中，我们使用原子检查点来表示模型状态。一个原子检查点包含了一个与张量操作相对应的模型状态的整合视图，例如，如果一个张量操作被分区到多个并行工作节点上，那么原子检查点就是这些碎片状态的合并状态。不同的模型状态（如权重和优化器状态）有各自独立的原子检查点文件，所有与一个参数相关的原子检查点文件被视为该参数的一个原子检查点集合（例如，一个目录）。遵循混合精度训练的标准实践，UCP将优化器状态以FP32格式存储，以保持数值稳定性并避免精度损失。虽然这会增加检查点的大小，但能确保训练过程的正确性。在不失一般性的前提下，假设训练使用Adam优化器【索引27，Adam: A Method for Stochastic Optimization，ICLR’15，2015】，对于每个张量操作，其参数的原子检查点集合包含三个原子检查点文件：
* model.pt: 一个代表参数权重的fp32张量。
* adam_m.pt: 一个代表一阶矩的fp32张量。
* adam_v.pt: 一个代表二阶矩的fp32张量。

原子检查点的特性。原子检查点比大多数现有的分布式检查点（这些检查点包含单个并行工作节点的完整本地快照）要精细得多。原子检查点不再与任何并行策略或硬件配置耦合，例如，它不包含任何rank ID、分区信息或特定分片策略（如为对齐而进行的填充）带来的额外数据。它仅仅是将模型中的每个参数映射到其权重和相应优化器状态的整合视图。

原子检查点的优势。原子检查点足以映射到广泛的并行算法，并且它通过检查点实现了灵活的并行策略重构。关键在于，当每个数据并行工作节点需要以逐个参数的方式加载模型状态的副本时，这种表示方式天然就是数据并行的。此外，正如我们在许多并行策略中识别出的分区模式（§5.3.1），定义一组原语来自动将原子检查点转换为特定的并行策略成为可能（§5.3.2）。因此，UCP充当了不同分布式训练技术之间的通用交换格式，用户无需为每种源到目标的并行配置实现单独的转换器，从而极大地增加了可重构并行的灵活性和覆盖范围。

对数据类型的特殊考虑。在实践中，开发者通常采用硬件友好的训练技术，如混合精度训练（MPT）【索引35，Mixed Precision Training，ICLR’18，2018】，其中同时维护IEEE float16权重和float32权重及优化器状态，以利用TensorCore【索引32，NVIDIA Tensor Core Programmability, Performance & Precision，IPDPSW’18，2018】等硬件特性的高计算吞吐量。MPT影响了可重构并行的设计，因为不同状态可能有不同的数据格式。此外，动态混合精度训练会以一种数据格式（如float16）保存检查点，并在不同训练阶段切换到另一种格式（如bfloat16或float32），这些格式在尾数和指数之间的位分配策略不同【索引25，A Study of BFLOAT16 for Deep Learning Training，arXiv preprint arXiv:1905.12322，2019】。因此，考虑检查点的数据格式以适应这些训练技术至关重要。在UCP中，我们将所有原子检查点（权重/优化器值）保持为float32格式，并支持以其他数据类型恢复训练，包括float32、float16和bfloat16，使其能够灵活地用不同数据类型重构并行策略。

5.3 基于模式的重构管道

自动化转换的挑战。原子检查点与特定的并行策略解耦。然而，另一个重要的设计选择仍未明确：如何自动地将各种源并行策略 $P_{src}$ 的分布式检查点转换为原子检查点，并将原子检查点映射到不同的目标并行策略 $P_{tgt}$？最重要的是，如何确保这个自动转换过程能覆盖广泛使用的并行策略以及来自最先进模型架构的新并行策略？这里的设计选择不会改变原子检查点的含义，但它们对于确保最终重构系统的实用性至关重要。为了回答这些问题，本节我们开发了一个基于模式的重构管道。

5.3.1 设计模式集

模式设计的原则。为了支持基于原子检查点的灵活并行重构，我们寻求一种能够自动将分布式检查点映射到原子检查点的方法。为此，我们研究了分布式检查点文件内部的模式，以便执行基于模式的转换。这里的关键考虑是如何设计和选择这些模式。好的模式应具备两个关键属性：覆盖性和灵活性。覆盖性是理想的，因为它能够将来自广泛并行策略和模型架构的分布式检查点映射到原子检查点。灵活性则便于后续的并行重构操作。基于这两个关键属性，我们提出了以下模式集。

UCP定义的模式。图3展示了UCP中预定义的模式。Unique表示一个张量唯一地与一个分布式检查点关联，这在诸如流水线并行等算子间并行生成的检查点中很常见。Replicate表示一个参数将在多个分布式检查点中被复制。例如，即使在张量并行中矩阵乘法在GPU间被分片，其他参数如LayerNorm和偏置项在GPU间是复制的，因为它们不是计算和内存密集型的。因此，如果我们将所有这些复制的参数都转换为原子检查点，将会导致冗余。Partial表示一个参数在GPU间独立更新，这对应于异步训练的模型。更复杂的模式存在于算子内并行策略中，这些策略沿着特定维度对参数进行分区，例如，张量并行中的Shard-V（列式分片）和Shard-H（行式分片），Shard-Hy（沿模型状态的多个维度分片），以及Shard-NC（非连续地分片模型状态）。

模式集的通用性。我们稍后将展示，我们选择的这套模式能够转换复杂的并行策略，如ZeRO-3【索引46，ZeRO: Memory Optimization Towards Training A Trillion Parameter Models，arXiv preprint arXiv:1910.02054，2019】和3D并行【索引38，Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM，arXiv preprint arXiv:2104.04473，2021】，同时足够灵活，也能支持新的模型架构，如混合专家模型（MoE）【索引23，Mixtral of Experts，arXiv preprint arXiv:2401.04088，2024】和分组查询注意力（GQA）【索引5，GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints，EMNLP’23，2023】。请注意，尽管这套模式覆盖了许多现有的并行策略，但这种方法是通用的，并且可以扩展以支持新的并行模式，而不仅限于本文中考虑的那些。

5.3.2 模式感知的重构操作

核心重构操作。给定模式集，UCP从分布式检查点中提取模式信息，并启动模式感知的重构操作。具体来说，表2列出了UCP中的主要重构操作，包括Extract、Union、StripPad、UcpInfo和Load。对于不同的模式，每个操作执行模式感知的转换。算法1演示了Extract和Union如何根据参数碎片的模式来整合它们。我们使用几个具体例子来说明UCP中的模式感知重构操作如何支持复杂的并行策略和不同的模型架构。
Table 2: UCP中的模式感知重构操作。

案例分析：ZeRO Stage 3。图4展示了将UCP的重构管道应用于ZeRO-3检查点的过程。如果ZeRO-3是源并行策略 $P_{src}$，每个数据并行（DP）rank会将其拥有的分片参数和优化器状态持久化到一个检查点文件中，当DP度为N时，总共会产生N个分布式检查点文件。UCP接着识别ZeRO-3的并行模式：ZeRO-3将每个张量展平为一维张量，并将该一维张量在不同的数据并行rank间平均分片。UCP将ZeRO-3分布式检查点文件中的这些张量识别为Shard-V模式。基于该模式，UCP对碎片参数运行Extract和Union操作来创建原子检查点，其中包含了参数和优化器状态的整合视图。ZeRO-3的一个复杂之处在于它会添加填充（padding）以确保一维张量可以被N个DP工作节点均匀整除以实现对齐。例如，对于一个形状为[1024]的参数，如果原始DP度为3，ZeRO-3会将其填充至[1026]，这样每个工作节点拥有的分片大小为[342]。UCP使用StripPad来移除填充，并使用Save将得到的原子检查点保存到持久化存储。
Algorithm 1 模式感知的重构

ZeRO Stage 3作为目标策略。另一方面，如果ZeRO-3是目标并行策略 $P_{tgt}$，每个DP工作节点会通过UcpInfo计算其新的分区元数据，然后按照层的顺序依次将原子检查点加载到每个rank，并为高性能添加对齐填充。使用前面的例子，假设 $N_{tgt} = 2$，UCP将恢复训练，每个rank拥有正确的参数形状[512]而不是[513]。一旦所有分区状态被加载到GPU中，例如加载到ZeRO-3的fp32_partitioned_groups_flat这个扁平化内存属性中，更新后的属性会被广播到其他必要的属性，如用于MPT的fp16_partitioned_groups_flat。通过这个过程，UCP在各种DP度下，能够灵活地将ZeRO-3作为 $P_{src}$ 或 $P_{tgt}$ 进行并行重构。

案例分析：3D并行。将3D并行作为 $P_{src}$ 或 $P_{tgt}$ 的重构过程与ZeRO-3类似。然而，由于3D并行由多种并行策略组成，它有其自身的复杂性。与ZeRO-3相似，在为3D并行保存分布式检查点时，每个GPU只保存其所拥有的模型状态的一个切片。但是，3D并行的参数模式要复杂得多：参数可以有Replicate模式（如LayerNorm），Shard-V和Shard-H模式（如矩阵乘法），当TP度 > 1时有Partial模式（如Alibi Embedding【索引44，Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation，arXiv preprint arXiv:2108.12409，2022】）和Replicate模式（如TiedEmbedding），以及当PP度 > 1时有Unique模式（大多数情况）。UCP通过模式匹配并使用Extract、Union、StripPad来创建无任何填充的整合原子检查点，极大地减轻了手动转换3D并行分布式检查点的负担。例如，根据是Shard-V还是Shard-H，UCP会相应地将来自TP的碎片参数沿行维度或列维度连接成一个单一的整合张量。与ZeRO-3类似，如果选择3D并行作为 $P_{tgt}$，首先会通过UcpInfo生成原子检查点和GPU rank之间的新映射，然后每个rank根据新的映射策略从原子检查点加载数据。

3D并行的额外复杂性。3D并行的额外复杂性还来自于流水线调度。例如，交错式1F1B（Interleaved 1F1B）【索引38，Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM，arXiv preprint arXiv:2104.04473，2021】是一种通过将模型的非连续层分配给同一个rank来减少流水线气泡的调度策略，这导致分布式检查点与特定的层分配顺序高度耦合。UCP的原子检查点与这种调度解耦，因为参数和优化器状态只与它们各自的参数名关联，从而允许重新配置每个rank与参数之间的映射。

案例分析：稀疏混合专家模型与不规则注意力机制GQA。图3中描述的模式不仅覆盖了密集型LLM的模型分片策略，如ZeRO-3和3D并行的例子所示，它还覆盖了稀疏MoE模型【索引23，Mixtral of Experts，arXiv preprint arXiv:2401.04088，2024】和不规则注意力机制如GQA【索引5，GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints，EMNLP’23，2023】等模型架构。这两种架构在LLM训练中受到了广泛关注，但也需要更复杂的并行策略。图5展示了其复杂性。这个例子中的MoE模型将MoE的FFN层的权重张量定义为[专家数量 × 输出隐藏层维度, 输入隐藏层维度]，这是一个融合的权重矩阵，不同于使用独立矩阵表示不同专家的标准MoE。有了这个融合矩阵，可以对该层应用TP。然而，与标准TP不同，分区是沿着输出隐藏层维度进行的，这不属于常见的Shard-V和Shard-H模式。在UCP中，这通过Shard-NC模式来处理，它允许UCP将其识别为一个3维张量，并应用模式感知的重构操作，如Extract和Union，从而仍然能从分布式检查点中获得融合矩阵的整合视图。

GQA的复杂性。在GQA【索引5，GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints，EMNLP’23，2023】的例子中，多头注意力（MHA）中的QKV矩阵也被融合成一个张量。然而，与传统MHA不同，GQA中的QKV大小不同，导致融合矩阵的形状不规则，为[q_size + k_size + v_size, hidden]，其中q_size != k_size == v_size。如果对这个张量应用张量并行（TP），它需要沿着第一个维度对每个Q、K和V进行分区，但分区大小不同。Shard-NC模式连同形状信息使得UCP能够识别这些可变大小的片段，并相应地应用模式感知的重构操作。

管道的灵活性总结。这些例子表明，UCP中基于模式的重构管道非常灵活，支持将各种并行策略作为 $P_{src}$ 或 $P_{tgt}$。同时，其精心设计的模式集使其不仅能支持广泛的复杂并行策略，还能很好地扩展到稀疏混合专家模型和具有不规则注意力的模型，以及密集模型。引入一个涉及超出UCP当前范围模式的新并行策略，将需要与上述例子相似的实现工作量。我们积极与社区合作，识别这些新模式，以扩展UCP的覆盖范围。

5.4 高效重构

可扩展性要求。UCP旨在用于高性能训练，因此§5.3中描述的操作需要能够很好地扩展到更大的模型和更多的节点。为实现低重构成本，UCP利用了三种优化：（1）嵌套并行重构，（2）冗余旁路加载，以及（3）惰性重构调用。

嵌套并行重构。受利用额外节点以提高重构效率的启发，我们将重构过程构建为一个MapReduce问题【索引12，MapReduce: Simplified Data Processing on Large Clusters，Communications of the ACM，2008】，从而引入了一种嵌套的并行重构方案。MapReduce最初是作为一种在分布式环境中处理大规模数据集的编程模型被引入的，并已广泛应用于Hadoop【索引50，The Hadoop Distributed File System，MSST’10，2010】等大数据处理框架中。MapReduce将一个问题分解为三个阶段：Mapper、Shuffler和Reducer。在我们的案例中，如图6所示，每个并行mapper读取一个分布式检查点，并根据其包含的张量碎片解析它。shuffler将mapper的输出，即代表参数碎片的张量，发送给被分配处理这些参数的相应reducer。reducer根据特定的模式（如我们§5.3中详细讨论的Shard和Partial）处理这些参数碎片。

负载均衡。需要注意的是，参数的大小差异巨大。例如，在一个176B的模型中，嵌入参数的大小为[50257, 12288]，而LayerNorm的偏置大小为[12288]。随机地将这些参数分配给工作节点会导致不平衡，因为计算时间和加载/保存时间都与参数大小成正比。这种不平衡会导致大多数工作节点空闲地等待处理最重任务的那个节点，从而主导了总转换时间。为解决此问题，我们引入了一种基于预计算的均衡方法。主工作节点根据参数元素数量（numels）将参数分成N组。例如，有4个工作节点时，主节点可能会将嵌入参数分配给工作节点1，同时将多个较小的参数（如LayerNorm的偏置和权重）分散到工作节点2、3和4，以使每个工作节点的总numel数相似。

嵌套并行的第二层。在工作节点间实现参数均衡分配的基础上，我们实现了第二层并行，即在每个工作节点内部利用多核处理。这种方法通过调动所有可用的CPU核心来并行处理参数，从而充分利用了工作节点内的计算能力。这种嵌套并行策略——先在工作节点间均衡，再在每个工作节点的核心间并行——最大化了资源利用率，并在节点和核心两个层面都最小化了空闲时间，从而为大型语言模型带来了可接受的转换时间。

冗余旁路加载。同一DP组内的每个工作节点共享或部分共享（除了ZeRO-3 DP）相同的模型状态，因此需要加载相同的原子检查点文件。UCP通过在同一DP组的工作节点间均匀分配加载工作负载来消除冗余的检查点加载。每个rank读取其分配的原子检查点文件到CPU内存，将它们传输到GPU内存，然后使用all-gather操作将数据分发给其他工作节点。这种优化通过消除从存储到CPU内存的冗余数据传输，显著减轻了IO带宽压力，并利用高带宽的GPU间互连（如NVLINK，速度达900 GB/s）在数据加载到GPU后重新分发。此外，UCP采用了内存高效的加载设计。UCP逐层加载原子检查点，并根据所选的并行策略将张量分配给GPU。一旦某一层被传输到GPU内存，其张量就会从CPU内存中释放，从而将峰值CPU内存使用量从完整模型检查点的大小降低到单层的大小。

惰性重构调用。细心的读者可能会想：“并行重构应该在何时触发？”或者“在保存检查点时，是否可以直接将分布式模型状态整合到一个单一的检查点文件中？”这种方法有两个缺点：将分布式模型状态整合到单一检查点会不可接受地减慢训练速度，并且在极端规模下是不切实际的，因为可能没有足够的内存来容纳整合后的模型状态。为了避免减慢正常训练，我们为UCP采用了惰性调用设计。UCP的重构仅在源并行策略 $P_{src}$ 和目标并行策略 $P_{tgt}$ 不同或硬件发生变化时才被调用。这样，分布式检查点保存逻辑无需任何更改，而当需要进行并行重构时，UCP仅产生低于0.001%的训练时间开销，因为重构不会增加关键路径长度，即正常的分布式训练过程。

---

A4 实验环境

• 工作负载：

  • 模型架构：实验主要围绕GPT风格的Transformer模型进行，覆盖了多种规模和类型：
    • GPT-3 medium (3.5亿参数): L=24, H=1024, A=16
    • GPT-3 7B (70亿参数): L=32, H=4096, A=32
    • GPT-3 风格 176B (1760亿参数): L=70, H=14336, A=112
    • Mixtral-7x8B 风格 MoE (420亿参数): L=32, H=3072, A=32, E=16
  • 数据集：使用Pile数据集【索引15，The Pile: An 800GB Dataset of Diverse Text for Language Modeling，arXiv preprint arXiv:2101.00027，2020】的一个子集进行训练，以评估UCP重构并行策略对训练损失的影响。

• 硬件配置：

  • 主要集群：64块 NVIDIA A100 40GB GPU，每节点配备256GB DRAM，10TB存储，200Gbps互联。
  • 176B模型集群：384块 NVIDIA A100 80GB GPU。
  • 1T模型效率评估集群：1024块 AMD MI250X 64GB GPU。
  • I/O带宽：NVIDIA GPU集群测得5 GB/s，AMD GPU集群测得3 GB/s，与消费级SSD或典型云服务商提供的带宽相当。

• 软件配置：

  • 实现：UCP在DeepSpeed框架中实现。

---

A4 实验结果

可重构并行的有效性验证

• 实验内容：验证UCP在重构并行策略后能否保持训练的准确性。
  1. 单源到多目标：使用一种源并行策略（TP=2, PP=2, DP=2, SP=1）训练GPT-3模型100个迭代，然后在第100个迭代时使用UCP将检查点恢复到多种不同的目标并行策略和硬件配置上，继续训练。
  2. 多源到单目标：使用多种不同的源并行策略训练GPT-3模型100个迭代，然后将它们在第100个迭代保存的检查点全部重构到同一种目标并行策略（TP=2, PP=2, DP=1, SP=1）上，继续训练。
• 实验结果：
  • 在“单源到多目标”实验中，所有从不同目标策略恢复的训练，其损失曲线都与源策略的预期收敛轨迹无缝衔接，保持一致（图7a）。
  • 在“多源到单目标”实验中，无论源配置如何，所有恢复后的训练任务都收敛到相同的损失曲线上，并且与源训练曲线在第101-200次迭代中完全匹配（图7b）。
• 分析结论：实验结果证实，UCP能够有效地在不同硬件和并行配置之间重构训练，同时完全保持模型的训练准确性和收敛性。这得益于原子检查点作为通用中间表示，以及基于模式的自动化重构管道。
  

  

与现有方法的比较

• 实验内容：将UCP与PyTorch DCP【索引20，Getting Started with Distributed Checkpoint (DCP)，pytorch.org，2024】和Megatron MCP【索引33，Megatron-Core Dist Checkpointing Package，docs.nvidia.com，2024】进行比较，评估它们在四种场景下的灵活性：（a）改变DP度，（b）切换到ZeRO-DP，（c）改变MP度（TP和PP），（d）应用于MoE模型。
• 实验结果：
  • DCP支持改变DP度和切换到ZeRO-DP，但在改变MP度时失败。
  • MCP支持改变DP和MP度，但在MoE模型和切换到ZeRO-DP时存在限制或失败。
  • UCP在所有四种场景下都成功恢复了训练，没有出现系统级错误（图8）。
• 分析结论：UCP相比现有的先进检查点系统，支持更广泛、更复杂的并行策略重构。
  

  

对不同模型架构的泛化能力

• 实验内容：在多种不同规模和架构的模型上验证UCP的泛化能力，包括GPT-3 7B、GPT-3 176B和稀疏MoE 42B模型。
• 实验结果：对于所有测试模型，UCP都能在训练中途无缝地恢复并切换到新的并行策略，训练收敛曲线与初始阶段保持一致（图9）。
• 真实案例：GPT-3 176B的训练曲线对应于BLOOM 176B【索引52，BLOOM: A 176B-Parameter Open Access Multilingual Language Model，arXiv preprint arXiv:2211.05100，2023】训练中的一个真实事件。当时训练集群从48个节点缩减到24个，UCP成功地在小规模集群上恢复了训练，避免了三个月训练成果的损失。
• 分析结论：UCP与模型架构无关，其全面的模式设计使其能够广泛适用于密集和稀疏等不同类型的模型。
  

  

重构效率评估

• 实验内容：评估UCP的保存成本和重构成本。
  1. 保存成本：由于采用惰性调用，UCP不改变标准的分布式检查点保存逻辑，因此无额外保存开销。
  2. 重构成本：分为两部分：从分布式检查点到原子检查点的转换开销，以及加载原子检查点的开销。
• 实验结果：
  • 转换成本：得益于嵌套并行重构，即使模型规模扩展到1T参数，转换时间也始终控制在约3分钟以内（图10）。与顺序方法相比，UCP的嵌套并行方法实现了高达257倍的时间缩减。
  • 加载成本：与加载标准分布式检查点相比，加载原子检查点仅增加约10秒的固定开销，这在不同模型规模下保持一致（图11）。冗余旁路加载技术根据数据并行度的不同，可带来3倍至20倍的加载速度提升。
  • 端到端成本：对于1T参数模型，总的端到端重构时间（保存、转换、加载）保持在5分钟以下，与数天甚至数月的训练时间相比可以忽略不计（表3）。
• 分析结论：UCP通过多项系统级优化（嵌套并行、冗余旁路加载、惰性调用）将重构开销降至最低，使其在实际应用中非常高效。
  

  

  
  

  

  
  Table 3: 不同模型和硬件配置下的端到端重构开销。
  

  

---

A7 补充细节

7 相关工作

检查点系统。检查点被广泛用作故障恢复的基本机制。虽然朴素的检查点方法会暂停训练，但异步检查点方法允许磁盘I/O和模型操作重叠【索引36，CheckFreq: Frequent, Fine-Grained DNN Checkpointing，FAST’21，2021】、【索引39，DeepFreeze: Towards Scalable Asynchronous Checkpointing of Deep Learning Models，CCGRID’20，2020】、【索引40，VeloC: Towards High Performance Adaptive Asynchronous Checkpointing at Large Scale，IPDPS’19，2019】。即时（JIT）检查点【索引16，Just-In-Time Checkpointing: Low Cost Error Recovery from Deep Learning Training Failures，EuroSys’24，2024】利用数据并行（DP）的状态冗余，仅在故障后创建检查点。为了解决远程持久存储低带宽的瓶颈，FastPersist【索引59，FastPersist: Accelerating Model Checkpointing in Deep Learning，arXiv preprint arXiv:2406.13768，2024】通过扩展本地NVMe SSD实现更高带宽，而内存中检查点方法【索引60，GEMINI: Fast Failure Recovery in Distributed Training with InMemory Checkpoints，SOSP’23，2023】则利用远程CPU内存，允许通过高速训练网络进行检查点。这些技术主要关注优化检查点保存效率，该环节位于整体训练过程的关键路径上。UCP不改变保存逻辑，与这些优化技术是正交的，可以无缝集成。

Spot实例训练。已有研究探讨使用spot实例进行DNN训练，因为它们成本效益更高。然而，这种动态环境使得重构开销成为一个关键因素。UCP已经引入了多项优化，如嵌套并行重构、冗余旁路加载、惰性重构调用，将重构成本从数小时降低到几分钟，使其成为在spot实例上进行更具弹性训练的潜在可行方案。然而，先前的工作，如Bamboo【索引55，Bamboo: Making Preemptible Instances Resilient for Affordable Training of Large DNNs，NSDI’23，22023】和Parcae【索引13，Parcae: Proactive, Liveput-Optimized DNN Training on Preemptible Instances，NSDI’24，2024】，表明许多云spot实例的动态性可能极高。因此，可能需要更高效的并行可重构性来适应这些环境。我们认为这是一个非常有趣的未来研究方向。

可重构机器学习系统。VirtualFlow【索引43，VirtualFlow: Decoupling Deep Learning Models from the Underlying Hardware，arXiv preprint arXiv:2009.09523，2021】和Singularity【索引49，Singularity: Planet-Scale, Preemptive and Elastic Scheduling of AI Workloads，arXiv preprint arXiv:2202.07848，2022】将深度学习任务与物理设备解耦，从而实现了运行时的灵活GPU映射，但需要复杂的驱动级虚拟化，并且不支持多维度的并行策略变更。Tenplex【索引58，Tenplex: Dynamic Parallelism for Deep Learning using Parallelizable Tensor Collections，SOSP’24，2024】是一个状态管理库，它将状态描述为一个可并行的张量集合（PTC），并在运行时为任务变更生成重构计划。虽然Tenplex支持改变部分并行策略，但它仅限于集群内的重构，并且将分区参数存储在主机内存中会产生内存压力，并因非优化的状态恢复而影响训练准确性。UCP支持跨集群迁移，同时完全保持训练准确性。

---

A5 结论

本文介绍了通用检查点（UCP），一个用于大规模DNN训练的、支持可重构并行的检查点系统。UCP的主要目标是通过分布式检查点实现并行策略的灵活重构。为此，UCP提供了一种名为“原子检查点”的新型检查点结构，并构建了一个基于模式的重构管道，该管道能够自动、系统地重构广泛使用的并行策略，同时保持模型准确性。为了降低重构成本，UCP引入了嵌套并行重构和惰性调用等优化，使得重构开销相对于端到端的训练成本可以忽略不计。我们对大规模LLM训练的评估表明，UCP能够以低成本重构一个比现有方案更广泛、更复杂的并行策略集合。UCP已通过一个PyTorch库实现并开源，用于加速大规模DNN训练，其有效性已在真实世界的大规模LLM预训练任务中得到验证。

---

方法细节中引用的参考文献汇总

以下是在“方法细节”（包含背景、挑战和UCP设计章节）中引用的参考文献列表，及其在文中的引用方式和描述：

1. [5] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. 2023, EMNLP’23.

   • 引用段落: §5.3.2
   • 引用描述: 用于说明UCP的模式集可以覆盖不规则的注意力机制，例如GQA。文中提到GQA中融合的QKV矩阵大小不同，UCP通过Shard-NC模式来处理这种可变大小分区的重构。

2. [12] MapReduce: Simplified Data Processing on Large Clusters. 2008, Communications of the ACM.

   • 引用段落: §5.4
   • 引用描述: UCP将重构过程构建为一个MapReduce问题，以引入嵌套并行重构方案，提高效率。

3. [18] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. 2019, NIPS’19.

   • 引用段落: §2, §3
   • 引用描述: 作为流水线并行（PP）的背景知识引用，说明PP是将模型状态划分为多个阶段部署在不同GPU上的并行策略。

4. [23] Mixtral of Experts. 2024, arXiv preprint arXiv:2401.04088.

   • 引用段落: §5.3.2
   • 引用描述: 用来说明UCP的模式集能够支持如稀疏混合专家（MoE）模型这样的新模型架构，并处理其复杂的并行策略。

5. [25] A Study of BFLOAT16 for Deep Learning Training. 2019, arXiv preprint arXiv:1905.12322.

   • 引用段落: §5.2
   • 引用描述: 在讨论数据类型时引用，说明不同的浮点格式（如bfloat16）在位分配策略上不同，UCP需要考虑这些以支持动态混合精度训练。

6. [27] Adam: A Method for Stochastic Optimization. 2015, ICLR’15.

   • 引用段落: §5.2
   • 引用描述: 在介绍原子检查点结构时，以Adam优化器为例，说明每个参数会保存权重、一阶矩和二阶矩三个文件。

7. [29] PyTorch Distributed: Experiences on Accelerating Data Parallel Training. 2020, arXiv preprint arXiv:2006.15704.

   • 引用段落: §3, §4
   • 引用描述: 作为数据并行（DP）的背景知识引用，并说明在DP中，模型状态被复制，导致检查点通常由一个rank保存。

8. [30] Sequence Parallelism: Long Sequence Training from System Perspective. 2022, arXiv preprint arXiv:2105.13120.

   • 引用段落: §2, §3
   • 引用描述: 作为序列并行（SP）的背景知识引用，说明SP是为训练长序列LLM而开发的并行策略。

9. [32] NVIDIA Tensor Core Programmability, Performance & Precision. 2018, IPDPSW’18.

   • 引用段落: §5.2
   • 引用描述: 在讨论混合精度训练时引用，说明MPT利用TensorCore等硬件特性来提高计算吞吐量。

10. [35] Mixed Precision Training. 2018, ICLR’18.

    • 引用段落: §5.2
    • 引用描述: 作为混合精度训练（MPT）的背景知识引用，说明UCP的设计需要考虑MPT中不同状态的不同数据格式。

11. [37] PipeDream: generalized pipeline parallelism for DNN training. 2019, SOSP’19.

    • 引用段落: §2, §3
    • 引用描述: 与[18]一同作为流水线并行（PP）的背景引用，并提及1F1B等流水线调度策略。

12. [38] Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. 2021, arXiv preprint arXiv:2104.04473.

    • 引用段落: §2, §3, §4, §5.3.2
    • 引用描述: 多次引用，作为3D并行的代表性工作，并提及交错式1F1B等高级流水线调度。在方法部分，用来说明3D并行的检查点模式复杂，以及UCP如何处理其调度带来的额外复杂性。

13. [44] Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. 2022, arXiv preprint arXiv:2108.12409.

    • 引用段落: §5.3.2
    • 引用描述: 在分析3D并行模式时，以Alibi Embedding为例，说明其参数在TP>1时具有Partial模式。

14. [46] ZeRO: Memory Optimization Towards Training A Trillion Parameter Models. 2019, arXiv preprint arXiv:1910.02054.

    • 引用段落: §2, §3, §4, §5.3.2
    • 引用描述: 多次引用，作为ZeRO数据并行的背景，并详细阐述了ZeRO-3检查点的分片特性（图1）以及UCP如何通过Shard-V模式和StripPad操作来重构ZeRO-3检查点（图4）。

15. [48] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2020, arXiv preprint arXiv:1909.08053.

    • 引用段落: §2, §3
    • 引用描述: 作为张量并行（TP）的背景知识引用，说明TP通过对张量操作进行分区来训练大型模型。

16. [50] The Hadoop Distributed File System. 2010, MSST’10.

    • 引用段落: §5.4
    • 引用描述: 在介绍嵌套并行重构时，以Hadoop为例，说明MapReduce模型在大数据处理框架中的广泛应用。

17. [51] Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model. 2022, arXiv preprint arXiv:2201.11990.

    • 引用段落: §2, §3, §4
    • 引用描述: 作为3D/4D并行和大规模模型训练的例子引用，并提及MT-NLG 530B模型及其巨大的检查点大小，以强调高效重构的重要性。