XPipe: Efficient Pipeline Model Parallelism for Multi-GPU DNN Training

1、 Authors:Lei Guan(NUDT), Wotao Yin(来自UCLA), Dongsheng Li, Xicheng Lu
论文来源：https://arxiv.org/pdf/1911.04610.pdf
2.、Main idea：

• workflow
   把一个minibatch分割为T个microbatch，也就是说这个microbatch就成为了整个pipeline中的最小data unit；每完成T个microbatch进行一次参数的更新。这里可以实现不同minibatch下的microbatch可以交叉训练(与pipedream中的1F1B相似)。
  
    当完成一个minibatch中的所有microbatch进行一次参数更新，而不是像pipedream中每完成一个microbatch就进行一次参数更新。但是这里存在两个问题：
    1）weight inconsistency，比如第一个GPU0在进行microbatch5的正向计算时，使用的是原始的权重，但是在进行microbatch5的反向传播时，权重版本已经进行了两次更新，即完成了两个minibatch的反向传播，模型参数完成了两次更新。
    2）weight staleness, 权重的陈旧性会影响模型的收敛，比如图(a)里面，GPU0在进行microbatch5的正向计算的时候，使用的是原始的权重版本，而在下游的GPU2进行microbatch的正向计算的时候，模型已经完成了一次权重更新，也就是说这个时候，使用的是新的权重，也就是说同一个microbatch在不同的GPU正向计算的时候使用的权重版本不一样。
• weight prediction
   为了解决上述异步管道训练出现的两个问题，基于Adam更新，提出了weight prediction。把一个minibatch中index最小的microbatch称为bellwether。具体的过程是：当一个minibatch喂入模型时，首先是第一个miicrobatch开始正向计算，也就是bellwether。那么在一开始的时候，就会先计算一个预测的权重W，此后，整个minibatch的所有microbatch都共享这份权重。那么如何计算权重？
   (1) 定义一个权重差异s，用来衡量当前管道单元和GPU 0上的第T个微批次完成其训练往返的管道单元之间的权重更新次数。正向计算时：
  
   反向传播时：
  
   其中T代表microbatch的个数，rank代表的是当前GPU的index.
   (2)计算weight prediction
  
   图中箭头所指位置即为训练第二个minibatch所需要计算weight prediction的位置，反向传播也是如此。当一个minibatch的bellwether完成权重预测之后，整个minibatch中的所有microbatch将共享这一份权重！！
   以T = 4为例，整个minibatch分为4个microbatch，当GPU0开始处理第二个minibatch的时候，首先计算weight version difference，再利用上面的公式计算weight prediction并保存，然后完成microbatch5的正向计算之后，向下游设备传递结果，GPU1再计算一次weight version difference，再计算weight prediction保存，依次向后完成正向计算和反向传播；之后，下一个microbatch继续利用当前的weight prediction进行正向计算和反向传播。也就是说，同一个minibatch是共享权重的！
  3、Conclusion
   受同步管道训练GPipe和异步管道训练PipeDream、SpecTrain等启发，提出了一种既能够把保证精度又能提高pipeline吞吐量的管道训练方法XPipe.最后的实验结果也表明XPipe在精度上比现有的所有异步管道训练方法都高，并且吞吐量也比同步管道训练方法要大，说明了该方法的有效性！