[论文解读] MXNET-MPI: Embedding MPI parallelism in Parameter Server Task Model for scaling Deep Learning

Amith R. Mamidala, Γεώργιος Κόλλιας+2 more|arXiv (Cornell University)|Jan 11, 2018

Stochastic Gradient Optimization Techniques参考文献 19被引用 18

一句话总结

本文提出 MXNET-MPI，一种混合框架，将 MPI 并行机制集成到参数服务器（PS）任务模型中，以实现可扩展的深度学习训练。通过将 MPI 集体操作——尤其是张量集合操作——嵌入 MXNet 的计算图中，该框架减少了网络争用和参数陈旧性，实现了每轮训练速度提升 6 倍，并在 ResNet-50 上实现了超过 0.72 的 ImageNet 顶部-1 准确率，优于纯 PS 和纯 MPI 方法。

ABSTRACT

Existing Deep Learning frameworks exclusively use either Parameter Server(PS) approach or MPI parallelism. In this paper, we discuss the drawbacks of such approaches and propose a generic framework supporting both PS and MPI programming paradigms, co-existing at the same time. The key advantage of the new model is to embed the scaling benefits of MPI parallelism into the loosely coupled PS task model. Apart from providing a practical usage model of MPI in cloud, such framework allows for novel communication avoiding algorithms that do parameter averaging in Stochastic Gradient Descent(SGD) approaches. We show how MPI and PS models can synergestically apply algorithms such as Elastic SGD to improve the rate of convergence against existing approaches. These new algorithms directly help scaling SGD clusterwide. Further, we also optimize the critical component of the framework, namely global aggregation or allreduce using a novel concept of tensor collectives. These treat a group of vectors on a node as a single object allowing for the existing single vector algorithms to be directly applicable. We back our claims with sufficient emperical evidence using large scale ImageNet 1K data. Our framework is built upon MXNET but the design is generic and can be adapted to other popular DL infrastructures.

研究动机与目标

• 为解决纯参数服务器（PS）和基于 MPI 的深度学习框架在可扩展性方面的局限性，这些框架常面临网络热点、参数陈旧或容错能力差的问题。
• 实现在单一框架内 MPI 与 PS 范式的无缝共存，支持在 PS 弹性与 MPI 性能之间灵活调优。
• 设计并实现新型通信避免算法，如 MPI Elastic SGD，以提升分布式 SGD 中的收敛速度。
• 通过将每 GPU 的向量组视为单一对象，优化全局聚合操作，提升带宽利用率并降低延迟。
• 在真实 HPC 与云集群上，展示在大规模 ImageNet 1K 训练中达到最先进性能。

提出的方法

• 通过创建独立的 MPI_COMM_WORLD 作业客户端与参数服务器接口，将 MPI 集体原语——特别是 allreduce 和 broadcast——嵌入 MXNet 的计算图中。
• 引入一种混合模型，其中 MPI 客户端数量可调，实现从纯 PS 到纯 MPI 的平滑过渡，中间配置兼顾容错性与性能。
• 设计一种新型 MPI Elastic SGD 算法，结合 MPI 通信器内的同步更新与外部的异步懒更新，降低陈旧性并提升收敛性。
• 实现张量集合操作，将 GPU 上的参数向量集合视为单一张量对象，使现有单向量算法可高效复用。
• 使用 GPU 加速内核（IBMGpu）优化全局规约操作，并与 NCCL 对比，采用基于环形和主机内存的规约策略，识别出最高带宽方案。
• 利用高带宽互连（如 NVLink）和多线程 CUDA 内核（112 个线程块）最大化规约与广播带宽，实现 allreduce 最高 30 GB/sec 的性能。

实验结果

研究问题

• RQ1 MPI 与参数服务器模型能否在单一深度学习框架中协同部署，以结合其在性能与容错性方面的优势？
• RQ2 如何高效地将 MPI 集体操作嵌入类似 MXNet 的基于任务的 PS 框架中，而不破坏其编程模型？
• RQ3 新型通信避免算法（如结合 MPI 的 Elastic SGD）是否能在大规模分布式 SGD 中显著提升收敛速度？
• RQ4 将每 GPU 的多个向量视为单一张量的张量集合操作，在大规模 allreduce 操作中能在多大程度上提升带宽并降低延迟？
• RQ5 混合 MPI-PS 模型能否在 ImageNet 1K 上实现最先进训练性能，同时保持可扩展性与容错性？

主要发现

• 与默认的 PS 方法相比，混合 MXNET-MPI 框架在 ImageNet 1K 上将每轮训练时间减少了六倍，同时提升了收敛速度。
• 采用 MPI Elastic SGD，该框架在通信频率降低的情况下，相比所有主流并行 SGD 方法，收敛速度提升超过 2 倍。
• 基于张量集合的 allreduce 实现使用 IBMGpu 内核达到峰值带宽 30 GB/sec，显著优于 NCCL（15 GB/sec）和百度环形实现六倍，在 Minsky 集群上表现优异。
• 在使用 ResNet-50 的完整 ImageNet 1K 数据集上，该框架实现了超过 0.72 的验证准确率，为发表时的最先进水平。
• 仅使用两个客户端并扩展工作节点的优化 MPI-ESGD 配置，优于使用 12 个工作节点的 dist-ESGD，表明减少客户端数量可缓解参数陈旧性并提升收敛性。
• 该框架的设计支持将 MPI 无缝集成到基于 Python 的深度学习工作流中，抽象了底层 MPI 复杂性，同时通过 LSF 作业重启实现高性能与容错性。