TF-Replicator: Distributed Machine Learning for Researchers

By P. Buchlovsky, David M. Budden, Dominik Grewe, Chris Jones, John Aslanides, Frederic Besse, Andy Brock, Aidan Clark, Sergio Gomez Colmenarejo, Aedan Pope and 2 more. 2019

20 citations recorded by the source corpus.

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 this https URL).

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, ML Systems & Hardware

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Sergey Levine
2. Pieter Abbeel
3. Chelsea Finn
4. Ruslan Salakhutdinov
5. Yoshua Bengio
6. Trevor Darrell
7. Nicolas Manfred Otto Heess
8. Quoc V. Le

In Computer Vision and Pattern Recognition

1. Luc Van Gool
2. Trevor Darrell
3. Dacheng Tao
4. Andrew Zisserman
5. Jitendra Malik
6. Andrea Vedaldi
7. Xiaogang Wang
8. Ming-Hsuan Yang

In ML Systems & Hardware

1. Tal Ben-Nun
2. Dan Belov
3. Torsten Hoefler
4. Jost Tobias Springenberg
5. Siqi Liu
6. Martin A. Riedmiller
7. Sean Luke
8. Liviu Panait

Citations in the corpus

• Large Scale Adversarial Representation Learning (2019, methodology): “All models were trained via TensorFlow [1] and Sonnet [32] with data parallelism on TPU pod slices [14] using 32 to 512 cores, coordinated by TF-Replicator [3].”
• V-MPO: On-Policy Maximum a Posteriori Policy Optimization for Discrete and Continuous Control (2019, methodology): “We implemented V-MPO in an actor-learner framework (Espeholt et al., 2018) that utilizes TF-Replicator (Buchlovsky et al., 2019) for distributed training on TPU 8-core and 16-core configurations (Google, 2018).”
• Efficient Acceleration of Deep Learning Inference on Resource-Constrained Edge Devices: A Review (2023, unknown): “TF replicator is another framework for leveraging distributed ML across TPUs [301].”
• Regularized Hierarchical Policies for Compositional Transfer in Robotics (2019, unknown): “Across all tasks, we build on a distributed actor-critic framework (similar to [Espeholt et al., 2018]) with flexible hardware assignment [Buchlovsky et al., 2019] to train all agents, performing critic and policy updates from a replay buffer, which is asynchronously filled by a set of actors.”
• Data Movement Is All You Need: A Case Study on Optimizing Transformers (2020, background): “…model parallelism (Van Es-sen et al., 2015; Jia et al., 2019d; Gholami et al., 2018; Dean et al., 2012; Chilimbi et al., 2014; Shazeer et al., 2018; Buchlovsky et al., 2019), pipeline parallelism (Chen et al., 2012; Li et al., 2018; Narayanan et al., 2019; Huang et al., 2019), microbatching (Oyama…”