PipeDream: Generalized Pipeline Parallelism for DNN Training

1.Title:
PipeDream: Generalized Pipeline Parallelism for DNN Training
Author: Deepak Narayanan,Aaron Harlap(Microsoft Research, Stanford University)
来源：SOSP 2019
2.Main idea:

• 流水线并行相对与批内并行的优势
  相对于DP，通信量会减小，在流水线并行中，每个worker只需要向下游worker发送激活信息或者是向上游worker发送梯度信息即可，而不用像DP中，必须聚合所有的梯度信息进行平均，再发给所有worker。
  使用流水线并行会使得计算与通信重叠。当前minibatch的跨阶段的前向激活与反向梯度的异步通信会与后面的minibatch的计算重叠。

• 采用的并行方式：数据并行 + 模型并行 + 流水线技术
  批内并行+批间流水线

• 自动划分机制 Work Partitioning
  DNN训练有一个特点：不同输入的计算时间几乎没有变化。于是PipeDream充分利用了这一事实，给定一个具有N层和M台可用机器的DNN，PipeDream首先在一台机器上分析模型，记录向前和向后过程所花费的计算时间，每层输出的大小以及每个层的相关参数的大小，最后输出为一个profile文件。
  然后PipeDream 优化器利用生成的profile文件，再结合当前内存、硬件以及worker数量等情况，合理地划分stage，保证每个stage大致执行相同的总工作量，这样才能保证整个流水线的平衡（不然总是在某一stage执行时间过长，流水线不稳定；整个流水线的吞吐量以最慢的stage的吞吐量为限制）
  

• 1F1B策略 Work Scheduling
  由于前向计算的activation需要等到对应的后向计算完成后才能释放（无论有没有使用 Checkpointing 技术），因此在流水并行下，如果想尽可能节省缓存 activation 的份数，就要尽量缩短每份 activation 保存的时间，也就是让每份 activation 都尽可能早的释放，所以要让每个 micro-batch 的数据尽可能早的完成后向计算，因此需要把反向传播的优先级提高，让micro-batch标号小的后向比 micro-batch 标号大的前向先做。因此，如果我们让最后一个stage在做完一次 micro-batch 的前向后，立马就做本micro-batch的后向，那么我们就能让其他的 stage 尽可能早的开始后向计算，这就是 1F1B 策略。
  
  以worker 3正向计算micro-batch 5为例，正向计算micro-batch 5 时，可以同时进行梯度的通信，把之前计算得到的micro-batch 3 的梯度信息反向传播给worker 2；之后反向传播micro-batch 4 时，可以同时进行激活（micro-batch 5 正向计算的结果）的通信，把激活值传给worker 4；
  所以说，计算和通信之间是有时间重叠的。
  

• 为了解决同一stage下针对同一minibatch的正向计算与反向传
  的权重版本不一致的问题，PipeDream使用了权重存储技术以及垂直同步。Effective Learning
  所谓的正向计算与反向传播使用不同的权重参数是指，worker1在对minibatch5进行正向计算时，使用的是minibatch1反向传播后更新的参数，而之后再对minibatch5进行反向传播的计算时，使用的是minibatch2,3,4反向传播后更新的权重参数，这就导致了同一stage下对同一minibatch进行正向计算与反向传播使用的权重参数不一致！
  权重存储技术维护多个版本的权重，每个活动的minibatch对应一个版本，每个stage使用最新的版本实现minibatch的正向计算，正向计算完成之后，pipedream会存储这个版本的权重，将来会重新使用这个版本的权重来进行minibatch的反向传播计算梯度信息。也就是说，权重存储确保了在一个stage中，相同版本的模型参数用于给定minibatch的正向计算和反向传播！但是并不能保证对于给定的minibatch在不同stage使用的版本参数的一致性。
  垂直同步：每个进入流水线的minibatch会与当前输入stage中已知的最新的权重参数版本产生关联，并且这个版本的权重参数会应用在这个minibatch的正向计算与反向传播的整个过程。这样就能保证minibatch在不同stage使用的权重参数的一致性！
  
  使用weight stashing之后，对于worker1，完成minibatch1的反向传播后，进行参数更新，假设更新后的版本为w1，那么之后worker1执行对minibatch5的正向计算时，会使用w1权重，并且会在worker1中保存这个版本的权重，当worker1执行minibatch5的反向传播时，会再重新应用w1进行梯度计算。这就实现了同一个stage下对同一minibatch正向计算与反向传播应用相同的权重参数。但同样会存在一个问题，对于同一minibatch，跨stage使用的权重参数是不一样的。
  未使用垂直同步，worker1在计算完minibatch1后，进行参数更新，设此时参数版本为w1，之后，minibatch5进入管道，使用版本w1，接着向下游正向传播，而worker2计算完minibatch2后又会更新本地的权重参数，设此时的版本为w2。那么worker2在正向计算minibatch5的时候，会使用参数版本w2，这就会导致同一个minibatch在不同stage下使用的权重参数不一样。
  使用垂直同步后，worker1在完成minibatch1的反向传播之后，会更新参数，设更新完后的参数版本为w1，那么之后，当worker1对进入管道的minibatch5进行正向计算的时候，会把这个这个版本的参数与minibatch5关联，每个worker都会在本地维护这一版本的参数。并应用到整个minibatch5的正向计算与反向传播过程，直至再回到worker1，此时更新权重，将w1从内存中删除。
  使用Vertical Sync之后，对于进入管道的minibatch b_i，会与当前最新的权重参数版本w_(i-x)相关联，并且会随着b_i的激活和梯度信息在整个管道中传播，直至该b_i的反向传播完成，再次进行参数更新，这时，删除w_(i-x),保存w_i.

3.Experiment
实验配置

• 与DP比较
  使用集群A的四台服务器，每台服务器有4张卡，比较了分别使用pipedream和DP对VGG-16模型进行训练的时间和准确性。由于使用PipeDream会明显减少worker之间的通信时间，这就使得使用PipeDream之后相对于DP的加速比达到了5.28
• 与除了DP的其他批内并行方法比较，如MP，HP混合并行
  （a）中，对于四个模型，只在模型并行的基础上增加流水线技术就可以获得近2倍的加速比，而对于VGG-16和AlexNet，再对某些stage进行replicate，放到不同的worker上进行DP，将会获得14.9和6.5的加速比。
  （b）中，主要是对于PipeDream和一些混合并行算法进行了比较，比如卷积层DP+FC层MP的方式；可以看出来，使用PipeDream也获得了不错的加速效果，
• 与批间并行方法比较（主要是跟Gpipe比较）
  针对模型GNMT，分别使用PipeDream和Gpipe进行训练加速，在Cluster A 上，Gpipe比PipeDream吞吐量低了55%，在Cluster B上，低了71%。这种吞吐量衰减的主要原因是，在Gpipe中，会频繁地向管道中注入microbatch，也就是频繁地进行管道刷新，这样会大大减少计算资源地利用率，也会降低整体地吞吐量！