The TensorFlow Partitioning and Scheduling Problem: It's the Critical Path!

TL;DR

This paper addresses the complex problem of partitioning and scheduling in TensorFlow's data flow graphs on heterogeneous devices, proposing heuristics that significantly improve execution time by focusing on the critical path.

Contribution

It introduces novel heuristic strategies for combined partitioning and scheduling in TensorFlow, emphasizing critical path minimization to enhance performance.

Findings

Critical path-focused heuristics outperform agnostic strategies.

Up to 4x speed-up achieved with the proposed heuristics.

Simulation results demonstrate effectiveness in communication-intensive workloads.

Abstract

State-of-the-art data flow systems such as TensorFlow impose iterative calculations on large graphs that need to be partitioned on heterogeneous devices such as CPUs, GPUs, and TPUs. However, partitioning can not be viewed in isolation. Each device has to select the next graph vertex to be executed, i.e., perform local scheduling decisions. Both problems, partitioning and scheduling, are NP-complete by themselves but have to be solved in combination in order to minimize overall execution time of an iteration. In this paper, we propose several heuristic strategies to solve the partitioning and scheduling problem in TensorFlow. We simulate the performance of the proposed strategies in heterogeneous environments with communication-intensive workloads that are common to TensorFlow. Our findings indicate that the best partitioning and scheduling heuristics are those that focus on minimizing the execution time of the critical path in the graph. Those strategies provide a speed-up of up to 4 times in comparison to strategies that are agnostic to the critical path, such as hash-based partitioning and FIFO scheduling.