Automated code transformation for distributed training of TensorFlow deep learning models

Abstract

Distributed training of deep learning models reduces training time by parallelizing training workloads across multiple GPUs. Distributed training frameworks, such as Horovod and DeepSpeed, provide APIs, and model engineers rewrite deep learning models using the APIs to parallelize their training. However, the rewriting is time-consuming and labor-intensive because it requires engineers to read and understand documents and examples of the frameworks as well as manual efforts to rewrite code.

In this paper, we propose an automated code transformation approach that transforms TensorFlow deep learning models designed for non-distributed training to models training on multiple GPUs with the Horovod framework. We closely inspect the Horovod document and code examples and identify four common training patterns of TensorFlow deep learning models. Then, we formalize code transformation rules for each training pattern. Using the rules, we implement an automated code transformation tool that takes a TensorFlow deep learning model written in Python and rewrites it with the Horovod APIs for distributed training. Through source-code level transformation, our approach enables developers to efficiently scale existing DL models to multiple GPUs. Our evaluation shows that the tool correctly transforms 15 out of 16 open-source TensorFlow deep learning models. To the best of our knowledge, our work is the first automatic transformation technique for distributing existing TensorFlow deep learning models at the source code level. We believe that our approach significantly reduces manual efforts to parallelize training of existing TensorFlow deep learning models.