【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就足够了。
实验
end-to-end测试，比较 read model, stop-and-start, NVIDIA MPS, PipeSwitch
- 推理任务到达。延迟、总开销、第一层的启动开销。
- 不同调度周期下，吞吐量、端到端延迟。
- 模型传送。测延迟。比较下面的：
以及剪枝策略的对照实验。
- 内存管理。比较：不用cuda自动管理（worker用cudaMalloc）；没有IPC优化的；没pin memory 的；cuda自动管理（worker掉用cudaMallocManaged）；PipeSwitch。测延迟。
- 上下文切换。比较：单进程、双进程、PipeSwitch。测延迟。
相关工作
借鉴了分布式DL训练的PipeDream 和 ByteScheduler。这俩着重于 inter-batch 流水线，以重叠计算和不同batch间的梯度传输，是面向同一DNN模型（任务）的训练的。PipeSwitch的创新处：使用intra-batch 流水线，交叠模型传输和计算，以减少不同DNN模型间（推理或训练）的切换开销。此外，不同进程间的内存管理、worker切换，也是要解决的问题。
vDNN和SwapAdvisor也有GPU内存管理，是针对大模型的单个训练任务的。
体会
个人感觉，不管设计还是行文都非常优雅。从最优分组传颂、内存管理、worker切换等多方面，逐点优化，在流水线效率和额外开销之间寻找平衡点，工程很饱满。