Theano-MPI: a Theano-based Distributed Training Framework

TL;DR

Theano-MPI is a scalable, GPU-based distributed training framework that accelerates deep learning model training across clusters, supporting both synchronous and asynchronous methods with reduced communication overhead.

Contribution

It introduces a flexible, open-source framework utilizing CUDA-aware MPI for efficient distributed deep learning training with novel communication optimization techniques.

Findings

Effective scaling of AlexNet and GoogLeNet from 2 to 8 GPUs

Demonstrated reduction in communication overhead

Achieved faster training times with the framework

Abstract

We develop a scalable and extendable training framework that can utilize GPUs across nodes in a cluster and accelerate the training of deep learning models based on data parallelism. Both synchronous and asynchronous training are implemented in our framework, where parameter exchange among GPUs is based on CUDA-aware MPI. In this report, we analyze the convergence and capability of the framework to reduce training time when scaling the synchronous training of AlexNet and GoogLeNet from 2 GPUs to 8 GPUs. In addition, we explore novel ways to reduce the communication overhead caused by exchanging parameters. Finally, we release the framework as open-source for further research on distributed deep learning