[Paper Review] User-transparent Distributed TensorFlow

Abhinav Vishnu, Joseph Manzano+2 more|arXiv (Cornell University)|Apr 15, 2017

Scientific Computing and Data Management6 references3 citations

TL;DR

This paper presents MaTEx-TensorFlow, a user-transparent distributed TensorFlow implementation that enables scale-out training on HPC systems without modifying user code. By integrating MPI directly into the TensorFlow runtime and abstracting low-level communication, it achieves near-ideal scaling on multi-node CPU and GPU systems using standard AlexNet, GoogLeNet, InceptionV3, and ResNet-50 models, with identical loss curves to sequential training.

ABSTRACT

Deep Learning (DL) algorithms have become the {\em de facto} choice for data analysis. Several DL implementations -- primarily limited to a single compute node -- such as Caffe, TensorFlow, Theano and Torch have become readily available. Distributed DL implementations capable of execution on large scale systems are becoming important to address the computational needs of large data produced by scientific simulations and experiments. Yet, the adoption of distributed DL implementations faces significant impediments: 1) most implementations require DL analysts to modify their code significantly -- which is a show-stopper, 2) several distributed DL implementations are geared towards cloud computing systems -- which is inadequate for execution on massively parallel systems such as supercomputers. This work addresses each of these problems. We provide a distributed memory DL implementation by incorporating required changes in the TensorFlow runtime itself. This dramatically reduces the entry barrier for using a distributed TensorFlow implementation. We use Message Passing Interface (MPI) -- which provides performance portability, especially since MPI specific changes are abstracted from users. Lastly -- and arguably most importantly -- we make our implementation available for broader use, under the umbrella of Machine Learning Toolkit for Extreme Scale (MaTEx) at { exttt http://hpc.pnl.gov/matex}. We refer to our implementation as MaTEx-TensorFlow.

Motivation & Objective

• Address the high barrier to entry for using distributed deep learning by eliminating the need for users to modify their TensorFlow code.
• Enable efficient execution of distributed TensorFlow on HPC systems with high-bandwidth interconnects like InfiniBand, which are unsuitable for gRPC-based solutions.
• Provide a production-ready, performant, and portable distributed TensorFlow runtime that supports data parallelism across multiple nodes.
• Ensure feature parity with standard TensorFlow by maintaining identical training dynamics and convergence behavior.
• Make the implementation broadly available through the Machine Learning Toolkit for Extreme Scale (MaTEx) to accelerate HPC-based deep learning adoption.

Proposed method

• Integrate MPI directly into the TensorFlow runtime to replace gRPC-based communication, enabling native support for high-performance interconnects like InfiniBand.
• Abstract MPI-specific logic from users by embedding it within the runtime, so users do not need to write MPI calls or modify their existing TensorFlow scripts.
• Implement data parallelism with synchronous stochastic gradient descent (SGD) by synchronizing gradients across replicas using MPI collective operations.
• Extend TensorFlow 1.0 with new user-operations and communication primitives to support distributed training while preserving backward compatibility.
• Use a custom data reader interface that supports multiple data formats, simplifying data loading in distributed environments.
• Leverage existing TensorFlow computational graphs and automatic differentiation, ensuring compatibility with all standard models (e.g., AlexNet, ResNet-50).

Experimental results

Research questions

• RQ1 Can a distributed TensorFlow implementation be built that requires no code changes to existing user scripts?
• RQ2 Can MPI be effectively integrated into the TensorFlow runtime to enable high-performance, scalable training on HPC systems with InfiniBand interconnects?
• RQ3 Does the proposed implementation maintain identical convergence behavior and loss curves compared to sequential TensorFlow training?
• RQ4 Can the system achieve strong scaling on both multi-core CPU and multi-GPU HPC platforms without sacrificing performance portability?
• RQ5 Is it feasible to provide a production-ready, user-transparent distributed deep learning system under an open-source HPC software umbrella like MaTEx?

Key findings

• MaTEx-TensorFlow achieves feature parity with sequential TensorFlow, as evidenced by identical loss curves during training on the ImageNet LSVRC12 dataset using AlexNet.
• The implementation scales effectively across multiple compute nodes on both Intel multi-core systems and NVIDIA multi-GPU systems using InfiniBand interconnects.
• The system enables transparent distributed training without requiring users to modify their TensorFlow code, significantly lowering the barrier to entry for HPC-based deep learning.
• By embedding MPI within the TensorFlow runtime, the approach provides performance portability and efficient use of high-bandwidth interconnects common in supercomputing environments.
• The solution is production-ready and publicly available under the MaTEx umbrella at http://hpc.pnl.gov/matex, supporting a wide range of deep learning models including GoogLeNet, InceptionV3, and ResNet-50.
• The approach successfully addresses the limitations of gRPC-based distributed TensorFlow, which is ill-suited for HPC systems due to reliance on Ethernet and lack of native RDMA support.