DOPPLER: Dual-Policy Learning for Device Assignment in Asynchronous Dataflow Graphs

TL;DR

DOPPLER introduces a dual-policy learning framework for device assignment in dataflow graphs, significantly improving execution time and training efficiency for complex machine learning workloads.

Contribution

It presents a novel three-stage dual-policy framework that incorporates system-aware heuristics, addressing limitations of prior reinforcement learning approaches.

Findings

Outperforms baseline methods in reducing execution time.

Demonstrates higher sampling efficiency and faster training.

Effective in complex machine learning workloads.

Abstract

We study the problem of assigning operations in a dataflow graph to devices to minimize execution time in a work-conserving system, with emphasis on complex machine learning workloads. Prior learning-based methods often struggle due to three key limitations: (1) reliance on bulk-synchronous systems like TensorFlow, which under-utilize devices due to barrier synchronization; (2) lack of awareness of the scheduling mechanism of underlying systems when designing learning-based methods; and (3) exclusive dependence on reinforcement learning, ignoring the structure of effective heuristics designed by experts. In this paper, we propose \textsc{Doppler}, a three-stage framework for training dual-policy networks consisting of 1) a $\mathsf{SEL}$ policy for selecting operations and 2) a $\mathsf{PLC}$ policy for placing chosen operations on devices. Our experiments show that \textsc{Doppler} outperforms all baseline methods across tasks by reducing system execution time and additionally demonstrates sampling efficiency by reducing per-episode training time.