【OSDI'20】PipeSwitch: Fast Pipelined Context Switching for DL

【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 总览

架构

image

  • 控制器。接收客户端的任务,控制内存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,实现等待传输、同步执行。

image

加载一层,计算……带来的系统开销:频繁的PCIe传输调用,有些层很小,PCIe调用开销显得大;传输和计算间的同步开销。为此提出:

model-aware 分组。把多层并到一个组,在per-group 粒度上流水线化。用小的组,传输和计算间的重叠更多,可提高流水线效率,但额外开销增大;用大的组,额外开销小,但流水线重叠少。因此,分组需 model-aware

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

image

$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) 。

image

剪枝2:Figure 3b。除了第一个组,可以将多个层打包在一起,而不会影响流水线的效率。假设固定第一组为层$0$ 到 $i$ ,用递归(1) 枚举第二组。只要第二组的传输在第一组的计算完成前完成,就可以把第二组的传输隐藏到第一组。要分组的最小层数为:

image

可以免去计算 $(i+1)$ 到 $j < j^*$ 分组,只枚举 $ j \ge j^* $ 的情况。

image

算法:

$B$ 表示已经分好的组,$x$ 表示还未分组的第一层。

image

$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 切换

image

使用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. 模型传送。测延迟。比较下面的:

    image

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

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

你可能感兴趣的:(人工智能操作系统深度学习)