Nvidia GPU热迁移-Singularity

1 背景

在GPU虚拟化和池化的加持下，可以显著提高集群的GPU利用率，同时也可以较好地实现弹性伸缩。但有时会遇到需要GPU资源再分配的场景，此时亟需集群拥有GPU任务热迁移的能力。举个简单的例子，比如某个新训练任务M需要独占8张A100，而集群中没有独立的A100资源，但总的A100的碎片资源超过了8张独立的A100资源，此时借助虚拟化、池化和热迁移有以下几种解决方案：

• 虚拟化：拆分任务M至n个独立的小任务，使其可以使用A100碎片资源
• 池化：拆分任务M至n个独立的小任务，使其可以使用A100碎片资源
• 热迁移：将其他A100上的任务迁移至同一张卡上，重新凑出任务M需要的A100资源

可以发现，基于热迁移方案，不需要用户对任务M进行修改，只需要调度侧重新进行GPU资源分配即可解决。

2 Singularity

Singularity是微软为AI工作负载提供的完全托管的全球分布式基础设施服务，支持各种硬件加速器。Singularity从头开始设计，在全球范围内扩展到数十万个GPU和其他AI加速器。Singularity的建立有一个关键目标：通过在给定的固定加速器容量池上最大限度地提高总有用吞吐量来降低人工智能的成本，同时为多个定价层提供严格的SLA。下图展示了Singularity的高级架构，包括其由全局、区域和工作负载范围的调度微服务组成的分层调度系统。

论文中涉及到比较多的内容，比如透明且解耦的机制、API拦截的分类、透明迁移和透明弹性等，本文主要关注透明迁移，即GPU任务的热迁移，对其他部分感兴趣的可以查看原文。

3 热迁移

在Singularity中，DNN任务的抢占、恢复和放缩主要涉及到以下4种类型的状态的一致性的检查点和恢复。

• CPU中的程序状态：如栈、堆、指令指针等
• GPU中的模型训练状态：如模型参数、优化器状态等
• CPU与GPU之间交互的控制状态：如活跃流、同步事件等
• 不同类型并行的gpu间和节点间通信状态：如数据、流水线、张量并行等

实现DNN任务的透明检查点具有挑战性，主要体现在以下几个方面：

• 多台分布式主机和GPU的一致切割，所有的分布式任务必须在集体通信方面（比如allreduce等）处于安全一致的状态
• CPU与GPU之间飞行中的状态（比如活跃句柄、存储在主机内存中的设备地址等）必须一致恢复
• 数百个分布式任务的检查点开销必须保持在较低的水平

3.1 检查点中CPU的程序状态

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

3.2 检查点中的设备状态

模型状态（比如参数等）由设备代理进程通过设备到主机之间的memcpy进行检查。一个挑战是，当在目的地进行恢复时，设备内存可能会映射到新的设备代理服务器地址空间中的另一个地址，从而使主机进程中的指针无效。为了避免这种情况，设备代理在启动时占用整个GPU内存，由设备分配器的服务执行mmap（比如cudaMalloc等），以始终映射到相同的CPU地址。

3.3 通信状态

DNN任务之间的大多数通信都是通过处理GPU-Direct（比如NVLink）和跨节点通信的集体通信库（比如NCCL）进行的。由于这些库是专有的，我们不能处理飞行中的通信。因此，在检查点时，我们暂停该任务，以确保没有正在运行的集体调用。Singularity采用一种新颖的分布式屏障算法，以完全透明的方式实现这一特性。

3.4 文件系统状态

DNN任务有时会安装本地包，并更新其他本地文件。迁移到新节点后需要保留这些文件。在容器范围内执行文件系统状态的差异（相对于干净的基础镜像）代价太大。主机libc文件系统API可以帮助实现这一点:每当以可写模式打开本地文件时，我们将文件名附加到日志中，并在检查点期间复制这些文件。通过使用内容校验和，在多个任务之间将非重复数据复制到远程存储。

3.5 检查点/恢复流

成功获得障碍后，通过对各个工作点执行CRIU检查点来检查任务。CRIU转储，以及GPU状态转储活跃张量，然后移动到远程存储。在新的目的地，在CRIU恢复中，进程从检查点的确切状态开始（即在设备代理客户端中获得障碍之后）。

4 其他

4.1 GPU算力化

具备GPU虚拟化、池化和热迁移的集群，留给集群的想象空间比较大，拥有着较好的调度和弹性的能力，但距离极致的调度和弹性的能力，还有一些距离。再以上述M任务需要8张A100为例，如果集群中不仅没有独立的A100资源，也没有足够的A100碎片资源，但是却有着大量的V100资源，此时似乎存在一种可能性，用足量的V100资源替代M任务需要的A100资源，但又会引入新问题，GPU卡之间怎么换算，如何做GPU算力归一化，如何做GPU算力售卖的透明性。