[论文解读] HyPar-Flow: Exploiting MPI and Keras for Scalable Hybrid-Parallel DNN Training using TensorFlow

Ammar Ahmad Awan, Arpan Jain+3 more|arXiv (Cornell University)|Nov 12, 2019

Advanced Neural Network Applications参考文献 26被引用 5

一句话总结

HyPar-Flow 是一个用户透明的可扩展框架，可使用 TensorFlow、MPI 和 Eager Execution 对任何基于 Keras 的 DNN 实现混合并行训练。它在 512 个 Frontera 节点上对 ResNet-1001 的训练实现了高达 481 倍的加速，通过动态协调模型并行与数据并行，并结合高效的基于 MPI 的通信和流水线技术。

ABSTRACT

To reduce training time of large-scale DNNs, scientists have started to explore parallelization strategies like data-parallelism, model-parallelism, and hybrid-parallelism. While data-parallelism has been extensively studied and developed, several problems exist in realizing model-parallelism and hybrid-parallelism efficiently. Four major problems we focus on are: 1) defining a notion of a distributed model across processes, 2) implementing forward/back-propagation across process boundaries that requires explicit communication, 3) obtaining parallel speedup on an inherently sequential task, and 4) achieving scalability without losing out on a model's accuracy. To address these problems, we create HyPar-Flow --- a model-size/-type agnostic, scalable, practical, and user-transparent system for hybrid-parallel training by exploiting MPI, Keras, and TensorFlow. HyPar-Flow provides a single API that can be used to perform data, model, and hybrid parallel training of any Keras model at scale. We create an internal distributed representation of the user-provided Keras model, utilize TF's Eager execution features for distributed forward/back-propagation across processes, exploit pipelining to improve performance and leverage efficient MPI primitives for scalable communication. Between model partitions, we use send and recv to exchange layer-data/partial-errors while allreduce is used to accumulate/average gradients across model replicas. Beyond the design and implementation of HyPar-Flow, we also provide comprehensive correctness and performance results on three state-of-the-art HPC systems including TACC Frontera (#5 on Top500.org). For ResNet-1001, an ultra-deep model, HyPar-Flow provides: 1) Up to 1.6x speedup over Horovod-based data-parallel training, 2) 110x speedup over single-node on 128 Stampede2 nodes, and 3) 481x speedup over single-node on 512 Frontera nodes.

研究动机与目标

• 通过支持模型并行和混合并行训练，解决超深层和大规模 DNN 因超出 GPU 内存限制而难以训练的挑战。
• 克服在多个进程之间显式通信以实现分布式前向/反向传播的复杂性。
• 提供统一的、用户透明的 API，无需修改模型定义或手动划分模型，可无缝切换数据并行、模型并行和混合并行。
• 在不牺牲模型准确率的前提下实现高可扩展性和高性能，即使对于 ResNet-1001 和实验性 5,000 层 ResNet-5000 等模型亦如此。
• 通过抽象用户层面的底层通信和编排任务，实现在 HPC 系统上对内存外模型的训练。

提出的方法

• 暴露一个单一的高层 API，接收 Keras 模型、模型分区数、副本数和训练策略（数据并行、模型并行或混合并行），且无需修改模型。
• 通过在进程间划分 Keras 模型并在节点间复制分区以实现数据并行，内部构建分布式模型表示。
• 利用 TensorFlow 的 Eager Execution，实现跨进程边界的动态、分布式前向和反向传播，支持显式通信。
• 使用 MPI 原语——通过 send/recv 在分区之间交换层数据和部分梯度，通过 allreduce 在副本之间聚合梯度。
• 通过优化的流水线阶段大小（128）对批处理分区实施流水线处理，以隐藏通信延迟并提高吞吐量。
• 在 Trainer 和 Communication Engine 中抽象系统级复杂性（分区、放置、通信），确保用户透明。

实验结果

研究问题

• RQ1 统一的、用户透明的框架能否高效支持在大规模 HPC 系统上对任意 Keras 模型实现混合并行训练？
• RQ2 如何有效结合模型并行与数据并行，以在不损害训练准确率的前提下实现高可扩展性？
• RQ3 在深度学习训练工作负载中，使用 MPI 原语进行分区间和副本间通信的性能影响如何？
• RQ4 HyPar-Flow 在多大程度上能够将超深层模型（如 ResNet-1001 和 ResNet-5000）的训练扩展至单 GPU 系统内存限制之外？
• RQ5 在现代 HPC 平台上，HyPar-Flow 的混合并行方法与传统数据并行方法相比，在加速比和效率方面表现如何？

主要发现

• 在大规模系统上，HyPar-Flow 对 ResNet-1001 的训练相比基于 Horovod 的数据并行训练最高实现了 1.6 倍的加速。
• 在 128 个 Stampede2 节点上，HyPar-Flow 对 ResNet-1001 实现了相比单节点训练 110 倍的加速，展示了出色的弱可扩展性。
• 在 512 个 Frontera 节点（28,762 个核心）上，HyPar-Flow 对 ResNet-1001 实现了相比单节点训练 481 倍的加速，凸显了其在混合并行训练中卓越的强可扩展性。
• 使用 HyPar-Flow 进行模型并行训练，对 ResNet-1001 相比顺序训练最高实现 3.2 倍加速，相比数据并行训练实现 1.6 倍加速。
• 该框架成功在仅两个节点上训练了实验性 5,000 层 ResNet-5000 模型，证明其能够处理现有系统无法支持的模型。
• 128 个阶段的流水线被证明为最优，显著通过重叠计算与通信提升了性能。