【AI sys】GPU 上 DL 工作 疾速上下文切换

PipeSwitch: Fast Pipelined Context Switching for Deep Learning Applications

https://www.usenix.org/confer…

（Johns Hopkins University & ByteDance）

简介

背景、动机

• DL 工作：吞吐敏感的训练任务、提早敏感的推理工作。为了保障推理的 SLOs，支流的设计是把它们离开部署在不同的 GPU 集群上。其弊病：推理工作负载低时（早晨），训练任务无奈利用推理集群闲暇的 GPU 资源；遇到flash crowd 时，推理工作无奈抢占训练集群的资源；为 SLOs 和限度不同工作间的推理，生产中常为一个推理利用调配单个 GPU。
• 若放同一 GPU 集群，则工作切换开销很大，可能是几秒，而推理的 SLOs 是数十数百毫秒。
• 以后应答切换开销的计划：空间共享 GPU 内存。但 MPS 和 Salus 等须要事后将过程的数据加载到 GPU 中，而 GPU 内存不足以预加载太多利用，且模型会增长。此外，这些计划没有给利用提供 GPU 强内存隔离。

要害思维

PipeSwitch 使得多 DL 工作在同一 GPU 上高效地分时复用，工作切换开销为毫秒级，且不耽误 SLO。

要害思维：DNN 层层重叠，计算是一层层的，PipeSwitch 利用这点，把模型传送、模型计算、GPU 内存替换的过程流水线化，以疾速切换上下文。同时，还要解决内存治理、worker 切换的问题。

奉献

• GPU 上 DL 工作的高效细粒度分时共享，毫秒级上下文切换开销，高吞吐量；
• 引入流水线上下文切换。蕴含对立内存治理，以及active-standby 机制，缩小切换开销、实现过程级隔离。
• 实现了零碎原型，并集成到 PyTorch 中。试验表明，工作启动 3.6-6.6ms，总开销比 NVIDIA MPS 好 10-50x，GPU 利用率靠近 100%。

PipeSwitch 总览

架构

• 控制器。接管客户端的工作，管制内存 daemon，管制 worker 执行工作。（TCP 线程、调度线程）
• 内存 daemon。治理 GPU 内存、DNN 模型。为沉闷 worker 调配 GPU 内存，将模型从主机内存传送到 GPU 内存。
• 沉闷 worker。执行以后 GPU 的工作。一个 server 蕴含一个沉闷 worker。
• 备用 worker。一个 server 有一或多个备用 worker。常闲暇。用来初试化新工作，或革除本身环境。

（worker 蕴含主线程、终止线程）

控制器和 worker 都须要用户注册模型。

工作执行

• 控制器接收客户端工作，调度工作，为抢占式（因为有推理工作）。
• 启动新工作：控制器告诉闲暇备用 worker 来初始化其环境；沉闷 worker 实现或中断当前任务后，控制器告诉内存 daemon 和备用 worker 去把工作加载到 GPU（流水线模型传送）；内存 daemon 分配内存该备用 worker，并把模型从主机内存传输到 GPU 内存。
• 备用 worker 变成新的沉闷 worker，执行新工作。
• 同时，先前的沉闷 worker 变成备用 worker，革除本身环境。

设计细节

GPU 上下文切换开销次要四局部：

• 工作革除。如开释 GPU 内存。
• 工作初始化。如启动过程，初始化 CUDA 上下文。
• GPU 内存调配。
• 模型传送。通过 PCIe，从主机 CPU 到 GPU。

各阶段工夫：数十毫秒到数秒。

流水线模型传输

推理工作只有前向船舶；训练任务的每次迭代都有前向流传、反向流传。两者工作都不须要等到整个模型加载到 GPU 中才开始计算，加载完一层即可。

PyTorch 中主动注入 hook，实现期待传输、同步执行。

加载一层，计算……带来的零碎开销：频繁的 PCIe 传输调用，有些层很小，PCIe 调用开销显得大；传输和计算间的同步开销。为此提出：

model-aware 分组。把多层并到一个组，在per-group 粒度上流水线化。用小的组，传输和计算间的重叠更多，可进步流水线效率，但额定开销增大；用大的组，额定开销小，但流水线重叠少。因而，分组需 model-aware。

如何找到最佳分组大小？两种剪枝办法。

$T(i, j)$：从层 $i$ 到 $ j$ 的一个组，传输工夫（蕴含调用 PCIe 的工夫）

$E(i, j)$：从层 $i $ 到 $ j$ 的一个组，执行工夫

$F(B, i)$：给定 $0$ 到 $i-1$ 层的分组 B，返回 $i$ 到 $n-1$ 层在最优分组策略下的工夫

$F$({}, $0$ ) = $ \min \limits_i $ $F$ ({$group(0, i $) }, $ i$ + 1)。（1）分为 $n$ cases。

剪枝 1：计算 lower bound。对应 Figure 3a。这是思考最佳状况：剩下的层放在一个组里，这样计算和传输就能够完满重叠，即计算能够在第一个组的计算实现后马上开始。若 case $i$ 的 lower bound 大于目前找到的最优工夫，则能够进行计算 $F$ ({$group(0, i $) }, $ i$ + 1)。

剪枝 2：Figure 3b。除了第一个组，能够将多个层打包在一起，而不会影响流水线的效率。假如固定第一组为层 $0$ 到 $i$，用递归(1) 枚举第二组。只有第二组的传输在第一组的计算实现前实现，就能够把第二组的传输暗藏到第一组。要分组的最小层数为：

能够免去计算 $(i+1)$ 到 $j < j^*$ 分组，只枚举 $ j \ge j^* $ 的状况。

算法：

$B$ 示意曾经分好的组，$x$ 示意还未分组的第一层。

$O(2^n)$。但剪枝可剪掉大部分的。

正确性证实。归纳法。

泛化。这个算法是在给定执行程序下，寻找最优的流水线分组法，也可利用于有环图模型。

对立内存治理

DL 模型很大，产生大量两头后果，很耗内存，且用 cuda 自带的内存治理接口会带来不必要的额定开销。

DL 工作的两类数据：模型自身（包含参数）、两头后果。前者是固定的（推理或训练都不会扭转模型构造）。两头后果的变动不会产生内存碎片：

• 推理工作。一层算完给下一层，下一层算完，上一层不再须要可清掉。
• 训练任务。前向流传的后果还要用于后向流传，两头后果是先进后出的，相似于栈。

栈式机制治理内存。

• 最小化内存调配开销。系统启动时，内存 daemon 应用 cudaMalloc 获取 GPU 内存，运行时动静分配内存（指针传给 worker）。daemon 保障一次只有一个 worker 占用 GPU 内存，以实现 worker 间内存隔离。一个 worker 应用一个内存池。
• 最小化内存 footprint，防止额定内存拷贝。内存 daemon 存模型，只需一份拷贝，且不须要把模型从专有过程拷到 worker。
• 最小化 IPC 开销。daemon 把相干的 GPU 内存句柄发给 worker。用 GPU 的来 IPC APIs 实现。daemon 和 worker 应用同样程序来为模型参数分配内存，这样许多参数的内存指针就会雷同，就只需一次 IPC 调用来初始化 worker。
• Pin memory。主机内存用作和 GPU 替换数据的页，pin 住，免得不沉闷时被替换到磁盘。

worker 切换

应用 active-standby 机制来实现疾速 worker 切换和过程级隔离。沉闷的执行以后 GPU 的工作，备用的在 CPU 期待下一个工作。 并行化：革除旧工作（沉闷 worker 执行），初始化工作（备用 worker 执行）。

注：革除，只清元数据，如 GPU 指针，而不开释内存。

控制器告诉沉闷 worker 进行，撤销对它的 GPU 内存调配，将内存分给新的沉闷 worker。

备用的太少，可能不能及时有闲暇的；太多，备用的保留上下文耗太多 GPU 内存。作者示意，两个 standby worker 就足够了。

试验

1. end-to-end 测试，比拟 read model, stop-and-start, NVIDIA MPS, PipeSwitch

   • 推理工作达到。提早、总开销、第一层的启动开销。
   • 不同调度周期下，吞吐量、端到端提早。
2. 模型传送。测提早。比拟上面的：

   以及剪枝策略的对照试验。

3. 内存治理。比拟：不必 cuda 主动治理（worker 用 cudaMalloc）；没有 IPC 优化的；没pin memory 的；cuda 主动治理（worker 掉用 cudaMallocManaged）；PipeSwitch。测提早。
4. 上下文切换。比拟：单过程、双过程、PipeSwitch。测提早。

相干工作

借鉴了分布式 DL 训练的 PipeDream 和 ByteScheduler。这俩着重于 inter-batch 流水线，以重叠计算和不同 batch 间的梯度传输，是面向同一 DNN 模型（工作）的训练的。PipeSwitch 的翻新处：应用intra-batch 流水线，交叠模型传输和计算，以缩小不同 DNN 模型间（推理或训练）的切换开销。此外，不同过程间的内存治理、worker 切换，也是要解决的问题。

vDNN 和 SwapAdvisor 也有 GPU 内存治理，是针对大模型的单个训练任务的。

领会

个人感觉，不论设计还是行文都十分优雅。从最优分组传颂、内存治理、worker 切换等多方面，逐点优化，在流水线效率和额定开销之间寻找平衡点，工程很丰满。