TF-Replicator: Distributed Machine Learning for Researchers

TL;DR

TF-Replicator is a flexible, scalable framework built on TensorFlow that simplifies distributed machine learning research across various architectures and models, enabling researchers to deploy and benchmark complex models efficiently.

Contribution

It introduces TF-Replicator, a new abstraction over TensorFlow that simplifies writing and deploying distributed machine learning models across different hardware setups.

Findings

Achieves strong scalability performance

Supports diverse models including ResNet-50, SN-GAN, and D4PG

Requires minimal distributed systems expertise

Abstract

We describe TF-Replicator, a framework for distributed machine learning designed for DeepMind researchers and implemented as an abstraction over TensorFlow. TF-Replicator simplifies writing data-parallel and model-parallel research code. The same models can be effortlessly deployed to different cluster architectures (i.e. one or many machines containing CPUs, GPUs or TPU accelerators) using synchronous or asynchronous training regimes. To demonstrate the generality and scalability of TF-Replicator, we implement and benchmark three very different models: (1) A ResNet-50 for ImageNet classification, (2) a SN-GAN for class-conditional ImageNet image generation, and (3) a D4PG reinforcement learning agent for continuous control. Our results show strong scalability performance without demanding any distributed systems expertise of the user. The TF-Replicator programming model will be open-sourced as part of TensorFlow 2.0 (see https://github.com/tensorflow/community/pull/25).