Source-linked AI summary
TensorFlow: A system for large-scale machine learning
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, Derek G. Murray, Benoit Steiner, Paul Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yuan Yu, Xiaoqiang Zheng
TL;DR
Existing systems limited extensibility for experimenting with new optimization algorithms and model architectures. TensorFlow introduces a unified dataflow model for computation and mutable state across heterogeneous systems, and demonstrates performant, scalable implementations for production and research workloads.
Problem
Parameter-server systems were insufficiently extensible for experimenting with new optimization algorithms and unconventional model architectures, limiting accessible large-scale machine-learning development.
Method
TensorFlow represents computation and mutable state in a unified dataflow graph and maps it across heterogeneous machines and devices for training, inference, and experimentation.
Results
TensorFlow implementations were performant and scalable across demonstrated examples, with sparse embedding step times ranging from 5 to 20 ms.
Takeaways & Limitations
TensorFlow provides a uniform programming model for harnessing large-scale heterogeneous systems in production tasks and experimentation with new approaches.
Takeaways & Limitations
TensorFlow lacks default policies that work well for most users, leaving automatic optimization and stronger consistency policies as ongoing development areas.
Abstract
from arXiv · showhide
TensorFlow is a machine learning system that operates at large scale and in heterogeneous environments. TensorFlow uses dataflow graphs to represent computation, shared state, and the operations that mutate that state. It maps the nodes of a dataflow graph across many machines in a cluster, and within a machine across multiple computational devices, including multicore CPUs, general-purpose GPUs, and custom designed ASICs known as Tensor Processing Units (TPUs). This architecture gives flexibility to the application developer: whereas in previous "parameter server" designs the management of shared state is built into the system, TensorFlow enables developers to experiment with novel optimizations and training algorithms. TensorFlow supports a variety of applications, with particularly strong support for training and inference on deep neural networks. Several Google services use TensorFlow in production, we have released it as an open-source project, and it has become widely used for machine learning research. In this paper, we describe the TensorFlow dataflow model in contrast to existing systems, and demonstrate the compelling performance that TensorFlow achieves for several real-world applications.
1 Introduction
TensorFlow is introduced as a system for developing, training, and deploying machine-learning models at scale. Its unified dataflow graph represents both computation and mutable state, supporting broad experimentation and large-scale training and inference.
- Motivation: Machine learning’s recent advances are attributed to sophisticated models, large datasets, and software platforms that make substantial computational resources easy to use.
- Contribution: TensorFlow enables researchers to experiment with new models, train them on large datasets, and move them into production.The system simplifies and generalizes Google’s first-generation DistBelief system to support a wider variety of ideas.
- Contribution: TensorFlow supports both large-scale training and inference, efficiently using hundreds of powerful GPU-enabled servers for fast training.
- Novelty: TensorFlow uses a unified dataflow graph to represent algorithmic computation and the state on which the algorithm operates.Unlike traditional dataflow systems, graph vertices can represent computations that own or update mutable state.
- Adoption and scope: More than 60 teams at Google used TensorFlow within a year, and the system was released as an open-source project.The paper focuses on neural network training, using image classification and language modeling as representative applications.
2 Background & Motivation
Large-scale machine learning requires distributed execution over large datasets and models, efficient accelerator use, scalable training and inference, and an extensible programming model. Existing single-machine, batch-dataflow, and parameter-server systems address parts of these needs but leave gaps in distributed execution, mutable model updates, or extensibility.
- Requirements: Distributed execution enables machine learning systems to use more data and larger models efficiently.Large datasets motivate data-parallel training, while large models motivate sharding parameters across machines.
- Requirements: Accelerator support is important because matrix multiplication and multidimensional convolution are highly parallelizable but require tightly coupled implementations.A single NVIDIA Titan X provides 6 TFLOPS peak performance, while Google’s TPU achieves an order-of-magnitude improvement in performance-per-watt over alternative state-of-the-art technology.
- Requirements: Production systems require scalable, high-performance inference alongside training, including low-latency, mobile, and distributed execution scenarios.Developers benefit from using the same code to define models for both training and inference.
- Related work: Batch dataflow systems struggle with mutable machine-learning models because immutable inputs and deterministic subcomputations make model updates expensive.SparkNet, for example, takes 20 seconds to broadcast weights and collect updates when training deep neural networks on Spark.
- Related work: Parameter servers meet most requirements but are insufficiently extensible because new optimization algorithms or unconventional architectures require modifying a C++ implementation.Single-machine frameworks provide extensible programming models, while the desired system should also let users scale the same code into production.
3 TensorFlow execution model
TensorFlow represents computation and mutable state in a dataflow graph whose partial, concurrent execution enables flexible coordination and distributed computation across heterogeneous devices. Its model lets users experiment with optimization, consistency, and parallelization strategies without modifying the runtime.
- Graph representation: A single dataflow graph represents an algorithm’s computations, state, parameter updates, and input preprocessing, making communication explicit for parallel and distributed execution.The graph can be partitioned across multiple distributed devices.
- Distributed execution: Dataflow with mutable state can mimic parameter-server functionality while enabling arbitrary subgraphs on parameter-hosting machines and experimentation with optimization, consistency, and parallelization.This flexibility is presented as a consequence of combining mutable state with dataflow execution.
- Graph representation: Vertices represent atomic operations, edges carry tensors, and operations may consume and produce variable numbers of typed tensor values.TensorFlow models data as dense n-dimensional arrays such as int32, float32, or string tensors.
- Partial and concurrent execution: Mutable variables and queues coordinate concurrent computations, supporting shared model parameters, data-parallel training, and diverse model architectures in unprivileged code.This partial and concurrent execution lets advanced users experiment without modifying TensorFlow’s runtime internals.
- Partial and concurrent execution: Clients declaratively select feeds and fetches for each step, allowing the runtime to prune the graph and execute multiple concurrent steps on the same graph.Stateful operations allow concurrent executions to interact through shared state.
- Distributed execution: Dataflow makes distributed execution explicit, allowing the same program to target GPU clusters, TPU clusters, and mobile inference devices.The runtime places operations on devices and partitions each executed subgraph into per-device subgraphs connected across device boundaries.
4 Extensibility case studies
TensorFlow’s unified dataflow graph lets users build extensions with simple primitives and user-level code rather than modifying the runtime. The case studies cover automatic differentiation, advanced optimization, sparse embeddings and parameter computation, fault-tolerant checkpointing, and synchronization methods.
- Automatic differentiation: TensorFlow’s user-level differentiation library derives backpropagation automatically from neural-network layers and loss functions.The algorithm identifies backward paths and sums their partial gradients; users can specialize gradients and implement batch normalization or gradient clipping.
- Optimization algorithms: Advanced optimizers such as Momentum can be implemented through TensorFlow’s graph primitives, avoiding parameter-server representation changes.Momentum accumulates a velocity for each parameter across iterations, whereas simple parameter-server writes directly express SGD.
- Sparse embeddings and parameter computation: TensorFlow composes primitive operations to implement sharded sparse embedding layers and colocates Gather with the variable it reads.Sparse embedding matrices can contain gigabytes of parameters, and TensorFlow also permits arbitrary computation on devices hosting shared parameters.
- Fault tolerance: User-level checkpointing uses Save, Restore, and Assign graph operations, with customizable variable policies and checkpoint retention.The implementation is reusable for fine-tuning and unsupervised pre-training, but concurrent checkpointing may produce inconsistent checkpoints that are suitable for asynchronous gradient descent.
- Parameter synchronization: TensorFlow’s graph lets users choose asynchronous or synchronous parameter synchronization schemes, including synchronous training with backup workers.Asynchronous updates improve utilization but use stale information; Figure 4 presents the three alternatives.
5 Implementation
TensorFlow is implemented as an extensible, cross-platform library with a thin C API separating multilingual users from a C++ core. Its runtime distributes and optimizes graph execution across heterogeneous devices while supporting user-defined kernels and production-oriented tools.
- Core library: A thin C API separates user-level code in various languages from TensorFlow’s extensible core library, which is implemented in C++.The implementation is open-source and supports Linux, Mac OS X, Android, and iOS, along with x86, ARM-based CPUs, and several NVIDIA GPU microarchitectures.
- Graph execution: The distributed master prunes, partitions, caches, and optimizes graph subgraphs before execution across participating devices.Its optimizations include common subexpression elimination and constant folding; pruning acts as dead code elimination.
- Graph execution: Approximately 2,000,000 null operations per second are dispatched by the dataflow executor, which schedules local kernels with low overhead and parallel execution.The executor dispatches kernels to local devices and can use multiple CPU cores or device streams.
- Runtime operations: The runtime provides over 200 standard operations and uses Eigen::Tensor, cuDNN, asynchronous copies, DMA, gRPC over TCP, and RDMA over Converged Ethernet for computation and communication.Kernels are generated for multicore CPUs and GPUs, while device transfers can overlap computation or reduce host pressure.
- User-facing features: Users build higher-level abstractions from standard operations, register additional C++ kernels when needed, and use serving, visualization, and profiling tools.TensorFlow prioritizes Python and C++ client support and includes tools for production inference, training progress, graph inspection, and distributed execution tracing.
6 Evaluation
TensorFlow exhibits low overhead at small scales and scales substantial computation across machines and GPUs for convolutional, embedding, vision, and language-model workloads. Performance limits arise from coordination, parameter-server contention, and diminishing returns, while sparse access, backup workers, and sampled softmax improve efficiency.
- Small-scale performance: On a six-core CPU with one Titan X GPU, TensorFlow has shorter convolutional-model step times than Caffe and performance within 6% of Torch.The similar TensorFlow and Torch performance is attributed to their shared cuDNN version for critical convolution and pooling operations.
- Scaling overhead: Synchronous coordination overhead rises from 1.8 ms with one worker to 8.8 ms with 100 workers for the scalar null model.These measurements represent synchronization overhead when each worker fetches one 4-byte value from each of 16 parameter-server tasks.
- Scaling overhead: Dense null-step time increases from 147 ms to 613 ms for a 100 MB model and from 1.01 s to 7.16 s for a 1 GB model as workers grow from 1 to 100.The parameters are sharded equally over 16 parameter-server tasks.
- Scaling overhead: Sparse embedding lookups take 5–20 ms regardless of whether the embedding matrix contains 1 GB or 16 GB of data.Each worker reads 32 randomly selected entries, showing that accessing only a parameter subset avoids dependence on total embedding size.
- Deep neural network training: Inception-v3 training reaches 2,300 images per second with 200 workers, but additional workers yield diminishing returns as parameter-server contention increases step time.Synchronous steps are longer than asynchronous steps because all workers wait for the slowest worker.
- Deep neural network training: For 50-worker Inception-v3 training, up to four backup workers reduce median step time, whereas a fifth slightly degrades performance.Backup workers reduce the chance that a straggler delays completion; four provide the shortest overall step time, while three are most resource-efficient.
- Language-model training: Distributed language-model training gains throughput from additional parameter-server tasks, while sampled softmax reduces softmax data transfer and computation by a factor of 78.The full softmax parallelizes multiplication and gradient calculation across parameter-server tasks; sampled softmax uses 512 classes per batch.
7 Conclusions
TensorFlow provides an extensible dataflow-based programming model that subsumes parameter-server systems and supports large-scale heterogeneous computing, production use, and experimentation. Initial adoption is encouraging, while default policies and several system capabilities remain active areas for improvement.
- Core contributions: TensorFlow’s dataflow representation subsumes existing parameter-server systems and provides a uniform model for large-scale heterogeneous computing, production tasks, and experimentation.The model is described as extensible and dataflow-based.
- Adoption: Over 8,000 people forked TensorFlow’s repository, its binary distribution was downloaded 500,000 times, and users published dozens of machine learning models using it.Google groups had also deployed TensorFlow in production, and research colleagues were using it for advances in machine learning.
- Limitations and future work: TensorFlow remains a work in progress because flexible dataflow enables excellent performance for power users, but effective default policies for most users are not yet determined.Further research on automatic optimization is intended to bridge this gap.
- Limitations and future work: The developers are actively developing automatic placement, kernel fusion, memory management, and scheduling algorithms, alongside improvements to mutable state and fault tolerance.The passage states that current mutable-state and fault-tolerance implementations suffice for applications, though the supplied text ends before specifying their limitation.
- Broader impact: The authors hope that sharing TensorFlow’s implementation and engaging with the research community will spur further research in distributed systems and machine learning.This is presented as an intended broader impact of the work.