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

1. Problem Statement

The paper addresses the device assignment problem in multi-GPU systems executing complex machine learning workloads (e.g., Transformers, Mixture-of-Experts) represented as static dataflow graphs. The goal is to map operations (vertices) to devices to minimize total execution time within a Work-Conserving (WC) system.

Unlike traditional bulk-synchronous frameworks (e.g., standard PyTorch or TensorFlow execution) where all devices wait for the slowest operation before proceeding, WC systems schedule operations asynchronously as soon as their data dependencies are met. While WC systems offer significant potential speedups by overlapping communication and computation, they introduce unique challenges:

• Stochastic Execution Order: Without global synchronization, the order of execution is dynamic and sensitive to hardware heterogeneity and resource contention.
• Complex Optimization: The execution time cannot be expressed as a closed-form function of the assignment due to this stochasticity.
• Competing Goals: An effective assignment must simultaneously balance GPU load and minimize inter-GPU communication, a task made difficult because load balancing is inherently temporal in asynchronous settings.

Prior learning-based approaches are limited by their reliance on bulk-synchronous assumptions, the use of a single placement policy without modeling system dynamics, and a lack of optimization during actual deployment.

2. Methodology: The DOPPLER Framework

The authors propose DOPPLER, a reinforcement learning (RL) framework that employs a dual-policy sequential decision scheme and a three-stage training pipeline.

2.1 Dual-Policy Formulation

Instead of learning a single monolithic policy, DOPPLER decomposes the assignment process into two distinct policies, modeled as an episodic Markov Decision Process (MDP):

1. SEL (Selection Policy): Selects the next operation (vertex) to assign from the set of candidate nodes. This policy traverses the partially assigned graph to approximate the non-deterministic flow of "time," effectively learning an execution order.
2. PLC (Placement Policy): Determines the specific GPU device for the selected operation, aiming to balance load and minimize communication.

Architecture:

• Both policies utilize Graph Neural Networks (GNNs) with message passing to encode the dataflow graph's structure, dependencies, and critical paths.
• SEL aggregates information from critical paths (bottom-level and top-level paths) and node features to select the next node.
• PLC combines node representations, device features (e.g., current load, end time), and the set of already placed nodes to select the optimal device.
• Efficiency Optimization: To avoid prohibitive training costs associated with re-running message passing at every step, DOPPLER performs message passing once per episode and updates device features dynamically without re-computing graph embeddings at every step.

2.2 Three-Stage Training Pipeline

DOPPLER employs a "learning-by-doing" paradigm across three stages to ensure convergence and adaptability:

1. Stage I (Imitation Learning / Offline): The policies are pre-trained using supervised learning to mimic a heuristic teacher (specifically, the Critical Path heuristic). This accelerates convergence by providing a strong initial policy before exploration begins.
2. Stage II (Simulation-based RL / Offline): The policies are refined using reinforcement learning within a software simulator (implementing the WC execution logic). The reward is derived from the simulated execution time. This stage allows for rapid iteration without the latency of real hardware.
3. Stage III (Real-System RL / Online): The trained policies are deployed in a real WC system. They are continuously refined via online RL using actual execution times as rewards. This stage adapts the policy to specific hardware characteristics (e.g., NVLink jitter, contention) that simulators cannot perfectly capture, with no additional overhead to the user.

3. Key Contributions

1. Problem Formulation: The paper investigates device assignment specifically within the context of asynchronous, work-conserving multi-GPU systems, addressing the stochastic nature of execution order.
2. Dual-Policy Approach: It introduces a separation of concerns where one policy learns the approximate traversal order of nodes and the other handles device placement, capturing both execution dynamics and hardware constraints.
3. Three-Stage Framework: DOPPLER proposes a hybrid training strategy combining imitation learning, simulation-based RL, and online real-system RL. This approach accelerates convergence and enables continuous optimization during deployment.
4. Scalable GNN Implementation: The authors propose an efficient message-passing approximation that reduces training time significantly without compromising convergence.

4. Experimental Results

Experiments were conducted on four NVIDIA Tesla P100 GPUs using various workloads: Chain Matrix Multiplications (CHAINMM), Feedforward Neural Networks (FFNN), and Llama transformer blocks/layers.

• Performance vs. Baselines: DOPPLER (specifically the fully trained DOPPLER-SYS) consistently outperformed baselines including the Critical Path heuristic, PLACETO (single-policy RL), GDP (graph embedding RL), and a custom enumerative optimizer.
  • DOPPLER achieved up to 52.7% lower execution times compared to the best baseline (GDP).
  • It reduced runtime by up to 78.2% compared to the Critical Path heuristic.
  • It outperformed the custom enumerative optimizer by up to 13.8%.
• Ablation Studies:
  • Dual Policy: Combining both SEL and PLC policies yielded the best results, particularly for complex models, demonstrating the necessity of learning both order and placement.
  • Training Stages: Training on the real system alone led to slow convergence and instability. The combination of imitation learning (Stage I) and simulation (Stage II) was crucial for stable and fast convergence before online refinement (Stage III).
• Transfer Learning: DOPPLER demonstrated strong generalization. Policies trained on simple graphs (FFNN) could be fine-tuned with few-shot episodes (2k–4k) to perform well on unseen architectures (Llama) and different hardware configurations (e.g., transferring from 4 P100s to 8 V100s).
• Scalability: The training and inference times of DOPPLER scale linearly with the size of the dataflow graph, outperforming other RL-based baselines in efficiency.

5. Significance and Claims

The paper claims that DOPPLER represents a significant step forward in optimizing modern ML workloads by moving beyond static, bulk-synchronous assumptions. Its significance lies in:

• Practical Efficiency: By enabling work-conserving execution, DOPPLER can save substantial GPU hours in large-scale deployments (e.g., saving millions of GPU hours annually for ChatGPT-scale workloads).
• Adaptability: Unlike static heuristics or purely simulation-based models, DOPPLER's online learning capability allows it to adapt to the specific, noisy realities of production hardware environments.
• Methodological Innovation: The dual-policy decomposition offers a more nuanced approach to the combinatorial explosion of device assignment problems, separating the "what to run next" from "where to run it," which aligns better with the temporal nature of asynchronous execution.

The authors emphasize that their approach does not require changes to the underlying ML frameworks but rather optimizes the scheduling layer, making it a practical solution for existing multi-GPU infrastructures.