Pipeline Parallelism GPipe And PipeDream

从模型并行讲起

模型并行有两种并行方式,一种是层间并行(inter-layer),一种是层内并行(intra-layer).

流水线并行属于层间并行的一种特殊情况,下文在讲解GPipe的时候会进行提及

层内并行可以看作是张量并行,张量并行的经典文章Megatron-LM是对transformer的三个维度之一[batch_size,sequence_length,hidden_dimmension]中的hidden_dimension进行划分,从而并行

后续也有一些文章是对sequence_length进行划分来达到并行的目的

不难看出,传统的数据并行是对数据集进行划分,即比如数据集有3200张图片,一个batch_size取32,则一个epoch有100个iteration,

GPipe

micro-batch

re-materialization

可挖掘的创新点

1. 在文章Performance Optimization这一部分提及了可以更早安排反向传递时的recompute,我们知道,要计算$B_k$层的梯度,需要$B_{k+1}$层的梯度以及$F_k$层的激活值,因此如果可以较为精准的提前开始重新计算$F_k$层的激活值,当$B_{k+1}$的值通过PCI-E也好,通过Infiniband,甚至用NVLink也好,要是可以适应各类设备,而遮掩掉recompute的时间的话(大约$\frac{1}{3}$),那跟别的模型比起来,等于白赚了recompute而节省的那一堆内存
2. 在文章Design Features and Trade-Offs这一部分中提及了GPipe假设的是单个层符合单个加速器的内存需求,因此**如果单个加速器无法满足单个层的需求,应该如何处理?**文中提出的是可以通过把矩阵乘拆分来解决
3. 不同层内存需求不平衡,因此需要更好的分区算法,以达到负载均衡,如果实在无法通过分区达到平衡,可以通过合适的卸载计算,将计算卸载到周边的计算节点处理,感觉这里可以借鉴MEC(移动边缘计算)一类的方法
4. micro-batch划分需要复杂的策略以支持需要跨micro-batch(即在mini-batch中)计算的层(如BN层)

PipeDream

weight stashing

vertical sync