Singularity: Planet-Scale, Preemptive and Elastic Scheduling of AI Workloads

Abstract

Singularity
- 微软的全球分布式调度服务，高效和可靠地执行深度学习训练和推理工作负载
- 核心：是一个新颖的、工作负载感知的调度器，它可以透明地抢占和弹性地扩展深度学习工作负载
  - 以提高利用率
  - 而且不会影响它们在全球 AI 加速器(GPU/FPGA)中的正确性或性能
- 所有作业都是可抢占的、可迁移的，并且在默认可以动态调整大小(弹性) : 一个活动的作业可以被动态和透明地
  - 被抢占和被迁移到不同的节点集、集群、数据中心或者区域，并且可以从执行被抢占的地方精确地恢复
  - 在给定类型的一组不同的加速器上调整大小
- 机制透明：不要求用户对代码进行任何更改，也不要求使用任何可能限制灵活性的自定义库。
- 可靠：利用Singularity可以获得效率和可靠性增益，而对稳态性能的影响可以忽略不计。
- 我们的设计方法是对DNN 网络架构不感知的，并且可以处理各种并行策略(数据/流水线/模型并行)

Introduction

Singularity的建立有一个关键目标: 通过在全球规模的加速器的固定容量池中最大化总有用吞吐量来降低人工智能的成本，同时为多个定价层次提供严格的 SLA。图1显示了 Singularity的高级体系结构，包括其分层调度系统，该系统由在全球、区域和工作负载范围内调度微服务组成。

设计目标

不闲置资源:
- Singularity将所有加速器（舰队）视为一个单一逻辑的共享的集群，并避免任何资源碎片化或静态容量保留。
- 机会性地使用空闲容量
提供job级别的 SLA: Singularity通过尊重作业级别的 SLA 来提供隔离。
例如，Singularity适合于增加推理作业的负载，通过弹性缩小或抢占训练作业来释放容量。
失败恢复: 作业从它们被抢占的地方恢复，从而最小化浪费

关键机制

两个关键的核心调度原语:
- 抢占和迁移: Singularity可以透明地checkpoint、抢占和迁移所有跨节点甚至跨集群和区域的 DNN 作业。checkpoint是通过使用一个有效的同步屏障来实现对分布式作业的所有Worker的分布式状态的一致切割。
- 调整大小/弹性: Singularity使所有作业都能够以透明的方式动态和弹性地调整大小，以便使用可变数量的 AI 加速器。
机制特点
- 透明
  - 含义：1）不需要修改用户脚本，2）不依赖于框架/库
  - 好处：为任意的深度学习工作负载提供一致的 SLA ，而不依赖于用户的任何合作来维护 SLA
  - 因此：checkpoint、迁移和弹性在默认情况下对所有作业都是启用的
- 保存工作
  - 迁移或调整大小的作业在程序执行的同一点恢复，其状态(例如，程序计数器、堆栈等)与被抢占或调整大小时完全相同。
Singularity中透明抢占、迁移和弹性的核心：自动将作业与加速器资源分离（解耦）
- Singularity中的Worker和加速器设备之间的动态绑定的，并且在工作的生命周期中不断变化。要扩展或缩小作业，我们只需更改Worker映射到的设备数量。
  这对用户来说是完全透明的，因为不管运行作业的物理设备的数量如何，作业的Worker的总数保持不变。
- Singularity使用了一种称为副本拼接的新技术：在同一个设备上对多个Worker进行时分，开销可以忽略不计，同时允许每个Worker使用整个设备内存。
  - 副本拼接依赖领域知识去利用分布式训练job的Worker之间在程序执行的特定点上的内存内容相似性
解耦的实现：设备代理
- 设备代理在自己的地址空间中运行，并且有一个到物理加速器设备的双射。当作业Worker初始化设备 API 时，它们被拦截并通过共享内存发送到设备代理进程，该进程的生命周期与作业Worker进程的生命周期解耦。
- 这种分离实现了两个关键的好处:
  - 主机地址空间可以保持Clean：不会有设备特定的映射和其他由 GPU 库(如 CUDA)创建带来的副作用 –> 从而可以使用CRIU来checkpoint和迁移主机进程
  - 它允许动态的，透明的，时分的多个Worker在同一个设备上工作，多路复用和Worker间调度交由代理设备执行

贡献

checkpoint机制：一种透明和健壮的机制来checkpoint那些不支持checkpoint的通用 DNN 作业，从而使所有作业可以自动checkpoint，抢占和可迁移。
分布式屏障：基于新颖的，语义感知的技术，以完全透明的方式实现了一个DNN job的所有Worker之间的分布式屏障。这是实现 DNN 训练 job (包括程序计数器、堆栈、 CPU 和 GPU 内存等)跨多台机器的分布式状态的一致性削减的关键，这样就可以精准地从被抢占时的点重新启动
副本拼接：该技术允许在同一 GPU 上透明地对 DNN 作业的多个Worker进行时分，开销可以忽略不计(2-3%) ，并且使用超轻量级上下文切换以一种健壮和通用的方式。
我们评估我们的核心机制的功效和端到端的效率，通过使用不同类型的并行性(数据，流水线，或张量并行性)的模型进行详细的实验
- 性能开销：对于大范围的工作负载，GPU 调用的动态拦截、透明分布式屏障算法和透明时分的稳态性能开销可以忽略不计，均在3% 以内
- 稳健性：尽管 PyTorch 和 CUDA 的版本发生了迅速变化，但这种方法仍然是健壮的和实用的。
- 延迟： Singularity中的迁移和弹性延迟是合理的(几十秒) ，没有重复计算，
- checkpoint大小：checkpoint大小与用户级checkpoint相当。

核心调度机制总览

为了提高利用率和可靠性，Singularity所有作业在默认情况下都可以抢占和调整大小。

分布在多个节点上的使用各种并行方式的作业都可以进行checkpoint，使用透明时分在一个潜在的不同的区域，不同数量的设备上在稍后恢复

在本节中，我们首先描述这些机制的三个关键方面:
- 透明性：对用户代码没有更改或约束
- 工作保存：作业从先前被抢占的同一个程序执行点恢复
- 解耦执行

用户透明

现有的checkpoint和弹性方法
- 依赖于用户直接编写代码来实现这些机制
  - 让用户负担了保存/恢复 python 程序状态的复杂性
    例如，循环变量、控制流、学习速率调度器、数据加载器状态、指令指针等
  - 需要确保程序在正确的点恢复
  - 当作业向上或向下扩展时，需要改变超参数
- 或者使用处理checkpoint和弹性的特定库(Pytorch Elasstic/Deepspeed)
  - 用户失去了灵活性，因为这些库控制了训练循环以保持checkpoint易于处理。这限制了用户的可定制性; 因此这种库的采用率很低。今天大多数 DNN 的训练工作量是不可checkpoint或可调整的。
Singularity的核心机制（checkpoint/迁移/弹性）不需要用户的任何合作，默认情况下是自动的。
- 它使得调度程序能够依赖这些机制作为所有作业的first-class constructs，以提供严格的 SLA
- 它完全隐藏了checkpoint的复杂性和弹性，使得用户可以专注于编写具有完全灵活性的代码

工作保存

Singularity使用的checkpoint是由作业的各个Worker的一致地址空间快照组成。
- 当这些快照捕获完整的程序状态(如指令指针、堆栈、堆等)时，作业从被抢占的位置准确地恢复，没有丢失任何工作。
- 相比之下，现存的checkpoint和弹性机制迫使程序从以前的模型checkpoint重新启动，从而重新执行初始化工作和自上一个checkpoint以来执行的工作

解耦执行

提高作业的可靠性和船队的效率/利用率的目标是高度协同的
- Singularity调度程序通过解耦作业与底层资源之间的映射来处理这些目标。Singularity调度器透明地虚拟世界大小和rank分配。
- 这种解耦是至关重要的：对于透明的checkpoint和抢占作业以及随后在不同的节点、集群、数据中心恢复它们来说，这些节点、集群、数据中心与之前的checkpoint状态具有相同或不同的 GPU 数量。
- 在作业的稳定状态执行期间，Singularity调度器还透明地将作业工作进程的训练逻辑与其与 GPU 的交互解耦。

用户利好

透明和节省工作量的checkpoint、迁移和灵活性从根本上增强了调度器的能力:
- 提高了容错性。任何作业(不管用户是否已经编写了checkpoint逻辑)都可以在发生硬件故障(GPU/节点/网络)时从系统自动获取的最新checkpoint精确地恢复，而不是从头开始。这显著提高了有用的舰队范围吞吐量。
- 机会性使用容量。透明的抢占和迁移允许作业在任何地方使用空闲资源，而不管集群/区域边界如何。此外，它使得具有较低 SLA 的作业可以机会性地使用空闲容量，并且当具有较高 SLA 的作业到达时，可以迅速抢占(而不会丢失作业)。透明弹性使就业机会能够扩大使用闲置产能，并在产能变得稀缺时收缩。
- 针对本地的后台碎片整理。用于容错(例如机架)、设备与设备互连、数据本地等的局部或拓扑域在作业进入和离开系统时会变得支离破碎，这使得用局部约束调度大型作业变得困难。小作业的迁移使调度程序能够对局部域进行碎片整理，以放置更大的作业。
- 在线升级。可以在不杀死作业的情况下进行全舰队范围的实时升级，因为在这些机器上运行的作业可以廉价且透明地迁移到不同的集群。

训练效率吞吐的SLA

传统的 SLA（延迟，5个9的可用性)，适用于推理工作量，他们不适合 DNN 训练。
Singularity引入了 GPU fraction度量：在抢占和灵活性面前量化吞吐量。
- image-20230317173405341
- 表1描述了 Singularity提供的多层 SLA。
  一个需要 N 个 GPU (基于软配额)的任务可能获得多于 N 个或少于 N 个 GPU，这取决于竞争集群负载。用户只需支付实际使用的费用，而不需要支付配额;
符号定义
- $T_{ideal}$ 是某一个job在 N 个专用 GPU 上非抢占式运行所需的现实世界完成时间
- $T_{real}$ 是在Singularity 中的实际完成时间
GPU time fraction of a job = $T_{ideal}$ / $T_{real}$。
- $T_{real}$ > = $T_{ideal}$，在允许资源超额订阅的情况下，因为作业在执行过程中可能被 Singularity抢占或缩小。
- 在Singularity中运行会相对慢一些，如果工作在一个专用容量设置中需要 H 小时
  - 高级 SLA 工作：最多需要 H/0.95小时
  - 标准 SLA : 最多需要 H/0.7小时
- GPU 分数 SLA 以小时粒度强制执行。Singularity中的调度策略旨在最大化舰队范围的吞吐量，同时最小化违反这些 SLA 的情况。

特定领域拦截的方法

设备代理的定义与组成
- 任何与加速器的交互都必须经过特定的库(例如，用于 NVIDIA GPU 的 CUDA、用于 AMD GPU 的 ROCm) 。 Singularity通过 LD_PRELOAD 机制动态拦截它。大多数此功能驻留在设备代理中。
- 设备代理可以被视为加速器设备的硬件抽象层服务
  - 组成：
    - 服务器组件：每个设备有一个
    - 客户端组件：嵌入在与设备交互的每个进程中
  - 主机调用的所有加速器特定的 API 都会被拦截并发送到设备代理服务器，该服务器在一个独立的地址 $space^1$中运行。
  - 在单独的地址空间中运行设备 API 有两个好处:
    - 它保持主机地址空间没有设备映射和其他类似的依赖关系，避免影响checkpoint实用程序，如CRIU
    - 它允许设备代理在弹性时分期间有效地跨多个工作进程共享
设备代理的通信与体系结构
- 主机进程和设备-代理之间的通信处于分派到设备的关键路径上，因此我们使用无锁共享内存通道，使得每次调用没有上下文切换开销，从而降低了延迟。
- image-20230317225820981
- 设备-代理中有两种类型的拦截器:
  - 分派拦截器($D_{Int}$ )
    - $D_{Int}$ 是语义无关的，它只处理将 API 跨地址空间发送到设备代理服务器，处理参数/响应的序列化/反序列化
  - 语义感知拦截器($SA_{Int}$)。
    - $SA_{Int}$ 在客户端或服务器端(分别称为客户端 $SA_{Int}$ 和服务器 $SA_{Int}$ )合并自定义逻辑，以实现诸如屏障、时间分片、内存管理等功能
  - 请注意，$D_{Int}$ 和 $SA_{Int}$ 并不相互排斥;
    例如，同一个 API 可能同时具有客户端 $SA_{Int}$ 、跨地址空间 $D_{Int}$ 和服务器 $SA_{Int}$ 。

限制$D_{Int}$的拦截面积

目标：为了在生产系统中实用和可维护
- 拦截层必须是完整的(作业调用的所有设备调用必须被拦截并发送到设备-代理服务器)
- 并且可伸缩(低工程成本)
背景：鉴于直接发布内核的新库(如 Apex、 Thust、 Deepspeed、 OpenAI Triton)的发展速度很快，拦截所有这些与设备交互的 API 的幼稚方法是不切实际的。
Singularity通过在低层拦截来限制 $D_{Int}$ 的表面积
- 低层：用于启动内核的驱动程序 API (例如，用于 NVIDIA GPU 的 cudaLaunchKernel)
- 好处：
  - 这确保了良好的覆盖范围，
  - 同时具有可伸缩性，因为其他定义客制化内核的库，最终都要通过launch API。
- Singularity有一个自动的代码生成器，可以生成所有 $D_{Int}$ 的存根; 它只需要一个特定加速器库(CUDA)的头文件列表，以及一些注释来指示状态变化调用

$SA_{Int}$ 的硬件抽象

虽然 $D_{Int}$ 的是自动生成的，但是大部分的自定义功能(例如，分布式屏障，用于时分的上下文切换)都驻留在 $SA_{Int}$ 中。
$SA_{Int}$ 中的大多数逻辑都是与设备无关的，并且使用一个映射到设备特定 API 的硬件抽象层(例如 nccl_allreduce)。
加速器的硬件抽象层(HAL)封装了加速器通用的关键功能。虽然当前的实现是特定于 NVIDIA GPU 的，但是要处理一种新的设备类型，只需要实现该设备的硬件抽象层，将该设备的特定 API 与 HAL 中的等效 API 映射在一起。
有三种与设备无关的功能需要使用$SA_{Int}$: 内存分配、通信和设备同步。
- 内存分配。内存分配API (例如 cudaMalloc、 cudaFree)需要一个 $SA_{Int}$ ，因为设备代理接管了内存分配。这使得设备代理对 GPU 内存中实际使用的区域具有完全的可见性，这有助于减少checkpoint大小。它还允许设备代理使用自定义内存分配机制，以帮助在同一 GPU 上对多个Worker进行透明的时分，从而获得弹性。
- 通信。多数加速器都有集体通信库(例如，NVIDIA 图形处理器的 NCCL，AMD 图形处理器的 RCCL)。在这些 API 上使用 $SA_{Int}$ 可以实现分布式屏障的算法，以便在分布式作业的多个Worker之间进行同步，从而获得一致的checkpoint。Singularity通过搭载相同的通信 API来提供了一个通用的障碍实现。这些 API 上的 $SA_{Int}$ 还有助于在弹性时分期间管理集合调用。
- 设备同步。设备同步 API 需要 $SA_{Int}$来处理透明弹性。Singularity中的时分是语义感知的，因为必须正确处理跨时间分片级别的通信。正确处理同步 API (例如 CUDA 中的 cudaStreamWaitEvent)对于时分的正确性和活性至关重要。

宿主机特定功能的$SA_{Int}$

Singularity还使用 $SA_{Int}$ 来选择 CPU 库
- 特别是 libc 的 I/O 库(例如，open、 read、 write 等)会被拦截来跟踪/记录作业对本地文件系统的更新，这样变异的文件就可以随着进程checkpoint一起迁移。
- 与设备库 API 的 $SA_{Int}$ 不同，主机 $SA_{Int}$ 没有相应的 $D_{Int}$ ，而是在主机地址空间中运行。

透明迁移的设计

在Singularity中，运行 DNN 作业的抢占、恢复和调整涉及到一致的checkpoint和恢复四大状态
- CPU 中的程序状态(例如堆栈、堆、指令指针等)
- GPU 中的模型训练状态(例如模型参数、优化器状态等)
- 处理 CPU 和 GPU 之间交互的控制状态(活动流、同步事件)
- 处理不同类型并行性(数据/管道/张量并行等)的 GPU 间和节点间通信状态
对于调度迁移，以及从计划外的故障中恢复，Singularity的透明checkpoint逻辑在两种模式下执行
- 当调度程序决定需要抢占作业时，基于扩展命令的随需应变
- 基于用户指定的间隔(时代级或基于时间)。
实现通用 DNN 作业的透明checkpoint是具有挑战性的，原因有
- 首先，在检查一个给定作业的时候，Singularity必须确保跨越多个主机和 GPU 的分布式状态的一致性; 一个分布式作业的所有Worker必须在集体通信方面处于一个安全和一致的状态(例如，allreduce)。
- 其次，CPU 和 GPU 之间的运行状态(例如，活动句柄，存储在主机内存中的设备地址)必须始终如一地恢复，尽管状态管理是由专有的闭源库(如 CUDA)完成的。
- 第三，checkpoint的空间开销必须保持在较低的水平，以适应有数百个Worker的大型分布式作业。

Checkpoint程序状态（CPU）

有多种系统提供地址空间迁移，其中 CRIU 是使用最广泛的。
- 然而，CRIU 的一个关键限制是它不处理使用 GPU 的进程的设备映射。
- 要使用 CRIU，主机地址空间必须与特定于设备的库隔离。
- 幸运的是，设备代理为我们提供了这种隔离。设备代理服务器大部分是无状态的，因此没有checkpoint; 它只是在目的地重新启动。

Checkpoint设备状态

保存方式：模型状态(例如，参数)由设备代理进程通过device to host的 memcpy 进行checkpoint。
特点：由于 Singularity中的内存分配 $SA_{Int}$，设备代理知道 GPU 内存的哪些区域实际上正在使用，因此显著减少了checkpoint大小。
挑战：在目的地恢复时，设备内存可能被映射到新设备-代理服务器地址空间中的不同地址，从而使主机进程中的指针失效。
内存地址映射：
- 为了避免这种情况，设备代理在启动时占用了整个 GPU 内存(被设备库跟踪的状态有些松弛)
- 有一个服务器$SA_{Int}$ ，服务于设备分配器(cudaMalloc)所执行的mmap，保证始终映射到相同的 CPU 地址。
设备句柄映射：保证CPU中设备句柄的一致性
- 与内存指针类似，主机地址空间还保留指向设备状态的其他句柄。
  例如，cudaStreamCreate 返回一个不透明的句柄，主机可以在后续的 GPU 调用中将其用作引用。但是，由于设备代理服务器在迁移后重新启动，句柄将无效。
- 句柄虚拟化：为了在迁移过程中保持这些句柄的保真度，我们对这些句柄进行虚拟化。设备代理不返回设备返回的实际句柄，而是返回一个虚拟句柄，并将此映射作为客户端状态的一部分进行记忆。
- 效果：恢复和重播后，物理句柄可能会更改，但虚拟句柄保持稳定。所有有状态的 API 调用(例如，创建上下文、流、事件等)都有注释，这些调用的 $D_{Int}$ 会自动记录它们，以便在恢复时重播。

通信状态

静默：DNN Worker之间的大多数通信都是通过集体通信库(例如 NCCL)进行的，我们无法处理飞行中的通信。因此，在checkpoint时，我们静默所有job，以确保没有飞行中的集体通信调用。
死锁产生：
- 要完成一个集体调用(例如 allreduce) ，需要所有参与的Worker都完成该调用
- 如果一个Worker在第 n 次 allreduce 调用已经返回后进行了checkpoint，而另一个Worker可能已经发出了第 n+1次调用，这个调用永远不会完成，从而导致死锁。
- 因此，在checkpoint之前，所有Worker必须完成同一个集体调用。Singularity使用了一种新的分布式屏障算法，以完全透明的方式实现这一特性。
分布式Barrier
- meta-allreduces：为了避免引入新的故障路径，Singularity中的障碍算法依赖于作业用于集体通信的同一个通信库，通过引入附加的meta-allreduces来交换障碍协议状态。需要确保附加的meta-allreduces相对于常规的allreduces的排序在所有节点间是一致的，以满足避免死锁的程序顺序要求。
- 算法：为了保证程序顺序一致性，在每一个数据allreduce操作执行之前，我们的算法会发射一个异步串联meta-allreduce。
  - Worker包含两个阶段
    - 第1阶段是稳定状态
    - 第2阶段是接收到障碍请求。
  - 串联 meta-allreduce 是有效载荷上的一个 SUM allreduce，由两个整数组成:
    - $Need_ {bar}$: 如果一个 worker 已经接收到了一个屏障命令，那么它会发送一个“1”，否则发送“0”。如果 SUM (need) > 0，则Worker知道有人已经启动了屏障协议，并切换到第2阶段。
    - $Ack_{block}$: 如果Worker已经切换到第二阶段，它将发送一个“1”，也就是说，它承认它已经直接或间接地看到了一个障碍请求，否则发送“0”。如果 SUM (ack) = = world _ size (即rank的总数) ，则Worker知道每个人都已经认可，并且可以安全地获得障碍。
  - 一旦Worker进入阶段2，它就进入了同步模式: 该Worker执行的每个集体调用都是同步的; 这确保了屏障协议的及时终止。
- 开销：开销很小障碍算法保证在最多两个小批内完成，并保证在屏障时没有飞行中的集体调用。它在阶段1的开销非常小，因为meta-allreduce只需要2字节，而且是异步的
- 局限性：适用于数据并行作业
关于tensor并行和pipline并行
- 张量并行和管道并行作业还有额外的复杂性，这些作业可以在不同的节点组之间执行多个 allreduce，此外还有对等调用，如 send/recv (用于流水线)。
- 解决方案：
  - 我们可以扩展上述算法，来推断出流水线和数据通信的相对顺序
  - 为了尽可能简单和减少checkpoint的大小，我们通过使用领域知识，确定一个小批的结束，在这一时间点上没有飞行中的通信，无论是tensor并行和pipline并行。
  - 我们使用与上面相同的串联 meta-allreduce 协议，但是在小批处理结束时只使用一次来实现障碍
- tradeoff：将障碍延迟到mini-batch结束(对于大型模型来说只需要几秒钟) ，带来了时间延迟，但是与在mini-batch执行中进行屏障相比，减小了checkpoint大小

文件系统状态

背景：Worker有时安装本地软件包，更新本地文件。迁移到新节点后需要保留这些内容。
挑战：执行容器维度的文件系统状态diff操作的代价太高。
解决方案：
- 使用Libc 文件系统 API 上的host $SA_{Int}$
- 每当以可写模式打开本地文件时，我们将文件名附加到日志中，并在checkpoint期间复制这些文件。
- 通过使用内容checksum，在远程存储的数据副本实现跨Worker去重。

Checkpoint/Restore工作流

工作流
- 在成功获得一个屏障之后，每个Worker使用CRIU执行checkpoint。
- CRIU镜像和GPU状态镜像，随后被移动到远程存储。
- CPU：在新的目的地上，使用CRIU的restore，程序从它被checkpoint的状态重新开始
- GPU：
  - 设备代理客户端执行的第一个操作是重新生成一个新的设备代理服务器，然后重播状态，使 GPU 恢复到checkpoint之前的状态。
  - 设备代理服务器还将 GPU 张量复制回checkpoint之前的相同地址的 GPU RAM。
设备-代理最终执行一个新的会合，worker可以发现彼此的新位置，并重新建立通信环。
checkpoint执行时机：除了由调度程序启动的随需checkpoint之外，Singularity中的每个作业都以用户指定的频率(例如，每30分钟)接受checkpoint，以处理计划外的故障。

压缩Checkpoints

Singularity采用多种技术来减小checkpoint的大小。

GPU镜像大小：
- Singularity执行每个缓冲区的内容校验求和，以便在Worker之间进行去重。
- 只有在没有其他Worker上传相同的缓冲区时才上传缓冲区
- 通过这种方式，Singularity中 GPU镜像的大小与用户级checkpoint的大小相似。
CRIU镜像的CPU 地址空间是在空间和时间两方面进行去重的
- 空间：在主训练过程和数据加载过程之间存在高度的内容重叠; 我们拦截 CRIU 发出的写调用来执行基于内容散列的页面分解
- 时间：在不同时间点采取的同一进程的checkpoint之间存在高度的重叠(因为地址空间变化很小) ; 时间维度上的去重使得之后的增量checkpoint比第一个 CRIU checkpoint小得多。

透明弹性的设计

Singularity的弹性实现是不改变世界大小的，改变的是worker和device之间的映射。
透明弹性建立在Singularity中的透明迁移支持的基础之上。
技术挑战：
- 时分复用：一个训练工作的多个Worker在同一个 GPU 上进行时分共享时，Worker之间的细粒度通信必须像在不同的 GPU 上一样
- 空分复用：对于大模型来说，每个Worker几乎可以利用 GPU 上的整个RAM; 在同一个 GPU 上运行多个Worker，需要将 GPU 状态在设备和内存之间来回交换，代价昂贵。
- 调度死锁：为了支持使用数据并行、流水线并行和张量并行相结合的作业，需要在 GPU 上仔细安排Worker，以便数据并行副本和相同的模型并行碎片在同一GPU上进行时分，并防止通信调度中出现死锁。

语义感知的时分

原因：
- 一个独立的Worker占有独立的内存是不可取的，因为大模型的每个Worker往往需要整个GPU RAM，这样的运行方式会耗尽内存
- 所以，时分需要语义感知
Singularity中的设备代理使得这样的时分成为可能。
- 因为设备代理与主机进程解耦，所以我们在多个主机进程(即多个rank)之间共享相同的设备代理。
- 当所有与 GPU 的交互都通过设备-代理进行时，它智能地安排多个rank，在给定的时间只允许一个rank在 GPU 上执行，然后选择特定的点来切换到另一个rank。
- 从概念上讲，在上下文切换的时候，设备代理将原始rank使用的 GPU 内存换出，然后换入新rank的GPU 内存，从而使每个Worker能够使用几乎整个 GPU 内存。
为了保持低开销，我们必须在绝对必要的情况下切换上下文。
- 在后向传播之后，数据-并行rank参与集体通信以交换梯度，这需要所有rank参与(并贡献其各自的梯度) ，必须进行上下文切换。
- 请注意，在单个小批处理中，为了重叠使计算与通信开销，框架可能会发出多个异步 allreduce 调用，然后在某个点进行同步。在这个同步点，设备代理切换到共享 GPU 的下一个rank，独占地运行，直到达到同步点，然后上下文切换到下一个rank，以此类推。
集体通信通过专有库进行。
- 通信器：NCCL 有一个通信器的概念，它被初始化为一个特定的参与rank环，随后的操作只是引用通信器。
- 解耦： 为了保持 NCCL 通信器与用户级时分的交互可控，我们将作业的逻辑数据并行世界大小与 NCCL 看到的世界大小解耦; 在我们的方法中，NCCL 只看到每个 GPU 一个rank。
- 梯度累计：在时分过程中，设备-代理透明地在抓取缓冲区中进行本地累积，并且只有共享 GPU 的最后一个rank使用本地累计梯度执行实际的nccl_allreduce
- 大小调整：因此，在一个调整大小的操作之后，NCCL 看到的世界大小发生了变化，恢复后由新的会和处理

内存共享的副本拼接

副本拼接使上下文切换更加便宜 –>5.2.1

基于checksum的动态数据去重

一项训练工作所消耗的 GPU 内存分为四类:
- 参数(P)。模型的每一层的权重/参数; 正向和反向传递在这些张量上运行。
- 优化器状态(O)。由优化器跟踪的状态，以计算每次迭代应用于参数的增量。跟踪历史状态(例如，第一和第二阶段的梯度)
- **梯度(G)。**每个副本都有与其迷你批处理相对应的渐变副本。在向后传递之后，对所有副本的梯度取平均值，然后使用这个平均值一致地更新权重
- 激活(A)。每一层的正向通道的中间输出; 在反向通道中用于计算相对于输入的梯度以进行反向传播。
副本拼接利用的洞见
- 在数据并行副本中，P和O在小批处理结束时由所有副本一致地更新，所有rank具有相同的平均梯度
- 在mini-batch结束后，因为后向传播已经结束， A会被框架释放掉
由于设备代理控制内存分配器，它对框架分配的每个缓冲区都具有可见性。
- 在上下文切换期间，设备代理为每个活动缓冲区计算内容校验和。
- 在换出过程中，它首先查看主机是否已经包含一个具有相同内容校验和的缓冲区，如果是，它避免换出，并简单地将 GPU 缓冲区标记为未使用
- 在换入过程中，它检查设备是否已经有一个带校验和的缓冲区; 如果有，它就避免了来自主机的换入。请注意，虽然内容匹配，在新的rank该缓冲区可能被映射到一个不同的设备地址。在这种情况下，设备代理执行设备到设备的移动，将缓冲区移到所需的地址
一个例子：
- 通过上述优化，如果4个rank共享一个 GPU，在上下文切换期间P和O缓冲区的交换只需要在第一个rank完成; 其他人会发现校验和已经存在于主机内存中，并忽略掉交换。
- 然而，对于每个rank来说换入仍然必须进行：当一个rank开始其时间片时，其本地状态包含来自前一个mini-batch的P和O，而前一个rank的P和O副本被更新到当前的mini-batch。
- 如果我们有空间存储两个额外的 P 和 O 版本在 GPU，我们就可以避免这个换入，这带来了两个挑战:
  - 大模型一般没有额外的空间容纳额外两组P和O –> 5.2.3
  - 我们仍然需要执行设备到设备的 P 和 O 副本在上下文切换，因为每个rank可能分配了在不同的地址分配了相同的缓冲区，D2D 复制开销仍然很大。–> 5.2.2

为持久分配使用领域知识

使用深度学习训练的领域知识来使地址保持一致
- 在数据并行副本中，根据定义，稳定缓冲区（如 P 和 O 在小批处理中保留的分配序列）的大小、顺序必须在所有副本中相同，因为它们具有相同的参数集。
- 然而，可能存在跨副本大小可变的其他分配(例如，大小取决于输入数据大小的激活，这种激活可能在不同的小批处理中有所不同)。由于这种可变大小的分配，内存分配器的状态在不同的副本之间发生差异，甚至导致稳定的缓冲区分配得到错误对齐的地址。
为了处理这个问题，Singularity中的设备代理使用一个双向内存分配器。
- 稳定缓冲区(如 P 和 O)在地址空间的高端分配，而其他缓冲区在低端分配。这确保了瞬态分配(例如激活)中的不稳定性不会影响高区域中的内存分配器元数据，从而确保稳定的缓冲区(例如 P 和 O)跨副本获得相同的地址。
- 为了识别稳定的缓冲区，比如 P 和 O，我们在分配器中添加了一个预先识别的堆栈跟踪列表(Python 和 C + +) ，这个列表与参数和优化器状态分配有关。
- 在分配时，最多有两个版本的 P 和 O 是活跃的-当前的小批量和以前的小批量，并且第三个副本需要作为草稿空间，以便当前的rank不会覆盖以前的小批量的原始版本的 P 和 O。

压缩选定的操作

使用了特定领域的洞察避免处理 P 和 O 的多个副本。
- 根据定义，所有数据并行副本将在完成小批处理之后使用一致的P和O缓冲区。
- 我们还知道，P和O缓冲区只有在所有副本的梯度allreduce的之后才会更新。
- 因此，如果我们能够识别更新参数和优化器状态的操作，我们就可以只在共享设备的其中一个rank中执行这些操作，并简单地在其他rank中“压缩”这些操作，因为
  - 它们无论如何都会导致相同的最终状态，
  - 缓冲区在队列中有相同的对应地址，所以随后的小批量计算将看到正确的数据。
- 为了压缩一个操作，设备代理只是省略了为这些操作向 GPU 发出 CudaLaunchKernel。通过这种压缩，我们避免了在前一个mini-batch的P和O的换入，因为它们不再由其他rank更新。
为了确保健壮性，我们遵循保守验证的方法
- 对照组：我们总是运行第一个mini-batch，同时禁用了压缩(因此产生了交换进/交换出成本) ; 这保证了正确的执行。
- 实验组：在验证mini-batch中，断言假设，验证使用缓冲区内容校验和来推断事后操作效果的方法：在验证mini-batch中，我们验证模型是否符合以下不变量:
  - 在压缩窗口期间的所有缓冲区更改必须在共享 GPU 的所有rank之间是相同的。
  - 在压缩窗口期间执行的设备到主机的副本必须在共享 GPU 的所有rank之间完全复制相同的数据。
- 如果上述验证失败，我们将模型视为压缩的不安全模型，并退回到基于交换的机制(如果必要，将作业“回滚”到验证成功的最后一个checkpoint)。
- 我们监视由于时分造成的开销，如果它超过一个阈值，我们将禁用该模型的时分。这将是一个罕见的场景，但仍然需要优雅地处理以获得健壮性。

处理模型并行工作

张量并联、流水线并联等模型并联作业的处理带来了新的挑战。
- 张量并行job中的每个矩阵乘法执行allreduce。如果我们为这样的 allreduce 进行上下文切换，那么副本拼接将无法工作，因为激活张量A仍然是活跃的。
- pipeline并行job对每个mini-batch执行 GPU/节点之间的端到端地发送与接收激活A和梯度G; 在mini-batch期间进行时分会由于活跃的A和G导致过多的交换。
为了应对这些挑战，Singularity使用了两种关键技术:
- 拼接感知的放置。通过拼接感知的放置，我们确保只有同一个的模型并行分区的数据并行副本在同一个 GPU 上进行时分。
  - 请注意，这要求 Singularity了解rank分配逻辑。对于使用具有不同rank分配策略的自定义启动程序的作业，Singularity为该作业提供了一个 API，用于传递所有rank的rank到拓扑的映射(例如，Rank 4是 DP0、 MP0、 PP1等)
- 推断集体调用的意图。设备代理推断集体通信的意图，并仅在数据并行维度上集体调用时触发时分。其他的集体调用只是简单地通过而没有上下文切换。推断方法：
  - Singularity利用集体通信的初始化路径(例如，ncclCommInitRank)来实现这一点。它强制在每个ncclCommInitRank 之后进行上下文切换，并且设备代理计数每个通信者。
  - 在完成一轮上下文切换之后，如果通信器的本地计数大于1，则设备代理将推断该通信器处于数据并行维度。因为拼接感知的放置。
  - 在集体调用期间，它只需要在通信器上查找一个map，就可以了解它是否是数据并行的。

处理ZeRO-冗余优化器

ZeRO切分数据并行状态，这样数据并行的Worker之间就没有冗余。这样的分区违反了压缩验证的不变量(5.2.3)。
为了解决这个问题，Singularity为 ZeRO 引入了部分分片的概念，它将分片因子(在 GPU 中拟合模型所需的最小值)与数据并行度(并行度)分离开来。
- 如果两者是相等的，那么根据定义，这个模型不能缩小到更少的 GPU，因为它不适合。
- 如果数据并行性因子更高，比如说，4倍于分片因子) ，那么我们可以支持多达4路的时分/缩小。在这个场景中，部分分片因子只是变成了模型的另一个维度——并行性，并且只有相同的 ZERO 分片的副本是时间分片的。
在 DeepSpeed 中引入部分分片非常简单

实现

实现方面的挑战
- 序列化不透明参数。在我们的基于拦截的设备代理中，CudaLaunchKernel 的 $D_{Int}$ 是具有挑战性的，因为它的签名是不透明的，使得序列化变得困难(签名是由 NVIDIA 的1 nvcc 在内部生成的，拦截器不可见)。为了处理这个问题，我们有一个自定义的服务器 $SA_{Int}$ ，它使用 cuObjecDump，CUDA 工具包中的一个二进制实用程序，解析生成的内核库并提取参数信息。为了避免高成本，我们缓存此信息，并仅在缓存丢失时运行 cuObjecdump。对于 JIT 内核，我们拦截 nvrtcCompileProgram 并通过解析生成的 PTX 提取参数签名。
- 隐藏调度延迟。设备代理中的跨地址空间调用发生在诸如 cudaLaunchKernel 和 cudaGetLastError 等操作的关键路径上，这会影响性能。我们对最频繁的调用使用特定于域的优化。对于 cudaGetLastError，我们在每次启动内核时都会在服务器上发布它，并将它与响应一起附加，这样当 PyTorch 发布它时，设备代理客户端就可以从缓存中返回它。对于 cudaLaunchKernel，我们执行延迟错误通知，调用在客户端返回，而不等待来自设备代理服务器的响应; 在向服务器发出下一个调用之前，对响应进行延迟读取，因此允许客户端 PyTorch 处理和服务器引起的延迟之间重叠; 因为当遇到这种(罕见的)错误时，PyTorch (和其他框架)崩溃，这不会影响作业语义。
- 隐藏上下文切换开销。在时分过程中从一个rank切换到另一个rank涉及到计算所有活动设备张量的校验和，并将它们与另一个rank的副本进行比较，如果需要，执行缓冲区移动(几毫秒的 CPU 活动)。此外，开关逻辑依赖于校验和计算的输出，校验和计算依次等待所有以前的 GPU 操作完成，隐含地强制设备同步。为了避免在关键路径中产生这种代价，Singularity执行下一级的急切分派。设备代理开始并行地服务于下一级(CPU 逻辑、设备操作分派)的有用工作和交换延迟。通过仔细使用诸如 cudaStreamWaitEvent 之类的异步排序原语，我们确保只有在切换完成之后，才能在 GPU 上执行对新排序的操作。

【论文笔记】Singularity论文笔记