Dynamic control flow in large-scale machine learning

By Yuan Yu, Martín Abadi, Paul Barham, Eugene Brevdo, M. Burrows, Andy Davis, Joseph Dean, Sanjay Ghemawat, Tim Harley, Peter G. Hawkins and 5 more. 2018

112 citations recorded by the source corpus.

Many recent machine learning models rely on fine-grained dynamic control flow for training and inference. In particular, models based on recurrent neural networks and on reinforcement learning depend on recurrence relations, data-dependent conditional execution, and other features that call for dynamic control flow. These applications benefit from the ability to make rapid control-flow decisions across a set of computing devices in a distributed system. For performance, scalability, and expressiveness, a machine learning system must support dynamic control flow in distributed and heterogeneous environments. This paper presents a programming model for distributed machine learning that supports dynamic control flow. We describe the design of the programming model, and its implementation in TensorFlow, a distributed machine learning system. Our approach extends the use of dataflow graphs to represent machine learning models, offering several distinctive features. First, the branches of conditionals and bodies of loops can be partitioned across many machines to run on a set of heterogeneous devices, including CPUs, GPUs, and custom ASICs. Second, programs written in our model support automatic differentiation and distributed gradient computations, which are necessary for training machine learning models that use control flow. Third, our choice of non-strict semantics enables multiple loop iterations to execute in parallel across machines, and to overlap compute and I/O operations. We have done our work in the context of TensorFlow, and it has been used extensively in research and production. We evaluate it using several real-world applications, and demonstrate its performance and scalability.

Subfields: Artificial Intelligence, Hardware and Architecture, Optimization & Training

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

In Hardware and Architecture

In Optimization & Training

Citations in the corpus

GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (2021, methodology): “As a result, we can wrap the repetitive modules of the MoE Transformer architecture in a while loop control flow statement (Abadi et al., 2016a; Yu et al., 2018) to reduce compilation time.”
Pathways: Asynchronous Distributed Dataflow for ML (2022, methodology): “Single-controller systems such as TensorFlow v1 (Abadi et al., 2016) offer a very general distributed dataﬂow model, including optimized in-graph control ﬂow (Yu et al., 2018).”
SwapAdvisor: Pushing Deep Learning Beyond the GPU Memory Limit via Smart Swapping (2020, methodology): “[57] is an on-demand swapping mechanism for TensorFlow.”
DistIR: An Intermediate Representation for Optimizing Distributed Neural Networks (2021, methodology): “Graph-based frameworks such as XLA [20] and ONNX Runtime [22] lower DNN models from frontends such as Tensorflow [3] or PyTorch [26] into IRs that are capable of expressing distributed computation [10, 14, 21, 35].”
The Deep Learning Compiler: A Comprehensive Survey (2020, methodology): “Models such as RNN and Reinforcement learning (RL) depend on recurrent relations and data-dependent conditional execution [42], which requires control flow.”
3LC: Lightweight and Effective Traffic Compression for Distributed Machine Learning (2018, methodology): “Our evaluation compares representative communication reduction schemes that we implement on TensorFlow: 32-bit float is the baseline that transmits 32-bit ﬂoating-point state changes without compression.”