[Paper Review] 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 Techniques19 references18 citations

TL;DR

This paper proposes MXNET-MPI, a hybrid framework that integrates MPI parallelism into the parameter server (PS) task model to enable scalable deep learning training. By embedding MPI collective operations—particularly tensor collectives—within MXNet's computation graph, the framework reduces network contention and parameter staleness, achieving 6x faster training per epoch and over 0.72 ImageNet top-1 accuracy with ResNet-50, outperforming both pure PS and MPI approaches.

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.

Motivation & Objective

• To address the scalability limitations of pure parameter server (PS) and MPI-based deep learning frameworks, which suffer from network hotspots, parameter staleness, or poor fault tolerance.
• To enable seamless coexistence of MPI and PS paradigms within a single framework, allowing flexible tuning between PS elasticity and MPI performance.
• To design and implement novel communication-avoiding algorithms, such as MPI Elastic SGD, that improve convergence rates in distributed SGD.
• To optimize global aggregation using tensor collectives, treating groups of vectors per GPU as single objects to improve bandwidth utilization and reduce latency.
• To demonstrate state-of-the-art performance on large-scale ImageNet 1K training using real-world HPC and cloud clusters.

Proposed method

• Embed MPI collective primitives—specifically allreduce and broadcast—into MXNet’s computation graph by creating an independent MPI_COMM_WORLD job client that interfaces with the PS.
• Introduce a hybrid model where the number of MPI clients is tunable, enabling a smooth transition from pure PS to pure MPI, with intermediate configurations balancing fault tolerance and performance.
• Design a novel MPI Elastic SGD algorithm that combines synchronous updates within an MPI communicator with asynchronous lazy updates outside, reducing staleness and improving convergence.
• Implement tensor collectives that treat a GPU’s set of parameter vectors as a single tensor object, allowing existing single-vector algorithms to be reused efficiently.
• Optimize global reduction using GPU-accelerated kernels (IBMGpu) and compare with NCCL, using ring-based and host-memory-based reduction strategies to identify the highest-bandwidth approach.
• Leverage high-bandwidth interconnects (e.g., NVLink) and multi-threaded CUDA kernels (112 thread blocks) to maximize reduction and broadcast bandwidth, achieving up to 30 GB/sec for allreduce.

Experimental results

Research questions

• RQ1 Can MPI and parameter server models be co-located within a single deep learning framework to combine their respective strengths in performance and fault tolerance?
• RQ2 How can MPI collective operations be efficiently embedded into a task-based PS framework like MXNet without breaking its programming model?
• RQ3 Can novel communication-avoiding algorithms, such as Elastic SGD with MPI, significantly improve convergence rates in distributed SGD at scale?
• RQ4 To what extent do tensor collectives—treating multiple vectors per GPU as a single tensor—improve bandwidth and reduce latency in large-scale allreduce operations?
• RQ5 Can the hybrid MPI-PS model achieve state-of-the-art training performance on ImageNet 1K while maintaining scalability and fault tolerance?

Key findings

• The hybrid MXNET-MPI framework reduces training time per epoch by six times compared to the default PS approach on ImageNet 1K, while improving convergence rates.
• Using MPI Elastic SGD, the framework achieves more than a 2x improvement in convergence rate compared to all major parallel SGD approaches, even with reduced communication frequency.
• The tensor collective-based allreduce implementation achieves a peak bandwidth of 30 GB/sec using IBMGpu kernels, outperforming NCCL (15 GB/sec) and Baidu’s ring implementation by a factor of six on the Minsky cluster.
• The framework achieves a validation accuracy of over 0.72 on the full ImageNet 1K dataset using ResNet-50, representing the state-of-the-art at the time of publication.
• The optimized MPI-ESGD configuration with only two clients and scaled workers outperforms dist-ESGD with 12 workers, demonstrating that reducing client count mitigates parameter staleness and improves convergence.
• The framework’s design allows seamless integration of MPI into Python-based deep learning workflows, abstracting low-level MPI complexity while enabling high performance and fault tolerance via LSF job restarts.