[论文解读] BaPipe: Exploration of Balanced Pipeline Parallelism for DNN Training

Letian Zhao, Rui Xu+6 more|arXiv (Cornell University)|Dec 23, 2020

Network Security and Intrusion Detection参考文献 34被引用 8

一句话总结

BaPipe 是一种新颖的分布式深度神经网络（DNN）训练框架，能够自动探索跨加速器的层内批处理流水线并行性，并实现负载均衡的模型划分。它采用自适应调度（例如 1F1B-SNO）和资源感知的负载均衡，将内存使用量减少高达 4 倍，并在 GPU 和 FPGA 集群上实现相较于最先进框架高达 3.2 倍的加速。

ABSTRACT

The size of deep neural networks (DNNs) grows rapidly as the complexity of the machine learning algorithm increases. To satisfy the requirement of computation and memory of DNN training, distributed deep learning based on model parallelism has been widely recognized. We propose a new pipeline parallelism training framework, BaPipe, which can automatically explore pipeline parallelism training methods and balanced partition strategies for DNN distributed training. In BaPipe, each accelerator calculates the forward propagation and backward propagation of different parts of networks to implement the intra-batch pipeline parallelism strategy. BaPipe uses a new load balancing automatic exploration strategy that considers the parameters of DNN models and the computation, memory, and communication resources of accelerator clusters. We have trained different DNNs such as VGG-16, ResNet-50, and GNMT on GPU clusters and simulated the performance of different FPGA clusters. Compared with state-of-the-art data parallelism and pipeline parallelism frameworks, BaPipe provides up to 3.2x speedup and 4x memory reduction in various platforms.

研究动机与目标

• 解决因内存瓶颈和高通信开销导致的数据并行化在训练大型 DNN 时的局限性。
• 通过支持具有统一权重版本的层内批处理调度，克服现有流水线并行化框架的低效问题。
• 自动探索最优的流水线调度与模型划分策略，以在异构加速器集群中平衡计算、内存和通信负载。
• 通过减少内存占用，实现更大 DNN 模型的训练，同时保持高训练吞吐量和收敛稳定性。
• 支持多种硬件平台，包括 GPU 和 FPGA 集群，具备自适应负载均衡和通信-计算重叠能力。

提出的方法

• 提出一种新型的层内批处理流水线并行策略，即每个加速器通过微批次以流水线方式处理网络的不同部分。
• 引入一种自动负载均衡探索策略，通过分析模型参数和硬件资源（计算、内存、通信）来选择最优的划分与调度方案。
• 采用多种流水线调度技术，包括 1F1B-SNO（一次前向传播、一次反向传播且无重叠），以平衡通信与计算，减少流水线空洞。
• 使用动态调度引擎，根据硬件约束和模型特性选择最佳配置（例如微批次大小、划分方式、调度模式）。
• 通过尽可能将模型权重和中间激活值存储在高带宽片上内存中，优化内存使用，尤其在 FPGA 集群中效果显著。
• 通过适配平台特定的内存层次结构和通信后端（如 GPU 的 GLOO，FPGA 的仿真），同时支持 GPU 和 FPGA 集群。

实验结果

研究问题

• RQ1 如何设计流水线并行化，以在保证训练一致性的同时实现在 DNN 训练中的高吞吐量？
• RQ2 在层内批处理流水线并行化中，哪些调度策略能最小化流水线空洞并最大化加速器利用率？
• RQ3 如何在考虑计算、内存和通信约束的前提下，自动实现异构加速器之间的负载均衡？
• RQ4 在不牺牲训练速度或模型准确率的前提下，DNN 训练中的内存使用量最多可减少多少？
• RQ5 所提出的框架在不同硬件平台（GPU 和 FPGA）上的可扩展性如何，体现在加速比和模型大小支持方面？

主要发现

• 在 GPU 集群上，BaPipe 相较于数据并行化（DP）实现高达 3.2 倍的加速，相较于最先进流水线并行化框架（GPipe）实现 1.4 倍的加速，适用于 VGG-16 和 ResNet-50 等小型 DNN。
• 在 GPU 集群上，由于通过负载均衡划分实现了有效的内存减少，BaPipe 支持的 DNN 模型规模比数据并行化大 4 倍，比 GPipe 大 2 倍。
• 在 FPGA 集群上，BaPipe 相较于 DP 实现高达 1.14 倍的加速，通过将权重存储在片上内存并最小化片外访问，内存使用量减少高达 4 倍。
• 该框架能自动选择最优的调度（如 1F1B-SNO）和划分策略，无需人工调优即可提升性能。
• 对于 ResNet-50，当通信开销超过计算收益时，BaPipe 在 GPU 集群上通过选择 DP 作为最优策略，优于 PipeDream。
• BaPipe 在 GPU 集群上支持的 GNMT-L 模型参数量最多可达 DP 的 4 倍，展示了其在超大 DNN 训练中的可扩展性。