Supporting Very Large Models using Automatic Dataflow Graph Partitioning

TL;DR

Tofu is a system that automatically partitions large deep neural network models across multiple GPUs by analyzing dataflow graphs, reducing memory usage and speeding up training.

Contribution

It introduces a novel operator semantics language and a recursive search algorithm for optimal partitioning to support large models on multiple GPUs.

Findings

Enables training of very large CNN and RNN models.

Achieves 25% to 400% speedup over alternative methods.

Reduces per-GPU memory footprint significantly.

Abstract

This paper presents Tofu, a system that partitions very large DNN models across multiple GPU devices to reduce per-GPU memory footprint. Tofu is designed to partition a dataflow graph of fine-grained tensor operators in order to work transparently with a general-purpose deep learning platform like MXNet. In order to automatically partition each operator, we propose to describe the semantics of an operator in a simple language which represents tensors as lambda functions mapping from tensor coordinates to values. To optimally partition different operators in a dataflow graph, Tofu uses a recursive search algorithm that minimizes the total communication cost. Our experiments on an 8-GPU machine show that Tofu enables the training of very large CNN and RNN models. It also achieves 25% - 400% speedup over alternative approaches to train very large models.