TensorFlow: Large‑Scale Machine Learning on Heterogeneous Distributed Systems – Overview and Implementation

TensorFlow: Large‑Scale Machine Learning on Heterogeneous Distributed Systems (2015.11) Authors: Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Gregory S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian J. Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Józefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mané, Rajat Monga, Sherry Moore, Derek Gordon Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul A. Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda B. Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, Xiaoqiang Zheng.

Programming Model and Core Concepts

TensorFlow represents computation as a directed acyclic graph (DAG). Nodes (operations) are connected by edges that carry tensors (multi‑dimensional arrays). Control‑dependency edges impose ordering without data transfer.

Operation (Op): An abstract computation with a name and attributes; polymorphic over tensor types.

Kernel: A device‑specific implementation of an Op.

Session: Client interface that runs parts of the graph, feeding inputs and fetching outputs.

Variable: A persistent mutable tensor whose lifetime spans multiple graph executions.

Device: Physical compute resource (CPU, GPU, etc.). Example device name: /job:localhost/device:cpu:0.

Tensor: Typed multi‑dimensional array (8‑bit to 64‑bit integers, floating‑point, complex, string) with reference‑counted backing store.

Implementation

Single‑Device Execution

All nodes run locally; each node maintains a dependency counter and is scheduled when the counter reaches zero.

Multi‑Device Execution

Two main challenges: deciding on which device each node runs, and managing cross‑device communication.

Node Placement

TensorFlow uses a heuristic cost model that simulates execution and greedily assigns nodes to devices that can execute the required kernel, preferring the device with the smallest estimated execution time. The placement algorithm is still an active research area.

Device‑to‑Device Communication

Cross‑device data transfer is realized with Send and Recv nodes that break an edge into two parts. This isolates the communication mechanism from the computation graph, allowing the master to issue only high‑level Run requests without managing inter‑worker traffic.

Distributed Execution

Extends multi‑device execution across multiple machines and adds fault tolerance. Errors in Send/Recv or keep‑alive checks cause the entire graph to abort and restart. To avoid recomputing already‑finished nodes, TensorFlow checkpoints state periodically using Save and Restore nodes that persist variables to distributed storage.

Extensions

TensorFlow includes built‑in automatic differentiation, advanced memory management (e.g., recomputation, swapping tensors between GPU and CPU), and support for partial graph execution.

Control‑flow primitives: Switch, Merge, Enter, Leave, NextIteration.

Input operation nodes that read data directly from files.

Queues (FIFO and shuffling) for asynchronous data feeding.

Containers to hold long‑lived mutable state across sessions.

Optimization

Common subexpression elimination.

Communication‑aware scheduling and memory usage control.

Asynchronous kernels.

Optimized kernel libraries (cuDNN, cuda‑convnet, cuBLAS).

Lossy compression (e.g., 16‑bit representation of 32‑bit floats) to reduce network traffic.

Practical Experience

Lessons from migrating models include understanding parameter counts, scaling from small to large models, verifying loss functions, starting on a single machine before scaling out, guarding against numerical errors, and cross‑checking results between systems.

Data Parallel, Model Parallel, and Concurrent Steps

Data Parallelism

Multiple replicas of the model process different data shards. Synchronous data parallelism aggregates gradients from all replicas before updating the model, while asynchronous parallelism lets each replica update parameters independently.

Model Parallelism and Concurrent Steps

Model parallelism splits a single batch across multiple devices, each handling a different part of the computation (e.g., layers of a deep LSTM). Concurrent steps pipeline multiple iterations on the same device set to improve throughput.

Tools

TensorBoard visualizes the computation graph, summaries, and training progress. Performance tracing helps identify bottlenecks.

Future Work

Reusable sub‑graphs (function‑like abstractions).

Improved placement and scheduling, possibly using deep reinforcement learning.

Just‑in‑time compilation.

Cross‑operation dynamic compilation frameworks.

Conclusion

TensorFlow is a data‑flow programming model that performs well when sufficient memory is available, but its speed can lag behind frameworks such as MXNet, leading some to nickname it “TensorSlow.” Nevertheless, its popularity stems largely from Google’s backing and extensive ecosystem.

Reference: [Translated] TensorFlow Whitepaper: http://www.jianshu.com/p/65dc64e4c81f