【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切换等多方面，逐点优化，在流水线效率和额定开销之间寻找平衡点，工程很丰满。