AI顶会ICLR 2022 | WPipe 蚂蚁集团大规模 DNN 训练的流水线并行技术 - OSCHINA

导言

ICLR，全称为 International Conference on Learning Representations (国际学习表征会议) 是三大机器学习领域顶会之一 (另外两个是ICML和NeuriPS)。该会议的主要创办者就包含了深度学习三大巨头的YoShua Bengio和Yann LeCun。ICLR的第一届是2013年举行，短短几年就得到的快速发展，逐步成为机器学习领域的顶会。我们这项工作就是ICLR 2022收录论文。

近年来，使用大规模深度神经网络 (DNN) 来提高性能的趋势，推动了用于高效 DNN 训练的并行流水线技术的发展，也催生出了一系列优秀的的Pipeline Parallelism(PP)技术，例如 GPipe、PipeDream 和 PipeDream- 2BW。然而，目前领先的 PipeDream-2BW 技术仍然存在两个主要缺点，即过多的内存冗余和权重更新延迟。在这个方向上，我们提出了一种名为 WPipe 的新的PP技术，它可以实现更好的内存效率和更新鲜的更新梯度。

WPipe 使用一种新颖的流水线方案，它将模型分为两组，并通过巧妙的移动操作，使得两个分组既可以无缝更新，也可以消除一半的更新延迟和额外内存占用。通过大型语言模型Bert和视觉模型ResNeXt的训练实验表明，与当前SOTA技术PipeDream-2BW 相比，WPipe 实现了 1.4x 的加速且内存占用减少了 36%，同时保障最终的模型准确度不受影响。