摘要

原文：https://arxiv.org/pdf/2202.07848.pdf

Singularity，微软的全局分布式调度服务，用于高效可靠地执行深度学习训练和推理工作负载。Singularity 的核心是一种新颖的、负载感知的调度器，它可以透明地抢占深度学习工作负载以及进行弹性伸缩，在不影响AI加速器（如 GPU、FPGA）的正确性和性能的前提下，提高利用率。

Singularity中的所有作业默认都是可抢占、可迁移、可动态调整（弹性）大小的：一个运行的作业可以动态透明地：

(a) 被抢占并迁移到不同的节点、集群、数据中心或区域，并从执行被抢占的位置恢复；

(b) 在给定类型的不同加速器上弹性伸缩resize；

我们的机制是透明的，因为它们不需要用户对代码做任何更改，也不需要使用任何可能限制灵活性的自定义库。此外，我们的方法显著提高了深度学习工作负载的可靠性。

结果表明，Singularity 提高了系统的效率和可靠性，对稳态性能的影响可以忽略不计。而且，Singularity 的设计与 DNN 架构无关，可以处理各种并行策略（例如，数据/管道/模型并行）。

 

1. 简介

Singularity 核心实现了 AI 任务的可抢占、可迁移、可动态调整，并且该实现与模型架构无关、与模型训练的并行策略无关，可以认为做到了用户无感。

1.1. 设计目标

Singularity 为了最大化整个集群的吞吐量，采用以下设计原则：

1. 不闲置资源：Singularity 将整个加速器集群视为单个逻辑共享集群，并避免任何资源碎片化或静态容量预留。Singularity 适时调度使用全球范围内的任何空闲资源，跨集群、AZ和工作负载边界（训练与推理）。
2. 提供作业级别的 SLA：在适时使用空闲容量的同时，Singularity 通过遵守作业级别的 SLA 来提供隔离。例如，当推理作业的负载增加时，Singularity 通过弹性缩小或抢占训练作业来释放容量。
3. 故障弹性恢复：DNN 训练作业运行时间长达数小时、数天甚至数周，因此从头开始成本损失巨大。在 Singularity 中，作业从它们被抢占的地方恢复，从而将故障重启成本最小化。

1.2. 关健机制

为实现上面的目标，整个 Singularity 系统的底层由两大重要的机制来支撑。它们分别是：

1) 抢占和迁移机制：Singularity 可以透明地设置检查点、抢占和迁移节点间甚至跨集群和区域的所有 DNN 作业。检查点是通过高效的的同步屏障 (synchronization barrier) 来实现分布式作业的所有参与者之间分布式状态的 一致性切分 (consistent cut)。

2) 伸缩和弹性机制：Singularity 使所有工作能够使用可变数量的 AI 加速器，以透明的方式 动态地、弹性地 伸缩资源。

Singularity 系统架构

1. 这里涵盖了所有的 AI 算力，包括 GPU、FPGA、CPU、ASIC 等不同的硬件形态。
2. 所有的算力资源都被容器化了
3. 硬件抽象层（HAL）竟然 在一层基础软件之上，这层基础软件包括 NVML、NCCL、CUDA，也就是设备管理、设备通信、设备计算这三类功能。
4. 硬件抽象层这里的 CUDA 指的是 CUDA Driver API f. 核心调度原语。高层原语包括：Failover、Suspend、Resume、Migrate、Scale Up/Down ，底层原语包括：Checkpoint、Restore、Distributed Barrier 。

Singularity是一个分布式系统，它的调度程序具有透明抢占、迁移和弹性功能。

在Singularity中，采用了新的方法称为 “自动解耦”，把作业和设备资源解耦，作业与加速器资源之间的绑定是动态的，并且在作业的生命周期中不断变化。为了扩展缩容，我们只需更改工作线程映射到的设备数量即可。这对用户来说是完全透明的，因为作业的总工作线程数不受运行作业的物理设备数量的影响。（训练任务的资源伸缩只改变物理 devices 数量，不改变 world-size，也就是 workers 的总数不变）。

Singularity调度程序的核心机制。为了实现这种解耦，Singularity引入了设备代理（device proxy）的概念。设备代理在自己的地址空间中运行，并且与物理加速器设备具有一一对应关系。当作业worker发起 device API时，它们被拦截并发送到共享内存中的设备代理进程，其生命周期与worker的生命周期解耦。

这种分离实现了两个关键好处：

(a)主机地址空间保持干净，没有由GPU库（如CUDA）创建的设备特定映射和其他副作用，从而使使用现有工具CRIU做检查点和迁移主机进程更容易；

(b)它允许在同一设备上对多个worker进行动态、透明的时间片切片，由设备代理在worker之间执行多路复用和调度。

Singularity使用一种称为“副本切分”的新技术，使得可以在同一设备上对多个工作线程进行时间片切片，而几乎没有开销，同时使每个工作线程都可以使用整个设备内存。副本切分依赖于领域知识，以利用分布式训练作业中特定程序执行点处工作线程之间的内存内容相似性。（资源伸缩，采用的核心技术是 replica splicing，实现了分时多任务（time-slice multiple workers）

 

2. 关健机制概述

为了提高利用率和可靠性，Singularity引入了新的机制，使所有作业默认情况下都可以抢占和动态伸缩。Singularity对分布在多个节点上并采用各种并行方式（如数据、管道或张量并行）的作业进行一致的检查点，并可以使用透明时间切片在稍后的时间点上在可能不同数量的设备上恢复作业。

这些机制的三个关键方面：(a)用户透明（不改变或限制用户代码），(b)持续运行（作业从先前抢占时停止执行的程序执行点恢复），以及 (c)解耦执行。

 

2.1. 用户透明

现有系统实现检查点创建和弹性伸缩一般都需要用户进行代码改动或者使用定制的库。

第一种方案的缺点是，将大量的改造负担留给了用户，例如复杂状态的保存和恢复，资源伸缩后的超参调整等。

第二种方案的缺点是，框架托管了训练的运行以方便创建检查点，这种实现剥夺了用户的灵活性，以至于在实际场景中很少采用。

当前的现状是，大多数 DNN 训练任务很难容错和动态伸缩。

用户透明之所以如此重要，基于如下两点判断：

1) 让调度的归调度，Checkpoint和Restore作为调度原语中的一等公民，在底层保证了任务的 SLAs

2) 让用户专注于代码不受束缚，这些机制对于用户完全透明，复杂性屏蔽在机制以下而不是让用户承担

2.2 持续运行 work-conserving

Singularity 提供的 Checkpoint 可以完整记录程序状态的快照，包括：指令指针、栈、堆内存等，任务可以完全恢复到被抢占前的状态，而不像当前既有方案一样恢复到上一个 epoch 的检查点。

 

2.3. 解耦执行

提高工作可靠性和集群效率/利用率的目标非常协同。Singularity调度程序通过将作业与底层资源之间的映射解耦来实现这些目标。Singularity调度程序透明地虚拟化了world size和rank分配。这种解耦对于透明地检查点和抢占作业以及随后在不同节点、集群、数据中心从先前检查点状态恢复它们，使用相同或不同数量的GPU，至关重要。在作业稳态执行期间，通过与GPU的交互，Singularity调度程序还透明地解耦了作业工作进程的训练逻辑（例如，使用PyTorch编写）。

解耦执行，这点对实现透明的抢占和迁移也至关重要。

 

3. 拦截机制/API Filters

Singularity 采用 Linux 系统下的 LD_PRELOAD 环境变量，实现动态连接库的调用劫持。

为了提供透明的检查点和弹性，Singularity利用了在CPU上执行和在加速器（如GPU）上执行之间存在的窄接口。与加速器的任何交互都必须通过特定的库（例如，NVIDIA GPU的CUDA，AMD GPU的ROCm等）进行，Singularity通过LD_PRELOAD机制动态拦截这些库。大多数功能都驻留在一个称为设备代理的组件中。设备代理可以被视为加速器设备的硬件抽象服务，并且具有一个服务器组件（每个设备一个）和嵌入在与设备交互的每个进程中的客户端组件。主机调用的所有加速器特定API都被拦截并发送到运行在隔离地址空间中的设备代理服务端。在单独的地址空间中运行设备API有两种好处：

(a)它使主机地址空间保持干净，没有破坏检查点实用程序（如CRIU）的设备映射和其他依赖项；

(b)它允许设备代理在时间片切片期间有效地共享在多个工作进程之间。

主进程和设备代理之间的通信在派发dispatch到设备的关键路径中，因此我们使用无锁共享内存通道使其低延迟，以便每次调用都没有上下文切换开销。上图显示了设备代理的高级架构。设备代理中有两种拦截器：DInt和SAInt。Dispatch拦截器DInt是语义无关的，只处理将API跨地址空间发送到设备代理服务器，处理参数/响应的序列化/反序列化。另一方面，语义感知拦截器SAInt在客户端或服务器端（分别称为客户端拦截器或服务端拦截器）上合并自定义逻辑，以实现例如屏障、时间片、内存管理等功能。请注意，调度拦截器和语义感知拦截器不是相互排斥的；例如，同一API可以同时具有客户端拦截器、用于跨地址空间的DInt和SAInt服务。

类似于 NVIDIA 提供的 MPS 方案，Singularity 提供了 devcie-proxy 服务代理。这层代理可以被看做是硬件抽象层(HAL)。这层代理实现了加速器相关 API 的拦截和转发。在拦截逻辑上，device-proxy 实现了两种拦截器，Dispatch Interceptors (DInt ) 和 Sematics-Aware Interceptors ( SAInt)。

DInt 是语义无关的，只负责不同地址空间的消息转发，处理参数及返回值的序列化和反序列化操作。如拦截cudaMalloc接管显存分配管理，拦截cudaLaunchKernel-在suqashing operations可以忽略部分rank的launchKernel操作。

SAInt 是语义相关的，需要与 client 或 sever 端的定制化逻辑相配合，实现诸如 barrier（log/replay）、time-slicing、显存管理等功能。

 

DInt：设备迁移等功能，拦截cudaStreamWaitEvent-同步，cudaMalloc-显存分配，cudaStreamCreate-流管理，cudaLaunchKernel-加载内核（Squashing selective operations），cudaGetLastError-返回结果（Hiding dispatch latency） 等。

设备迁移/句柄映射：cuda相关context、stream、event 的创建都会被DInt标记，涉及到这些的（状态更改调用）调用都会自动记录下来，并在恢复的时候进行重放

三类 SAInt 设备无关操作：内存分配（cudaMalloc、cudaFree），通信（nccl_allreduce-distributed barrier实现一致的checkpoint, ncclCommInitRank-推断集合通讯意图/DP)，设备同步(cudaStreamWaitEvent-正确处理时间切片跨rank的通信)。

以及host相关的SAInt：文件迁移：选择性地拦截CPU库；特别是libc I/O库（例如open、read、write等）被拦截以跟踪/记录作业对本地文件系统所做的更新

3.1. 限制DInt拦截面（Limiting surface area of DInt）

为了保证生产环境的可用性和可维护性，劫持层必须做到 完备 和 可扩展。考虑到相关的库更新非常快，拦截全部的 API 是不切实际的。Singularity 采用的方案是拦截底层的 Driver API。此方案在保证完备性的基础上兼顾了扩展性，因为无论何种上层库发出issue kernel，最终都会通过 Driver API。同时，Singularity 引入了代码生成器，自动化生成 DInt 代码的 Stubs，它只需要来自特定加速库（例如CUDA）的头文件列表，并带有一些注释来指示状态更改调用。

 

3.2. HAL（Hardware Abstraction Layer） for SAInt

虽然DInt自动生成了一些代码，但 Singularity 设备代理的大部分自定义功能（例如分布式屏障、时间片上下文切换）都驻留在SAInt中。大多数逻辑在SAInt是与设备无关的，使用了映射到设备特定 API（例如 nccl_allreduce）的硬件抽象层。加速器的硬件抽象层封装了跨加速器共同的关键功能。虽然当前实现是针对 NVIDIA GPU 的，但要处理新的设备类型，只需为该设备实现硬件抽象层，将设备的特定 API 映射到 HAL 中的等效 API即可。有三种不依赖于设备的功能需要SAInt：内存分配、通信和设备同步。

 

Singularity 定义了三类 SAInt 操作：内存分配（cudaMalloc、cudaFree），通信（nccl_allreduce-distributed barrier实现一致的checkpoint, ncclCommInitRank-推断集合通讯意图/DP)，设备同步(cudaStreamWaitEvent-正确处理时间切片跨rank的通信)。

1）内存分配：Replica splicing for memory sharing – 5.2 减少checkpoint大小

内存分配API需要一个SAInt，因为设备代理接管了内存分配。这使得设备代理能够完全看到GPU内存中实际使用的区域，有助于减少检查点大小。它还允许设备代理使用自定义内存分配机制，以帮助在同一GPU上的多个工作程序之间进行透明的时间切片。

2）通信: distributed barrier – 4.3 在作业发出每个数据allreduce之前发出异步串行元allreduce

大多数加速器都有一个集合通信库（例如NVIDIA GPU的NCCL，AMD GPU的RCCL）。这些API（如NCCL）上的SAInt使Singulariy可以实现分布式屏障（distributed barrier），以便在分布式作业的多个工作程序之间进行同步，以获得一致的检查点（例如，在任何rank上没有集合调用时）。Singularity通过利用相同的通信API提供了通用屏障实现。在这些API上使用屏障还有助于管理弹性时间切片期间的集合调用。

3）设备同步: 语义感知时间切片– 5.1 在单个小批量中，框架可能会发出多个异步 allreduce 调用(以重叠计算与通信)。在发出小批量的 allreduces之后，PyTorch 等框架通常会在 GPU 上执行同步操作 (cudaStreamWaitEvent )，然后再复制回平均梯度。在这个同步点，设备代理切换到共享 GPU 的下一个级别，并让它独占运行直到它到达同一点，然后上下文切换到下一个级别，依此类推。

设备同步API需要一个SAInt来处理透明弹性。Singularity中的时间切片是语义感知的，因为必须正确处理跨时间切片rank的通信。正确处理同步API（例如CUDA中的cudaStreamWaitEvent）对于时间切片期间的正确性和活性（liveness）至关重要。。

3.3. Host-specific 功能

除了设备硬件抽象层，SAInt 在 Singularity 中还被用来处理 CPU 的某些 API。

特别的，对 IO 库中的 open / read / write 操作进行劫持，以用来迁移用户任务生成的 track / log 文件。

除了设备的硬件抽象层外，Singularity还使用了一个SAInt来选择性地拦截CPU库；特别是libc I/O库（例如open、read、write等）被拦截以跟踪/记录作业对本地文件系统所做的更新，以便可以将修改的文件与进程检查点一起迁移。与设备库API的SAInt不同，主机SAInt没有相应的DInt，并在主机地址空间中运行。虽然特定于领域的拦截是与设备无关的，但为了简单起见，下面的部分重点介绍NVIDIA GPU。

4. 透明迁移机制

在Singularity中，运行DNN作业的抢占、恢复和弹性伸缩及到四种广泛类型的状态的一致检查点和恢复：(a) CPU中的程序状态（例如堆栈、堆、指令指针等），(b) GPU中的模型训练状态（例如模型参数、优化器状态等），(c) 处理CPU和GPU之间交互的控制状态（活动流、同步事件），以及(d) 不同类型并行性（数据/管道/张量并行等）的节点间和GPU间通信状态。

对于已调度的迁移以及从非计划故障中恢复，Singularity的透明检查点逻辑以两种模式执行：(1) 基于外部命令的按需执行（调度器），当调度程序决定需要抢占作业的时候 (2) 基于用户指定的间隔（epoch级别或基于时间）。

对于通用DNN作业的透明检查点是具有挑战性的，原因有几个。首先，在检查点给定作业时，Singularity必须确保跨多个主机和GPU跨多台机器的分布式状态的一致切分；分布式作业的所有worker必须在集合通信方面处于安全和一致状态（例如allreduce）。其次，尽管CUDA等专有闭源库管理状态，但必须一致地恢复CPU和GPU之间的正在进行的状态（in-flight state）（例如活动句柄、存储在主机内存中的设备地址）。第三，对于具有数百个worker的大型分布式作业，必须保持检查点空间开销较低。

4.1. 检查点的程序状态

有多个系统提供地址空间迁移，其中CRIU是最广泛使用的。但是，CRIU的一个关键限制是它不能处理使用GPU的进程的设备映射。要使用CRIU，必须将主机地址空间与设备相关的库隔离开来。幸运的是，设备代理架构为我们提供了这种隔离。设备代理服务大多是无状态的（下面有一些例外），因此不会被检查点记录；它只是在目标处重新启动。

4.2. 检查点的设备状态

设备的状态保存：通过设备到主机（device-to-host）的memcpy，设备代理进程通过检查点保存模型状态（例如参数）。由于在Singularity中进行了内存分配拦截SAInt，因此设备代理知道GPU内存的哪些区域实际上正在使用，device-proxy 拥有全局的显存分配视图，从而显着减小了检查点大小。

恢复：一个挑战是，在目标处恢复时，设备内存可能会映射到新的设备代理服务器地址空间中的不同地址，从而使主机进程中的指针无效。为了避免这种情况，设备代理在启动时独占整个GPU内存（留有一些余地用于由设备库跟踪的状态），并且具有由设备分配器（例如cudaMalloc）执行的mmap的server SAInt，以始终映射到相同的CPU地址。Singularity 采用了显存托管 + memory map 的方法，这一方法保证存放在物理地址上的内容映射到中断前的虚拟地址空间上。

 

设备句柄映射

与内存指针类似，主机地址空间还保留对设备状态的其他句柄。例如，cudaStreamCreate返回一个不透明句柄，可以由主机用作后续GPU调用中的引用。但是，由于迁移后重新启动了设备代理服务器，因此该句柄将无效。为了在迁移过程中保持这些句柄的可信度，我们对这些句柄进行虚拟化。设备代理不返回设备返回的实际句柄，而是返回虚拟句柄，并将此映射记为客户端状态的一部分。恢复和重放后，物理句柄可能会更改，但虚拟句柄仍然不变。所有状态相关的 API，比如 context、stream、event 的创建都会被DInt标记，涉及到这些的调用都会自动记录下来，并在恢复的时候进行重放。另外采用了领域知识规则对日志进行压缩。

 

4.3. 通信状态

大多数 DNN 任务中的集合通信操作都会采用集合通信库，例如：NCCL。我们无法做到通信过程中的状态保存，所以在状态保存之前，我们对任务实行静默，以保证没有在执行的通信操作。

静默 操作无法由每个 worker 独立执行。在一个集合通信操作(e.g., allreduce）中，所有参与方必须完成同一个操作。如果一个 worker 在第 n 次迭代完成后进行了检查点的创建，另一个 woker 在第 n + 1 次迭代完成后进行了检查点的创建，那么这次状态保存会发生死锁而永远无法完成。在创建检查点之前，所有的 wokers 必须完成了同一个集合操作。为了解决这个问题，Singularity 创新的使用了分布式 barrier 算法，完全透明的实现了这一特性。

分布式屏障Distributed barrier

分布式 barrier 是实现分布式训练热迁移的关键。Singularity的核心是一种新颖的、工作负载感知的调度器，它利用与作业使用的通信库相同的通信库进行集合通信，通过引入额外的meta-allreduce来交换屏障协议状态，从而避免引入新的故障路径（例如通过带外通道TCP、远程存储等进行协调）。该算法需要确保额外的元allreduce相对于作业执行的常规allreduce在所有工作进程中具有相同的顺序（集合通信的程序顺序要求，以避免死锁）；我们的屏障算法在作业发出每个数据allreduce之前发出异步串行元allreduce。这可以轻松地确保一致的程序顺序。该算法分为两个阶段：第一阶段是稳态，第二阶段是收到屏障请求时。串行元allreduce是一个SUM allreduce，其有效负载由两个整数组成

• need_barrier：如果接收到 barrier 指令，worker 发送 ‘1’ ，否则发送 ‘0’。如果 SUM(need) > 0，则 worker 知道某处触发了 barrier 协议，切换到 Phase 2
• ack_barrier：如果切换到 Phase 2 则 woker 发送 ‘1’，它表明 woker 直接或间接地看到了 barrier 请求，否则发送 ‘0’。如果 SUM(ack) == world_size (i.e., the total number of ranks),则 worker 知道其它合作者都回应了请求，接下来可以安全的获取 barrier。

 

一旦worker进入Phase 2，它就进入了同步模式：该工作执行的每个集合调用都是同步的；这确保了屏障协议的及时终止。屏障算法保证在最多两个小批次内完成，并保证在屏障时没有未完成的集合调用。在稳定状态下，它的开销非常小，因为在Phase 1 期间，微小（2字节）的meta-allreduce是异步的。

虽然上述算法适用于数据并行作业，但张量并行和管道并行作业存在额外的复杂性，可能会在不同节点组之间执行多个allreduce，除了点对点调用（用于流水线）之外。虽然我们可以扩展上述算法以推断管道和数据通信之间的相对顺序，但我们当前的设计优化了简单性和检查点大小，方法是使用领域知识：我们确定小批次的结束时间，在此时刻张量并行或管道并行维度中都没有正在进行中的通信。我们使用与上面相同的串行meta-reduce协议，但仅在小批次结束时使用一次（而不是每个数据allreduce使用一次）来实现屏障。权衡是屏障被延迟到小批次结束（大型模型需要几秒钟），但与在小批次中获取屏障相比，它给我们带来了更小的检查点大小。

 

4.4. 文件系统状态

DNN 训练任务有时会安装或更新本地文件，这些在热迁移过程中都应当保留下来。执行容器维度的文件系统状态保存代价太大。Singularity 采用了 host SAInt 劫持 libc 文件系统 APIs 帮助解决：无论何时本地文件以可写模式被打开，我们就将这个文件名记录到 log 中，在 checkpoint 的时候，将这些文件拷贝过来。远程存储使用内容校验在worker之间对数据副本进行去重。

4.5. Checkpoint / Restore 流程

成功获取屏障后，通过对各个worker执行criu checkpoint来完成检查点。然后将CRIU转储以及活动张量的GPU状态转储移动到远程存储中。在新目标节点上，在criu恢复时，进程从准确的检查点开始执行（如设备代理客户端获取屏障之后的状态开始）。设备代理客户端执行的第一个操作是重新生成一个新的设备代理服务，然后重放状态更改调用（replay state-changing calls）以将GPU带回到检查点之前的状态。设备代理服务还将GPU张量复制回GPU RAM，地址与检查点之前相同。最后，设备代理执行新的合约（rendezvous），以便ranker可以发现彼此的新位置并重新建立通信环。除了由调度程序启动的按需检查点外，Singularity中的每个作业都按用户指定的频率（例如每30分钟）进行检查点以处理未计划的故障以防止异常中断。

 

4.6. 检查点压缩

Singularity采用了几种技术来减小检查点大小。首先，Singularity对每个设备缓冲区执行内容校验和以在worker之间进行去重。在上传设备缓冲区之前，worker计算其内容校验和，并仅在没有其他worker上传相同缓冲区时上传缓冲区。因此，Singularity中的GPU转储大小与用户级检查点的大小相似。

CPU地址空间的CRIU转储在空间和时间上进行去重。首先，主要训练过程和数据加载器进程之间的内容重叠很高；我们拦截CRIU write calls 以执行基于内容哈希的页面去重。其次，在不同时间点对同一进程的检查点之间存在很高的重叠度（因为地址空间的变化很小）；在时间维度上进行去重使得后续增量检查点比第一个CRIU检查点小得多 。

 

5. 透明弹性机制

Singularity引入了一种新的能力，可以调整任何DNN训练作业以使用不同数量的GPU，而无需更改用户代码并且不会影响作业的语义。在Singularity中调整作业是完全透明的：对于用户而言，作业始终以相同的world size（即ranks/workers数量）运行。调度程序可以将每个工作者一对一地映射到物理GPU（完全扩展），也可以使用多对一映射，其中物理GPU在多个worker之间进行虚拟化和时间分片（缩容）。相比之下，DeepSpeed 或PyTorch Elastic 等库会将弹性暴露给用户，因为作业在从先前检查点调整大小后会以不同的world size重新启动，导致浪费（例如，自上一个检查点以来的初始化和迭代被重做）。

透明弹性建立在Singularity中透明迁移支持的基础上。例如，要将作业从4个GPU缩小到1个GPU，我们只需对3个rank进行CRIU检查点，并将这些进程迁移到单个GPU上进行时间分片。由于CRIU检查点的属性，woker从完全相同的程序状态恢复而不需要重作（redo）任何计算，因此伸缩是工作连续的。

透明和动态伸缩面临着几个技术挑战。首先，在同一GPU上对训练作业的多个worker进行时间分片时，worker之间的细粒度通信（例如allreduce）必须像在不同GPU上运行的worker一样进行；这需要时间分片也具有语义感知性和细粒度性（每个小批次多个上下文切换 context switches）。其次，对于大型模型，每个worker可能会利用几乎整个GPU上的RAM；将多个worker放在同一个GPU上需要将GPU状态来回交换到主机内存中，这是一项昂贵的操作（例如3-10倍开销）。第三，为了支持使用数据并行、管道并行和张量并行组合的作业的透明弹性，需要仔细地将worker放置在GPU上，以便只有相同模型并行分片的数据并行副本在同一GPU上进行时间分片，并防止通信调度中的死锁。

1. 当多个 worker 在同一个 GPU 上工作时，如何实现和多个 GPU 上一样的通信？
2. 当模型很大时，如何在多个 worker 之间进行显存切换？
3. 如何支持不同的并行模式，数据并行，张量并行，流水线并行，而不发生死锁？

 

5.1. 语义相关的时分复用 time-slicing

为了在多个工作进程之间共享GPU，可以简单地在同一GPU上独立运行这些进程。但是，这种方法不可行：在大型模型中，每个worker几乎使用了整个GPU RAM，因此运行多个这样的worker会耗尽内存。因此，在Singularity中需要进行语义感知的时间切片。

Singularity中的设备代理使得这种时间切片成为可能。因为设备代理与主机进程解耦，所以我们可以在多个主机进程（即多个ranks）之间共享相同的设备代理。由于所有与GPU的交互都通过设备代理进行，它会智能地调度多个rank，只允许一个rank在给定时间在GPU上执行，并选择特定的点来上下文切换到另一个rank。在概念上，在上下文切换时，设备代理会换出swap out原始rank使用的GPU内存（即复制到主机内存），然后换入swap in新rank的GPU内存，从而使每个工作进程几乎使用整个GPU RAM。当然，这种swap out/swap in会非常昂贵；

1）降低切换频率

为了降低开销，我们必须仅在绝对必要的情况下进行上下文切换。当一个rank只执行其各自的数据上的计算（例如，前向和后向传递操作，如矩阵乘法）时，就没有必要进行上下文切换。在反向传递之后，数据并行rank参与集合通信（例如，allreduce）以交换梯度，这需要所有rank参与（并贡献其各自的梯度），从而需要上下文切换。请注意，在单个小批量中，框架可能会发出多个异步allreduce调用（以重叠计算和通信）。在小批量的all reduce都已发出之后，例如PyTorch的框架通常会在GPU上执行同步操作（cudaStreamWaitEvent），然后将平均梯度复制回来。在此同步点上，设备代理切换到共享GPU的下一个rank，并让它独占运行，直到它达到相同的点，然后切换到下一个rank，依此类推。

2）解耦作业world size和NCCL 看到的world size

集合通信（例如allreduce）通过专有库（例如NCCL）进行。NCCL具有通信器的概念，该通信器为参与rank的特定环初始化，并且随后的操作（例如allreduce）仅引用通信器。为了使NCCL通信器与用户级时间分片管理的交互可管理，我们将作业的逻辑数据并行world size与NCCL看到的world size分离; 在我们的方法中，NCCL仅看到每个GPU一个rank。在时间分片期间，设备代理透明地将本地累积到其临时缓冲区scrach buffer中，并且仅最后一个共享GPU的rank使用本地累积梯度的结果执行实际nccl_allreduce。因此，在resize操作之后，NCCL看到的world size发生了变化（由恢复后的新约会rendezvous处理）。

 

5.2. 显存副本分片（显存换入换出开销优化）

上面介绍的时间分片机制可以保证弹性伸缩的正确性，但性能太慢了。要知道，V100 拥有 32 GB 显存（而最新的 A100 高达 80 GB）。在大模型中每个 single worker 会基本占满显存。每个 mini-batch 上下文切换伴随的显存换入换出需要 2~4 秒的时间，而 mini-batch 的运行之间可能都远小于这个值（一般为几百毫秒），这样带来了 5~10 倍的开销。

1）基于校验和的动态去重

我们把训练中占用显存的内容分为以下四类。

1. Parameters (P) 存放模型每一层的权重和参数，前向和后向过程都会运行在这些 tensor 之上。
2. Optimizer State (O) 存放每次迭代的 delta 值以及历史状态（例如，first and second moments of gradients）。
3. Gradients (G) 每个工作副本拥有自己的梯度拷贝。在反向传播过程中，不同副本的梯度将取平均值，然后更新为统一的权重。
4. Activations (A) 临时存放前向过程中每一层的输出；反向传播中被用于计算梯度。
5. P 和 O 都会在 mini-batch 结束的时候进行同步，且在所有副本上保持一致
6. A 会在每一个 mini-batch 结束的时候被框架释放。

Replica splicing 利用的关键洞察是，在数据并行的副本中，参数（P）和优化器状态（O）是同步lock-step进行的，即它们在每个 mini-batch 结束时由所有副本一致地更新，使用（相同的）平均梯度。因此，在 mini-batch 结束时，与共享 GPU 的rank对应的张量 P 和 O 将在各个rank之间相同。此外，在 mini-batch 结束时，激活张量 A 由框架释放，因为反向传递已完成。我们利用这些洞察用以下的方式有条件的进行换出/换入。

 

如何避免不同 rank 的 P 和 O 发生频繁换入呢？答案是：增加副本

如果显存有足够的空间存放额外两组 P 和 O 的副本，就可以避免换入。运行过程中最多有两个 version 的 P 和 O 被激活使用，分别是当前 mini-batch 的和前一个 mini-batch 的。第三份副本用来做草稿空间scratch buffer，避免破坏前一个 mini-batch 的内容

由于设备代理控制内存分配器，因此它可以看到框架分配的每个缓冲区。在上下文切换期间，设备代理为每个活动缓冲区计算内容校验和。在换出期间，它首先查找主机是否已经包含具有相同内容校验和的缓冲区；如果是，则避免换出，并将 GPU 缓冲区标记为未使用（在新rank需要新分配时lazy GC，所以可能有机会缓存多个版本在设备中）。同样，在换入新rank数据时，它检查设备是否已经具有具有该校验和的缓冲区；如果是，则避免从主机换入。请注意，尽管内容匹配，但该缓冲区可能映射到新rank中的不同设备地址；在这种情况下，设备代理将该缓冲区从设备移动device-to-device move到所需地址，这比从主机进行换入（HBM 带宽 vs. 主机内存带宽）更快。

如果4个rank共享一个GPU，则在上下文切换期间换出P和O缓冲区只需要为第一个rank执行换出；其他rank会发现校验和已经存在于主机内存中，并忽略。但是，请注意，当rank开始其时间切分，其本地状态包含来自前一个小批次的P和O，而上一个rank的副本已更新为当前小批次，必须为每个rank执行换入；如果我们在GPU内有足够的空间来存储两个额外版本的P和O（对于任何时间分片因子> 2），则可以避免换入。这带来了两个挑战：

（a）两个额外副本的P和O的额外空间对于大型模型是不可接受的

（b）我们仍然需要在上下文切换期间执行P和O的设备到设备device-to-device复制，因为每个rank可能已经分配了不同地址的相同缓冲区。

D2D copy代价仍然很大，由于（很多）源和目标缓冲区之间的循环依赖关系，会强制复制以阶段方式进行，从而限制了并行性。下面我们将介绍如何解决这两个挑战。

 

2）显存分配一致性领域相关的知识

如何解决空间问题和地址不同问题？

首先来看地址不同的问题。这是因为每个 worker 负责分配自己的 buffer 空间，而每个迭代之间有些 buffer 的 size 是会发生变化的，比如随着 mini-batch 的 input size 不同，A 的 buffer 会有不同的 size。这样就导致了不同的迭代步之间的地址无法对齐。

Singularity 使用高低地址分离的显存分配器来解决这个问题。稳定 size 的 buffer 被分配到高地址（例如 P 和 O），其它 buffer 被分配到低地址。

在数据并行的情况下，每个副本都会执行自己的内存分配，因此会导致同一 P 和 O 缓冲区在不同的副本中具有不同的地址。为了解决这个问题，Singularity 使用了关于深度学习训练的领域知识，使得地址保持一致，而无需在副本之间进行显式协调。我们使用了一个高低端内存分配器。稳定size的缓冲区（如 P 和 O）在地址空间的高端进行分配，而其他缓冲区则在低端进行分配。这样可以确保瞬态分配（如激活A）中的不稳定性不会影响高区域中的内存分配器元数据，从而确保稳定缓冲区（如 P 和 O）在所有副本中具有相同的地址。我们已经在各种模型和 PyTorch 版本上进行了实证验证。为了识别 P 和 O 等稳定缓冲区，我们使用了一个预先确定的堆栈跟踪列表（Python 和 C++），用于参数P和优化器状态O分配；这个列表需要每个 PyTorch 版本更新一次。当分配内存时，设备代理客户端会获取堆栈跟踪并与列表进行匹配。请注意，在极端情况下，如果注解Annotation不正确并且我们无法获得一致的P和O地址，则只会影响性能（可测量），但不会影响正确性。

 

3）挤压可选择操作 Squashing selective operations

为了避免处理多个P和O的副本（这会导致换入成本或额外的GPU内存），我们使用另一个领域特定的洞见。根据定义，所有数据并行副本在完成小批量后将到达P和O缓冲区的相同版本。我们还知道，只有在梯度在副本之间进行allreduce之后，才会更新P和O缓冲区。因此，如果我们可以识别更新参数和优化器状态的操作，我们可以仅在共享设备的一个rank（root rank）中执行这些操作，并简单地“挤压”其他rank中的这些操作，因为（a）它们最终会导致相同的最终状态，（b）缓冲区在rank之间具有相同的对应地址，因此随后的小批量计算将看到正确的数据。为了挤压操作，设备代理只需省略向GPU发出cudaLaunchKernel的操作即可。通过这种挤压，我们避免了换入P和O的先前版本，因为它们不再由其他rank更新。

 

请注意，通过挤压，我们利用领域特定的洞见来改变执行顺序，因此我们需要确保它不会影响正确性。虽然合理的模型符合我们在挤压中的假设，但我们应该防止病态模型因违反这些假设而遭遇静默损坏/不正确执行。为确保鲁棒性，我们采用保守验证方法；如果验证成功，则保证挤压是安全和正确的，但如果验证失败，则保守地禁用挤压。因此，违反我们假设只能影响性能（在这种情况下，我们禁用该模型的时间分片），但永远不会影响正确性。

 

保守验证：禁用挤压运行第一个mini-batch（周期性的K个mini-batch），保证正确执行；在设备代理识别变异mutated buffer是很挑战的，因为GPU操作（如 cudaLaunchKernel）参数地址是indirect-addressed via multiple levels of GPU buffers，对设备代理是模糊的。因此验证依赖一个新的方法，通过缓冲内容的校验和来推断操作结果。

 

5.3. 如何处理模型并行作业

前面讨论的都是数据并行作业；模型并行有新的挑战：例如，张量并行作业在前向和后向传递中的每个矩阵乘法中执行allreduce。如果我们为这样的allreduce进行上下文切换，副本切分将无法工作，因为激活张量仍然处于活动状态。同样，管道并行作业在每个微批次的GPU/节点之间执行激活和梯度的点对点发送和接收；在微批次期间进行时间切片会因为梯度和激活处于活动状态而导致过多的交换。

Singularity使用两种关键技术来解决处理模型并行作业（如张量并行和管道并行作业）时遇到的新挑战：splice-aware placement和推断集合调用的意图。使用splice-aware placement，Singularity确保只有同一模型并行分区的数据并行副本在同一GPU上进行时间切片。例如，要在4个物理GPU上运行具有4路流水线和2路数据并行性的8-rank作业，Singularity会将同一流水线阶段的两个数据并行副本放置在每个GPU中。对于3D并行作业也是如此；在同一设备上进行时间切片的rank将属于同一流水线阶段和相同的张量并行分区。请注意，这需要Singularity认知到（aware）rank分配逻辑。两个流行的库NVIDIA-Megatron和DeepSpeed 在并行性维度上具有相同的rank分配逻辑，并且这种逻辑在Singularity中会镜像。对于使用具有不同rank分配策略的自定义启动器的作业，Singularity提供了一个API，用于作业通信所有rank-to-topology 映射所有的ranks（例如，Rank 4是DP0、MP0、PP1等）。

其次，设备代理推断集合通信的意图，并仅在数据并行维度上对集合调用进行时间切片。其他集合调用只是简单地透传而不进行上下文切换，这是正确的，因为这些调用的完成仅取决于在其他GPU上执行的rank，并不需要来自在同一GPU上进行时间切片的其他数据并行副本的输入。但是，透明地推断特定allreduce调用的意图是非常困难的，因为每个用户模型都可能具有其自己的控制流和跨多个并行性维度的执行顺序。Singularity利用集合通信的初始化路径（例如ncclCommInitRank）来实现这一点。它在每个ncclCommInitRank之后强制进行上下文切换，并且设备代理（与使用相同设备的所有rank共享）保留每个通信器计数。在完整的一轮上下文切换后，当一个通信器的本地计数大于1时，设备代理会推断该通信器处于数据并行维度。在集合调用期间，它只需查找通信器上的映射即可知道它是否为数据并行。

 

5.4. Handling ZeRO-redundancy optimizer

ZeRO通过消除数据并行worker之间的冗余来分割数据并行状态，从而使数据并行状态不具有冗余。这种分区违反了我们压缩验证的不变量。为了解决这个问题，Singularity引入了ZeRO的部分分片概念，将分片因子（最小需要适合GPU的模型）与数据并行度（用于并行）解耦。如果两者相等，则该模型无法缩小到更少的GPU，因为它无法适配。如果数据并行因子更高，例如分片因子的4倍，则我们可以支持高达4路时间切片/缩容。在这种情况下，部分分片因子只是模型并行性的另一个维度，并且只有相同ZeRO分片的副本才会被时间切片。在DeepSpeed中引入部分分片非常简单（大约30行Python代码）。

 

6. 系统实现

透明检查点、抢占、恢复和弹性，这些机制是Singularity调度程序的的一部分。接下来我们重点突出一些实现上的挑战。

1）如何进行 kernel launch / Serializing opaque parameters

在我们的基于拦截的设备代理中，cudaLaunchKernel的DInt很具有挑战性，因为它的签名是不透明的，使得序列化变得困难（签名是由NVIDIA的nvcc在内部生成的，并且对拦截器不可见）。为了解决这个问题，我们有一个用于cudaLaunchKernel的自定义server SAInt，它使用CUDA工具包中的二进制实用程序cuObjDump解析生成的内核库并提取参数信息。为了避免高成本，我们缓存这些信息，并仅在缓存未命中时运行cuObjDump。对于JIT内核，我们拦截nvrtcCompileProgram并通过解析生成的PTX提取参数签名。

2）如何隐藏分发延迟

设备代理中跨地址空间调用发生在操作的关键路径上，例如cudaLaunchKernel和cudaGetLastError，这会影响性能。我们针对最频繁的调用使用特定领域domain-specific的优化。对于cudaGetLastError，我们在每次内核启动时机会性地在服务端上发出issue它，并将其与其响应一起捆绑发送，以便当PyTorch发出issue它时设备代理客户端可以从缓存中返回它。对于cudaLaunchKernel，我们执行延迟错误通知，调用在客户端返回而不等待来自设备代理服务器的响应；在向服务端发出下一个调用之前懒惰地lazily读取响应，从而允许PyTorch处理客户端和服务端产生的延迟之间的重叠；因为当遇到这种（罕见）错误时PyTorch（和其他框架）会崩溃，所以这不会影响作业语义。

3）如何隐藏上下文切换延迟

在时间片切换期间从一个rank切换到另一个rank涉及计算所有活动设备张量的校验和，并将其与另一个rank的副本进行比较，如果需要则执行缓冲区移动（几毫秒的CPU活动）。此外，切换逻辑取决于校验和计算的输出，而校验和计算又等待所有先前的GPU操作完成，隐式地强制设备同步。为了避免在关键路径中产生这种成本，Singularity执行下一个rank的急切调度eager dispatch。设备代理与切换逻辑并行地开始为下一个rank提供服务，从而通过下一个rank（CPU逻辑、设备操作的调度）与切换延迟重叠有用的工作。通过仔细使用异步排序原语，如cudaStreamWaitEvent，我们确保新rank的操作仅在切换完成后在GPU上执行。

 

7. 性能评估

7.1. device-proxy 开销

引入 device-proxy 后，平均开销可以控制在 3% 以内。部分模型甚至有正向收益，主要归功于

cudaLaunchKernel 优化。

 

7.2. 检查点大小

checkpoint 可以分为三个部分，GPU 的状态转储、首次 CRIU 的状态转储、后续增量 CRIU 的状态转储。

CRIU 转储部分是大头，这部分的内存占用会随着 worker 数量的增加而现象增加。从这份数据来看，在实际生产环境还是可以接受的。

7.3. 副本分片弹性伸缩

大部分模型在引入 Singularity 的弹性缩容后，开销小于 3%。甚至在小模型上，开销也小于 5%。小模型的 mini-batch 时间更短，上下文切换的延迟占比更大。

7.4. 迁移和 resizing 延迟

大部分模型迁移的端到端延迟为几十秒的时间，当前的时间包含了块存储的网络传输时间，如果使用 P2P 的方式直接传输，延迟会进一步降低。不过相比于训练一般持续时间为天级别甚至周级别，几十秒的延迟是完全可以接受的。

7.5. device-proxy 的鲁棒性和可维护性

已经在 PyTorch 1.6 / 1.7 搭配 CUDA 10.1 / 11.0 / 11.1 这些版本组合上完成了测试，结合实现来看，device-proxy 的鲁棒性和兼容性是有保证的。

 

8. 相关工作

1) DNN 调度

2) 弹性机制

3) 迁移机制

 

9. 结论

Singularity在调度深度学习工作负载方面取得了重大突破，将弹性等小众特性转化为主流、始终开启的特性，调度程序可以依靠这些特性来实现严格的SLA。通过新颖的机制，无需修改作业便可具有可抢占和可伸缩的能力，性能开销可以忽略不计，Singularity实现了前所未有的工作负载可互换性水平，使作业能够利用全球分布式集群中任何空闲容量，同时仍然保持SLA。Singularity通过一个非常简单的用户体验实现了所有这些：用户只关注于ML任务，不需要考虑检查点或弹性；这些机制是基础设施优化，对用户完全透明。