DOPPLER: Dual-Policy Learning for Device Assignment in Asynchronous Dataflow Graphs
Authors: Xinyu Yao
Daniel Bourgeois
Abhinav Jain
Yuxin Tang
Jiawen Yao
Zhimin Ding
Arlei Silva
Chris Jermaine
Venue: ICLR 2026 Poster (2026)
Keywords: Distributed Systems; Reinforcement Learning; Graph Scheduling;
Primary area: reinforcement learning
TL;DR: We can do better device assignment of dataflow graphs on multi-GPU systems using dual policy networks, training with real system during deployment, and other techniques.

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 approaches face three limitations: (1) reliance on bulk-synchronous frameworks that under-utilize devices, (2) learning a single placement policy without modeling the system dynamics, and (3) depending solely on reinforcement learning during pre-training while ignoring optimization during deployment. We propose Doppler, a three-stage framework with two policies—$\mathsf{SEL}$ for selecting operations and $\mathsf{PLC}$ for placing them on devices. Doppler consistently outperforms baselines by reducing execution time and improving sampling efficiency through faster per-episode training. Our results show that Doppler achieves up to 52.7\% lower execution times than the best baseline. The code is available at https://github.com/xinyuyao/Doppler.

Published as a conference paper at ICLR 2026
DOPPLER: DUAL-POLICY LEARNING FOR DEVICE ASSIGNMENT IN ASYNCHRONOUS DATAFLOW GRAPHS
Xinyu Yao
Daniel Bourgeois'
Abhinav Jain'
Yuxin Tang'
Jiawen Yao!
Zhimin Ding'
Arlei Silva 1.2
Christopher Jermaine'
Rice University {xy38, dcb10, aj70, yt33, jy75, 2d21, arlei, cmj4/erice.edu
Rice Ken Kennedy Institute
ABSTRACT
We study the problem of assigning operations in dataflow graph to devices minimize execution time work-conserving system, with emphasis on complex machine learning workloads. Prior learning-basee approaches face three limita- tions: (1) reliance on mulk-synchronous frameworks that under-utilize devices, (2) learning a single placement policy without modeling the system dynamics, and (3) depending solely on reinforcement learning during pre-training while ignoring opti- mization during deployment. We propose DOPPLER, three-stage framework with two policies-SEB for selecting operations and PLC for placing them on devices. DOPPLER consistently outperforms baselines by reducing execution time and im- proving sampling efficiency through faster per-episode training. Our results show that DOPPLER achieves up to 52.7% lower execution times than the best baseline. The code is available at https: :/p::://// thub. com/ kinyuyac/Dopples
INTRODUCTION
Existing systems for multi-GPU computing such PyTorch Paszke al. 2019). TensorFlow (Abadi et al. 2016), and the JAX-based Google stack (Frostig et al. 2018) proceed through a computation in lock-step, level-wise fashion. Consider the three-matrix multiplication chain Y Z in Fig. la, each matrix is partitioned into four submatrices to be distributed to eight GPUs to be implemented using the optimal 3D algorithm (Agarwal al.) 1995). The resulting additions, multiplications, and transfers form the dataflow graph in Fig. 1b.
In this graph, X Y is first computed through pairwise matrix multiplications (e.g., X11 x Yi1) distributed across the GPUs. Once these partial results are pro- duced, they must be aggregated and moved, typically using collective communication library (Woolley] 2015). Then, Z is com- puted from X Y. which requires another set pairwise multiplies, followed by an aggregation step.
X11
X12
Y
Y12
Z11
Z12
X21
Xx2
Y21
Yza
Zz1
Zz2
(a)
MAdd
Medd
Maset
MMU
Mndu
MM
Mida
Z11
Z
Z
Zez
Key:
GPU1
GPLI
In most deep learning systems, such steps executed bulk-synchronously (Valiant 1990). For example, until all pairwise mul- tiplications are completed, aggregation can- not begin. This leads to idle resources and lost opportunities for speedup: one slow multiplication delays the entire step (Li ct al., 2020c). Moreover, the aggregation phase communication-dominaeee during which GPUs are severely underutilized.
MMa
OPUS
GPLIS
GPLS
@PU4
X12
Xx1
Y12
Y22
GPUT
GPUB
(b)
Figure 1: A dataflow graph (b) corresponding to the decomposed chain from (a). Vertices correspond to computation kernel calls, edges data dependencies, and colors the mapping of computations GPUs.
"Corresoonding Author.

---

Published as a conference paper at ICLR 2026
A more efficient approach would overlap "reduce" with computation by scheduling operations (data transfers and kernel calls) asynchronously, as soon as they can be run. A scheduler that never willingly allows resources to sit idle and schedules them dynamically is called work-conserring (WC) [Kleinrock] 1965). Table 1 shows the potential speedup of a WC system on two workloads: a chain of matrix multiplications and additions, CM, and feedforward neural network, FFNN (configuration details are provided in Appendix F). On a GPU server, the WC system achieves a reduction of 46.3 ms (33%) for CHAINMM and 26.7 ms (53%) for FFNN, compared to PyTorch. In the context of long-term deployment, these small per-query millisecond reductions accumulate into substantial GPU-hour savings. For instance, a 24.2 mss reduction per query for running Llama model (Yaadav 2024) amounts to more than 2.4 million GPU hours saved annually at ChatGPT-scale workloads (assuming one million queries per day). Additional results are shown Section 6.
While WC systems can improve efficiency, they introduce significant challenges, particu- larly for GPU assignment. In contrast to bulk- synchronous settings that fix execution order via all-reduce, WC systems rely asynchronous point-to-pon boint-to-point communication; without global synchronization, execution order is harder to control, and performance becomes highly sen- sitive to hardware heterogeneity and resource contention. An effective device assignment must capture this information and balance two competing goals: (1) maintaining GPU load balance and (2) minimizing inter-GPU communication (Harlap et al. 2018 Lu et al. 2017). Traditional placement methods emphasize communication minimization, but under WC scheduler, the stochastic execution order makes load balancing especially challenging.
MODEL
WC SYSTEM
SYNCHRONOUS
CHAINMM
139
185.3
FFNN
50.2
76.9
Table I: Execution time (in milliscconds) for ex- ecution in a work-conserving system and a syn- chronous system.
as load balancing is inherently temporal (Saha et al. 2019).
In this paper, we tackle the device assignment problem in an asynchronous WC system by introducing DOPPLER, a reinforcement learning-based framework that adopts a cearinngbyy-doing paradigm to optimize device placement. While prior work Addanki et al. (2019; Zhou et al. (2019); Mirhoseini al. (2117) learns single placement policy, DOPPLER learns efficient kernel assignments through dual- policy sequential decision scheme. The first policy selects the next kernel (vertex in the dataflow graph) to assign by traversing the partially assigned dataflow graph in a manner that approximates the nondetteriinistic flow of "time," " while the second policy determines the GPU placement of that kernel balance load and minimize communication. This separation captures both execution dynamics and hardware constraints, producing assignments tailored to stochastic WC execution.
Moreover, DOPPLER employs three-stage framework for training its dual policies. Stage I (offline) uses supervised learning to train the policies to follow simple heuristics, such as co-locating neighbor- ing vertices on the same device. Stage II (offline) transitions to reinforcement learning: the policies generate assignments that are "executed" in a simulated WC system, with rewards computed from the simulated runtime. Stage III (online) deploys the trained policies in a real WC system, where they are continuously refined through reinforcement learning using rewards derived from observing assignment runtimes, and the policies are recursively updated as the system executes.
Our contributions are the following: (1) We investigate the device assignment problem in multi- GPU system under a WC system; (2) We introduce a dual-policy learning approach to first learn the approximated traversal order of nodes before assigning them to devices; (3) We propose DOPPLER, three-stage training framework to improve scheduling efficiency on the fly by continuously training during deployment, along with two pretraining stages to accelerate convergence during deployment; (4) Our experiments show that DOPPLER achieves up to 52.7% lower execution times compared to the best baseline. DOPPLER also achieves significant runtime reduction compared to a stronger baseline that we designed ATVOOPIIII,E by up to 13.8%.
DEVICE ASSIGNMENT PROBLEM IN A WORK-CONSERGE
SYSTEM
In this paper, we study the device assignment problem for static dataflow graphs, which capture modern ML workloads. Sfate-of-the-ar architectures such as Transformers (Vaswani et al. 2017) and Mixture-of- Experts models (Jacobs et al.) [991] Shazeer et al.) (017) are defined over largely fixed computation graphs during both training and inference (Shi al. 2024). Crucially, while the
2

---

Published as a conference paper at ICLR 2026
graph structure is static, its execution is fully asynchronous: the runtime schedules and executes operators in parallel, constrained only by data dependencies the graph.
Formally, consider dataflow graph G (V.E), where notes set of vertices representing compu- tations, and & {el, e2..... Cml denotes directed edges representing data dependen- cies. Given a set of devices D, our goal is to compute device assignment A that maps each vertex U E V to device in D. We denote by A, the device assigned to vertex e under assignment A. The objec- tive is to find an assignment. A that mini- mizes the overall execution time, denoted by ExecTime(A))
{V1, 1/2, ..., )n de-
Algorithm 1 ExecTime(A)
% rdy(u, df is T iff the result of vertex % v is on device d; initially, nothing is ready rdylu, cl % except inputs: available everywhere
v(r V e D)
rdylu, df T Y(u V. D) s.t. (v', &
t + 0 % exec begins at time
S () % schedule is empty
while 3(v E V) s.t. rdyle, A,] = do
tasks + EnumTaske (rdy, A, S)
lask +- ChooseTask
A, S, lasks)
task = mull then
For example dataflow graph, consider the matrix multiplication chain Y Z, which can be decomposed to server with eight GPUs by sharding each matrix four ways in Fig. 1a. The resulting fine-grained dataflow graph (Fig. lb) con- tains eight submatrix multiplies associated with the two original multiplies, and four matrix additions to aggregate the results. candidate assignment of the graph to eight
% if no task is chosen, just wait (1, task)) ~ P(.[S, t) % which task done?
S S + (task, end) % save completion rdy rdy(vertex(task)) devi ce(tssk) <-
else
S S + (lask, L beg) % record initiation end if end while return
GPUs is shown in Fig. 1b. This assignment achieves low execution time as (a) the expensive matrix multiplications that tend to run in parallel are load-balanced, and (b) communication is minimized by co-locating neighboring nodes.
Therefore, key question is how to de- fine the execution time of an assignment ExecTime (A) in WC systems? In prac- tice, is difficult to derive a closed formula for ExecT me(A), (A), given the stochasticity of WC systems. As operations are issued dynamically, based on the state of the sys- tem, different runs of the same assignment can have very different execution times.
Algorithm 2 EnunTaskksrdyy.44 A,
output +- {
get all potential transfers for (V1, (1), E E do
rdy(ui, Ang] = false and rdy[e), A = true and transfer(Vi, A Ang) & S then
Add transfer(v,, Auy Ang) to outpul end if end for Algorithm describes how an assignment A is executed in WC system (the sub- routine EnumTasks (rdy, A, S) in Algo- rithm enumerates the tasks that can be taken at each step by the scheduler). The algorithm stochastically simulates the ex- ecution of the assignment A via WC dy- namic scheduler and returns the total exe- cution time. It works by repeatedly asking the scheduler to choose the next task to schedule; when there is task that can be scheduled, the algorithm waits until an event is stochastically generated. In the algorithm, a schedule the complete list of events that have occurred up to tix- An event is a (task, time, eventtype) triple, where task is either an data transfer transfer between devices or an execution (exec) of node on device, time records when the transfer exec event occurs, and eventtype specifies the recorded time as either beg or end of an event. EnumTasks enumerates all t.llt.a transfer and exec tasks that are ready when EnumTass Taaks called.
% get all potential ops to exec for 1/2 E V do rdy(vi, AngJ true Voj s.t. (v1, 02) E and exec(u2, Ang) & then
Add exec (12, Aug) to output end if end for return oulput
Algorithm has two key generic components that allow it to serve reasonable proxy, or digital twin, for real-life scheduler executed on real-life hardware. First, the distribution P P((fout; a)
3

---

Published as a conference paper at ICLR 2026
Current state
RL agent
Next state
Device embedding
Algorithm 3 ALSIGN(SEL,,PLCA
SEL, Policy Natwork
A+ for l' E [I, do Choose t SELg(A.9)
PLC, Policy Network
Choose d t (A, U,
Node embedding
Add (v - d) A end for return A
device
In device
Candidate node
Step
Step
Figure 2: (Left) The ASSIGN algorithm, which sequentially produces an assignment A using SEL, policy and placed using PLC, policy. (Right) A graphical depiction of the algorithm' implementation.
(n)) governs the "next completed task." Given schedule S and current time Lin, P is joint distribution over the next task to complete and the time tout at which this task completes. Second, the function ChooseTast encapsulates the underlying scheduling algorithm that is implemented by the WC system. It may choose any lask from a set of tasks. As described, it may operate first. (seeking to probe deeply into 9), breadth-first, or may employ any other applicable strategy.
In practice, Algorithm is implemented by either (a) simulator, where the distribution P((fout: a) S, f.)) is realized by a model that takes into account factors such as the number of floating point operations in the underlying operation (e.g.. CUDA kernel executions) or the number of bytes to be transferred (in the case of GPU-to-GPU communication), thereby simulating the event-driven behavior of a real system, (b) deploying the assignment A in real-world WC system and observing the runtime. We use option (a) in Stage of DOPPLER and option (b) in Stage DOPPLER (see details in Section 5).
3 A SOLUTION VIA REINFORCEMENT LEARNING
As there is no closed-form objective function, bearningbby-doing is promising approach for deter- mining the optimal assignment A" Stages II and III of DOPPLER formulate choosing the assignment as bandit problem. For an assignment A at time tick I, we obtain reward r; RA where RA is the reward distribution for (in practice, Fi is sampled by invoking ExecTime (A) and observing the runtime). Let R RA where A arg m E R[RA]. Our goal is to minimize the regret:
T
[E[R*] r2)
This is bandit problem with DI arms. Our problem is not amenable to classic solutions because the number of arms is large. For example, if we are producing an assignment for 8-GPU server that is to execute dataflow graph with 100 vertices, there are z 2300 possible assignments.
Fortunately, there combinatorial structure the assignment problem that allow deal with the very large set of possible assignments (Cesa- Bianchi & Lugosi| 2012 Chen et al. 2013). Note that if two assignments A₁ and Ag differ only in how a few vertices have been assigned to devices, it is likely that the two reward distributions RA, and RA, will be similar. Thus, it may be possible to systematically search the possible assignments.
Our approach builds the assignment at each time tick using sequential process controlled by two policies SEL, and PLCp. If we use ASsIGN(SELa, PLC₃) (Algorithm 3) as the mechanism for choosing the assignment A at each time tick t of the bandit problem in Equation 1, this process can be reformulated as an episodic Markov decision process (MDP), as shown in Liu et al. (2024), where each episode executes ASsIGN(SELL, PLC,) (shown in Fig. 2) and reward is obtained once the assignment is completed at the end of each episode. Therefore, any suitable reinforcement learning algorithm can be used to learn the policies SEL, and PLC,. In our implementation, apply graph
4

---

Published as a conference paper at ICLR 2026
neural networks (GNN) along with message passing to encode graph G and use feedforward neural networks to decode actions for SELg and PLCg. Details on GNN architectures are in Section 4.2.
DOPPLER DUAL POLICY IMPLEMENTATIONS
We describe our episodic Markov Decision Process formulation (Section 4.1), along with the graph neural network architectures that used implement the dual policy learned during DOPPLER training (Section 4.2) and an efficiency analysis of the GNN message-
message-passing passing (Section 4.3).
4.1 EPISODIC MDP FORMULATION
We formulate device assignment as episodic Markov Decision Process (S,A. A. H.P. P, P.)., where P is the transition function, and is the horizon, which equals the number of nodes in the graph.
States Each state 8;, S is tuple (Xg, Cn, Cn, XD,),, where Xg (Xv, Xg) represents the static graph features, including node and edge features such as bottom-level paths (i.e., to entry nodes), top-level paths (i.e., to exit nodes), and communication costs. Cn. is the dynamic nodes, and Co is defined as the set of entry nodes in the graph. Finally, XD.A represents the dynamic device features (e.g., total computing time and the end time for computations on each device). Details about the sets of features can be found in Appendix E.
of candidate
Actions At each time step h, the agent takes an action a₁ A. where an (Un, dr). Each action selects node and places it on device d, using policies SELØ and PLC#, shown in Fig. 2. Each episode is composed of | iterations i.e., one iteration per node in V.
Reward We calculate rewards using the execution time Re, = (-1) x ExecTime(sH) derived from either the real system or simulator. The reward computed at the end of each episode (h H). with intermediate rewards set to zero for efficiency. To enhance stability, we subtract baseline reward equal to the average execution time observed across all previous episodes (Reg). The final reward computed as тH
Rag Rag
4.2 DUAL POLICY GRAPH NEURAL NETWORK ARCHITECTURES
Device assignment requires reasoning over dataflow graphs whose behavior is shaped by dependen- cies, critical paths, and communication edges. Graph neural networks naturally provide the right inductive bias for this setting: they propagate information along dependency edges, capture local structure, and generalize across graphs of varying sizes. The dual-policy decomposition is motivated by heterogeneity, since decisions depend jointly on both the node being scheduled and the device which it may run. This modular design reflects the semantics of the dynamic environment and allows the same architecture to adapt across different workloads and hardware configurations.
We further describe the details of the policy networks for computing SEL, and PLC, in Algorithm 3. The symbol 0 is used here emphasize that these functions have parameters that will be optimized as part of the training process. DOPPLER applies a graph neural network (GNN) to encode node information in the dataflow graph. Our GNN is a message-passing neural network (Gimmer al. 2017) that learns node representations for each node via K successive iterations:
h. M]
(hlk- "It hift-), ewv.))
(2)
(EN(2)
where h!d are representations learned at the k-th layer, h Xv(v]. ie and é are functions, N(v) operator. We will use GNN(G,
are the neighbors of and + is a pernutation-invant hh [K), 1K [.k. to refer to the representations of all nodes in g.
We also apply an I-llayer feedforward neural network (FFNN) to encode node information: x! [[ [[-1]] + hldl where x!" representations at the I-th layer, and Will and blll are weights and
biases, respectively. Let FFNN(X) = |x[ [.. x [M]] be the representations of all nodes in G. with !0) X(v).
5

---

Published as a conference paper at ICLR 2026
Node policy network (SELg): Selects node from the candidate set C based the observed graph state Xg using the e-greedy approach. Let b(e) and t()) be the b-path and t-path for v; where b-level (t-level) path for is the longest path from to entry (exit) node in G. We aggregate information from these critical paths via GNN embeddings H [u] for each node u along them. Nodes are selected according probabilities estimated from graph embedding matrix Hg, which is the result of the concatenation of critical path (he,s and h), GNN (H [e), and feature (2(p)) representations.
H (NN(G,,e)) Z FFVN(XY), h.
[[o] hag bu.e |2(e]]
aar maxy Qo(v), [random E.C. C, =
= 1 1-6,
Qg (v) softmax(FFNN(h.)) SELo(G,
(4)
where p the probability the event and e a parameter.
Device policy network (PLCg): Places node into one of the devices in D based the composed state observation Xp. Xg). Devices are selected based on an embedding matrix Hd for the set of devices, generated by concatenating representations for the node (H [v]), node features (Z(u)), device features (Y (d) and nodes already placed on the device (ha).
H GNN(G, Xg), Y FFNN(XD), Z FFNN(XXg), p,d [D] ha Y[d] [
(5)
HD ha,1i.... ho,m]:
(6)
11 LeaNY
QD(d)
(7)
arr max, Qp(d),
p p-1-c
P-)x(-) (random E D, p= 6.
(8)
4.3 EFFICIENT MESSAGE PASSING APPROXIMATION
Implementing the MDP described in Section 4.1 requires performing message- paasssggg on the dataflow graph when calling SEL, and PLC, policies each MDP step h. We found this be prohibitive for large graphs since we may apply up to 8k episodes x 261 steps 2m steps in our experiments. PLACETO (one of our baselines) suffers from this issue, being very inefficient during training as it performs one message-passing round per MDP step. Instead, we propose performing message passing on the graph only once per MDP episode and encoding updated assignment information at each step h in XD,,, without message passing. We found empirically that this modification has negligible impact on DOPPLER'S 's convergence but leads to significant reduction in training time, especially for large neural networks (we show an ablation study in Appendix G.3).
Imitation Learning Stage
Simulation-based RL Stage
Real System RL Stage action
action
Dual Policy
Dual Policy
Dual Policy
SELECT,
PLACE, Palles
SELECT,
PLACE, Palez
SELECT
PLASE Palky
New state
New state
Complete Assignment
New state
Complete Assignment
IL Loss
Heuristic
Reward
Simulatorr
Reward
Real
Figure Three-stage framework for cost-effective training DOPPLER, combining imitation learning, iouulltiio-bbsee reinforcement learning, and real-system reinforcement learning.

---

Published as a conference paper at ICLR 2026
5 DOPPLER: COST-EFFECTIVE TRAINING
DOPPLER adopts a three-stage framework to train and deploy its dual policies (SEL, and (LC)): (Stage I) imitation learning, (Stage II) simulation-based reinforcement learning, and (Stage III) real- system reinforcement learning (RL) as shown in Fig. 3. Imitation learning serves as a pre-training step that accelerates convergence in subsequent RL stages by leveraging an existing list-scheduling heuristic. The two RL stages further realize a "learning- boniingggb-oiin paradigm, enabling the policies to be continuously refined through interaction with simulated and real systems.
Imitation Learning stage (Stage I). We propose using imitation learning for teaching the dual policy to replicate the decisions of existing heuristic (teacher) before deploying it to the real system. During imitation learning, the computation graph emits the current assignment state at each step, which is fed to both the dual policy and the heuristic. We then apply the decision of the CRITICAL PATH heuristic (Ocp) to obtain the objective function:
J(0)= Beee (Ilep()), ( (Vo log [Tpppp)))]
(9)
Simulation- reinforcement learning stage (Stage II). Even after pre-training with teacher, we may have policy that is too low quality for deployment in the real system, where longer running times or slow convergence due exploration may be unacceptable. Thus, also train DOPPLER simulator that implements Algorithm 1. The simulator is invoked only at the end of an episode, after we obtain the assignment for all nodes. The dual-policy networks are updated using the policy gradient method (Sutton et al. 1999) by maximizing the following objective function:
using a software-
[(] = (o) Ve (a(ls))
(10)
Real-system reinforcement learning stage (Stage III). This stage is analogous to sim-to-real adaptation in RL-based hardware scheduling as no simulator can perfectly capture dynamic GPU contention, NVLink jitter, or system-level scheduling noise. Because the first two stages already produce high-quality dual policy, deployment does not subject users to long warm-up periods or unstable exploratory behavior. Stage III performs lightweight online refinement: once deployed, the policy adapts to the target system's s specific hardware configurations and workload characteristics. Importantly, the reward signal for optimizing Equation 10 is obtained "for free": the system naturally observes real execution times (ExecTimee Time) while serving requests. This allows DOPPLER to continuously improve in production with no additional overhead or disruption to users.
4 GPUS
RUNTIME REDUCTION
MODEL
CRIT.
PATH
PLACETO
GDP
ENUMÖPT.
DOPPLER-SHP SIM DOPPLER-SYS BASELINE ENUMÖPT.
CHAINMM
230.4 4.3 137.1 2.2 198.0 = 139.0 10.0
122.5 + 4.0
123.4 2.5
10.7%
11.9%
FENN
217.8 11.3 126.3
100.3 i 3.2 50.2 2.5
4.19
4.7
52.7%
5.6%
LLAMA-BLOCK
230.9
411.5
19.733665. i 184 172.7
191.5 6.0
160.3 1 4.3
30.6%
7.2%
LLAMA-LAYER 292.6 = 5.8 295.1 7.0 231.5 5.1 1748 4.7
167.0 3.4
150.6 + 4.2
48.5%
13.8%
Table 2: Real engine execution time (ms) for assignments produced by our methods (ENUMOPT., DOPPLER) and prior baselines. We also show the runtime reductions achieved by DOPPLER-SYY compared with previously proposed basclines and ENUMÖPT.
EXPERIMENTS
We study the following questions, focusing on the quality of assignments produced by DOPPLER (Q1-Q5) and the scalability of DOPPLER'S 's dual-policy networks (Q6). Q1: How does DOPPLER compare with alternative approaches in terms of the execution time achieved by its generated assignments? Q2: What are the individual contributions of the SEL, and PLC, policies to assignment performance? Q3: How do imitation learning (Stage I), iimulation-based RL (Stage II), and real- system RL (Stage III) jointly improve the performance of DOPPLER'S 's training? Q4: How can exccution time be explained based on the assignments produced by DOPPLER? Q5: Do DOPPLER's learned qual-policy generalize from one dataflow graph to unseen graphs and hardware architectures? Q6: How do DOPPLER'S policy training cost and inference overhead scale with graph size?
7

---

Published as a conference paper at ICLR 2026
6.1 EXPERIMENTAL SETUP
Neural network architectures. We test dataflow graphs from four types of neural network archi- tectures in our experiments: feedforward neural network (FFNN), chain matrix multiplications (CHAINMM), a Llama transformer block (LLAMA-BL.OCK) and a complete Llama transformer layer wwwttaalllbbbeanaaeeeeeee
(LLAMA-LAYER). FFNN and CHAINMM capture the dense common across classical neural networks Markidis et al. (2018), while transformer blocks represent the dominant architecture contemporary large language models Vaswani al. (2017). Further details can be found in Appendix D.
workloads
GPU systems. We compare the assignment approaches using four NVIDIA Tesla P100 GPUs with 16GB of memory each. Moreover, we conducted ablation studies on (1) restricted GPU memory (8GB out of 16GB) and (2) eight NVIDIA V100 GPUs with 32GB memory, and we show the results in Appendix H.
Baselines. We compare our approach against four baselines. CRITICAL
4 GPUs
PATH Kwok & Ahmad (1999) a
MODEL
SYS
SEL
PLC
popular (non-learning) heuristic for
DAG device assignment. PLACETO
CHAINMM
123.4+2.5
127.0=0..8
121.6=0.77
Addanki et al. (2019) is recent RL- based alternative that applies a sin- gle (device) policy and is trained us- ing simulations (see the ablation study on the simulator in Appendix G.1). GDP Zhou et al. (2019) is another recent RL-based method that consists of graph embedding and sequential at- tention. ENUMERATIVEOPTIMPEE structure of sharded tensor computation to produce high-quality assignment (described in detail in Appendix B).
FFNN
47.410.7
59.117.6
6..2
160.3
175.6#4.1
172.9+4.3
LLAMA-LAYER
150.6
161.7.14.1
159.514499
Table 3: Ablation study showing real engine execution times (ms) on GPUs. SYS, SEL, and PLC denote DOPPLER- SYS, DOPPLER-SEL. and DOPPLER-PLC, respectively.
is a baseline we developed-it our best effort at exploiting the
plyperparameters For RL-based methods, we run 4k episodes for CHAINMM and FFNN, and 8k episodes for LLAMA-BLOCK and LLAMA-LAYER. We tried different learning rate schedules for each method (initial values {le 3, 1e 4, 1e-1e) 5)) and found that 1 decreasing linearly to le-6 works best for PLACETO and le-4 linearly decreasing to le-7
7 works best for all versions of DOPPLER and GDP. We apply 0.5 exploration rate linearly decreasing to 0.0 for PLACETO and 0.2 linearly decreasing to 0.0 for DOPPLER and GDP. An entropy weight of 1c-2 is applied for all RL methods. For CRITICAL PATH, we run 50 assignments and report the best execution time. The reported execution time (and standard deviation) using the real system is the average of 10 exccutions.
4 GPUS
TRAIN MODEL TARGET MODEL ZERO-SHOT 2K-SHOT
4X-SHOT
DOPPLER-SYE SYS CRIT. PATH PLACETO
ENUMÖPT.
FENN
LLAMA-BLOCK
251.0 165.3 = 159.4 4,8
160.3 = 4.3 230.9
8.7 411.5 = 19.7 172.7
CHAINMM
LLAMA-BLOCK
242.3
1849
1740
444
160.3 1 4.3
230.9
101.5:1100 19.7 172.7 411.
FENN
LLAMA-LAYER
206.1 4.5 158.2
155.8
5.0
150.6
4.2
292.6
5.8
+ 174.8
CHAINMM
LLAMA-LAYER
338.2 5.0 164.4
3.3
156.44444
150.6 = 42 292.6 5.8 295.1
7.0
7.0 174.8 4
Table 4: Real engine execution time (ms) for assignments produced by DOPPLER under different few- shot transfer settings (zero-shot, 2k-shot, 4k-shot) compared against full training (DOPPLER-SYS) and baselines. DOPPLER achieves performance comparable full training with only 4k-shot transfer.
6.2 RESULTS AND DISCUSSION
Comparison between our solutions and existing alternatives (Q1). Table 2 reports execution times across neural network architectures on four GPUs. DOPPLER-SY. -YS adopts all three training stages described in Section 5, and DOPPLER-SIM adopts only Stages and II. The results show that DOPPLER-ST SYS outperforms all baselines in most settings, with DOPPLER-SIA often ranking second best. For example, DOPPLER-SSS reduces runtime by up to 78.2% over CRITICAL PATH,
8

---

Published as a conference paper at ICLR 2026
62.5% over PLACETO, and 52.7% over GDP, while both DOPPLER- SYS and DOPPLER- SIM surpass ENUMERATIVEOPPIIIZEE by up to 13.8% (see ablation studies on random seeds in Appendix G.2).
Ablation study for select and place policies (Q2). Table 3 presents an ablation study isolating the effects of the node and device policies. In the ablated variants, we replace our policies with assigns selected nodes to the earliest-available device,
CRITICAL PATH strategies: DOPPLER-SS while DOPPLER-PLO selects nodes with the longest path to an exit. Overall, combining both policies yields the best performance. For CHAINMM, DOPPLER-PLC slightly outperforms DOPPLER-SYS by a few milliseconds, but for more complex models, the combined policies provide clear gains.
Improving training using Stage I, II, and III (Q3). Fig. 4 shows per- formance for DOPPLER-SYS when trained using different combinations of imitation learning (D), simula- tions (II), and real-system execu-
iiitationSimullation-eee System Execution Time Comparison
500
Ok-Ok-Bk
4k-4k 8k
400
0k-St-8k
2k-OK-Bk
Execution Time
2k-4K-8k
2k-8k-8k
tions (III) for the LLAMA-LAYER
300
dataflow graph. As we hypothe- sized, training on the real-system only leads to slower convergence due to the need to explore from a poor initial model, resulting in un- stable performance. Imitation learn- ing and simulations enable faster convergence and lower execution times. See ablation studies on pre- training PLACETO in Appendix G.4.
200
1000 2000
3000
4000
5000
6000
7000
8000
Episode ID
Figure 4: Real engine execution times (in milliseconds) for DOPPLER-SYS using different combinations of three training stages for the LLAMMA-LAYE LAYER dataflow graph.
Visualizing DOPPLER-SYS's as- signments (Q4). Fig. 5 shows the assignments produced by DOPPLER-SYS for the FFNN dataflow graph, which achieves both GPU load balancing and communication minimization along the critical path. By further profiling how the assignments are scheduled in the system (see Appendix C), we find that DOPPLER's assignments often enable overlapping communication and computation across GPUs, minimizing stalling and maximizing utilization. We provide additional studies in Appendix A, including detailed assignment visualizations, GPU profiling results, and analyses of the node-selection order generated by SEL
Figure 5: Assignments for FFNN produced by DOPPLER. Colors indicate the mapping of computa- tions to GPUs DOPPLER effectively balances load while minimizing inter-GPU communication.
DOPPLER'S transfer ability across graphs and architectures (Q5). We evaluate transfer learning (1) from simple graph (FFNN, CIAINMMM to Llama-structured graphs on the same hardware in Table 4 and (2) across hardware architectures for the same graph. With 2K fine-tuning episodes, DOPPLER adapts to new architectures and outperforms baselines. With 4K episodes (less than half of the original training), it finds assignments comparable to full target training. In our hardware adaptability experiments, we train a policy for the FFNN graph with four P100 GPUs and transfer to eight V100 GPUs The zero-shot setting yields 82.7% of communication the same GPU group (with all-to-all NVLink), and 10.6% across GPUs without direct NVLink. After 2K episodes, the policy improves assignments to 94.7% (t 12.0%) intra-GPU, 1.9% (4 4.8%) within NVLink groups, and only 3.4% (4 7.2%) across GPUs without NVLink (detailed results in Appendix J). which justifies performance improvement.
GPU, 6.7% within
DOPPLER training and inference scalability (Q6). We analyze DOPPLER's scalability in both training and inference time for dataflow graphs with increasing size in Fig. 6. The figure shows that

---

Published as a conference paper at ICLR 2026
DOPPLER scales linearly with the size of graphs and achieves the lowest training and inference times, compared to RL-based baselines such as GDP. In terms of performance for much larger dataflow graphs, they are not substantially more complex than our experiment setups, modern large-scale training and inference rely heavily on data parallelism (Li et al. 2020b) and pipelining (Narayanan et al. 202T). More discussions are provided in Appendix
7 RELATED WORKS
DOPPLER (Inference)
17.5
DOPPLER (Policy Update)
Time Per Episode (seconds)
Placeto (Inference)
Classical Approaches for Device Place- ment and Scheduling. List scheduling (LSS) heuristics, such CRITICAL PATH, decompose the problem of computing schedule into a sequence of select and place steps (Kwok & Ahmad 1999). DOPPLER can be seen a neural LS heuristic that learns to select and place directly from observations using an MDP. Our experiments show that DOPPLER outperforms CRITICAL PATH. Graph par- titioning (Kernighan & Lin 1970; Kirk-
15.0
Placeto (Policy Update)
GDP (inference)
12.5
GDP (Policy Update)
10.0
7.5
5.0
2.5
100
150
200
250
300
350
400
# of Nodes in the Computation Graph
patrick et al. 1983] Fiduccia & Matthey-
ses 1988 Johnson al. 1989 |Hagen &
Figure 6: Inference time and RL policy update time as the number of nodes in the dataflow graph increases.
Kahng 1992, Karypis) 1997) can also be
applied device placement, but previous
work has shown that RL is a better alternative for this problem MMihhoseini et al. 2017
Reinforcement Learning for Combinatorial Optimization. Traditional algorithms for combinato- rial optimization problems often rely on hand-crafted heuristics that involve sequentially constructing a solution. Recently, there has been growing interest in applying RL (and deep learning more broadly) to learn heuristics for these problems. RL can learn constructive sequential heuristics for large combi- natorial problems and exploit structural similarity between related instances, enabling generalization (Mazyavkina et al.) (0221). For instance, (Bello al.) 2017) introduced policy gradient method for the Traveling Salesman Problem (TSP). Subsequent studies extended RL to problems beyond TSP
(Khalil al. 2017 Cappart et al. 2019: Drori al. 2020 Emami & Ranka 2018 Lu al. 2019
Mazyavkina kina et al.) N Nazari et al. [2018] Kool et al. |019 Abe et al. 2011 Manchanda et al.
2018 Gu & Yang [2020] Cai et al. (019). 2020 Chen & Tian 2019 Li et al. 2020a]] Laterre et al.
Our work is unique in leveraging combinatorial structure, list-ccheduling heuristics, and direct access to the target system during training address the device assignment problem using RL.
RL for Device Placement and Scheduling, Early work introduced sequence-to-sequenee RNN trained with policy gradients for device placement, showing that RL can outperform heuristics (Mirhoseini et af.] 2017; Pellegrini (207)). PLACETO (Addanki et al.) (2019) replaced the RNN with a GNN, while Paliwal al. (2019) combined RL with genetic algorithm. though high evaluation cost. Zhou et al. (2019) proposed a graph-embedding approach with sequential attention and single placement policy. DOPPLER introduces dual-policy learning with three-stage training for faster convergence and continuous optimization. More recent work includes end-to-end optimization from graph construction to placement (Duan et al. 2024) and improved node representations using cosine phase position embeddings (Han al. 2024 which are complementary to our approach.
CONCLUSION
In this paper, we have considered the problem of assigning computations in dataflow graph to devices to minimize execution time in a work-conserving system. We have proposed DOPPLER, dual-policy learning framework for learning device assignment in three stages. Some of the key innovations are (1) DOPPLER explicitly tries to learn an approximate node traversing order, to make the assignment problem easier, (2) DOPPLER adopts two pre-training stages using imitation learning and simulation basse learning to speed up policy convergence, and (3) DOPPLER continues dual-policy training during deployment, achieving gradual reduction in execution time time.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We provide details of experimental settings and hyperparameters used for training in Section 6.1 to reproduce our results and experiments. We have included our code and data used for running all experiments at https:/gghhbbbommiiyyyyyy//ppppp
ACKNOWLEDGMENT
We acknowledge the support by the US Department of Transportation Tier-1 University Transportation Center (UTC) Transportation Cybersecurity Center for Advanced Research and Education (CYBER- CARE) (Grant No. 6943552348332), by NSF grants 1918651, 2008240, 2131294, and 2212557, by the NIH under CTSA #UMITR004906 and by the Rice Ken Kennedy Institute.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. (TensorFlow): system for (Large-Scale) machine learning. In 12th USENIX symposium on operating systems design and implementation (OSDI 16), pp. 265-283, 2016.
Kenshin Abe, Zijian Xu, Issei Sato, and Masashi Sugiyama. Solving np-hard problems on graphs with extended alphago zero. arXiv preprint arXiv: /905. /6623, 2019.
Ravichandra Addanki, Shaileshh Bojja Venkatakrishnan, Shreyan Gupta, Hongzi Mao, and Mo- hammad Alizadeh. Placeto: Learning generalizable device placement algorithms for distributed machine learning. Advances in Neural Information Processing Systems, 32, 2019.
Ramesh C Agarwal, Susanne M Balle, Fred G Gustavson, Mahesh Joshi, and Prasad Palkar. A three- dimensional approach to parallel matrix multiplication. IBM Journal of Research and Development 39(5):575-588 582,, 1995.
Irwan Bello, Hieu Pham, Quoc Le, Mohammad Norouzi, and Samy Bengio. Neural combinatorial optimization with reinforcement learning. In International Conference on Learning Representations (ICLR), 2017.
Daniel Bourgeois, Zhimin Ding, Dimitrije Jankov, Jichui Li, Sleem Mahmoud Abdelghafar. Yuxin Tang. Jiawen Yao, Xinyu Yao, and Chris Jermaine. Eindecomp: Decomposition of declaratively- specified machine learning and numerical computations for parallel execution. Proc. VLDB Endow. 18(7):2240-2D) 2253, 2025.
Qingpeng Cai, Will Hang, Azalia Mirhoseini, George Tucker, Jingtao Wang, and Wei Wei. Rein- forcement learning driven heuristic optimization. arXiv preprint /rX1v:1906.06639, 2019.
Quentin Cappart, Emmanuel Goutierre, David Bergman, and Louis-Martin Rousseau. Improving optimization bounds using machine learning: Decision diagrams meet deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 1443-1451, 2019,
Nicolo Cesa-Bianchi and Gâbor Lugosi. Combinatorial bandits. Journal of Computer and System Sciences, 78(5): 100-1122, 2012.
Wei Chen, Yajun Wang, and Yang Yuan. Combinatorial multi-armed bandit: General framework and applications. In International conference on machine learning. pp. 151-159. PMLR, 2013.
Xinyun Chen and Yuandong Tian. Learning to perform local rewriting for combinatorial optimization. Advances in neural information processing systems, 32, 2019.
Zhimin Ding, Jiawen Yao, Brianna Barrow, Tania Lorido Botran, Christopher Jermaine, Yuxin Tang, Jiehui Li, Xinyu Yao, Sleem Mahmoud Abdelghafar, and Daniel Bourgeois. Turnip: A" nondeterministic" gpu runtime with epu offload. arXiv preprint arXiv:2405. 16283, 2024.
11

---

Published as a conference paper at ICLR 2026
Iddo Drori, Anant Kharkar, William R Sickinger, Brandon Kates, Qiang Ma, Suwen Ge, Eden Dolev, Brenda Dietrich, David Williamson, and Madeleine Udell. Learning to solve combinatorial optimization problems real- world graphs in linear time. In 2020 19th IEEE International Conference on Machine Learning and Applications (ICMIA), pp. 19-24. IEFE, 2020.
Shukai Duan, Heng Ping, Nikos Kanakaris, Xiongye Xiao, Panagiotis Kyriakis, Nesreen K Ahmed, Peiyu Zhang, Guixiang Ma, Mihai Capotă, Shahin Nazarian, et al. structure-await framework for learning device placements on computation graphs Advances in Neural Information Processing Systems, 1777-81772 2024.
Patrick Emami and Sanjay Ranka. Learning permutations with sinkhorn policy gradient. arXiv preprint (005.08800000000 2018.
Charles M Fiduecia and Robert M Mattheyses. A linear-time heuristic for improving network partitions. In Papers on Twenty-five years of electronic design automation, pp. 241-247. 1988.
Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. Systems for Machine Learning, 4(9), 2018.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp.
1263-1272. PMLR, 2017.
Shenshen Gu and Yue Yang. A deep learning algorithm for the max-cut problem based on pointer network structure with supervised learning and reinforcement learning strategies. Mathematics, 8 (2):298, 2020.
Lars Hagen and Andrew Kahng. New spectral methods for ratio cut partitioning and clustering. IEEE Iransactions on computer-aided design of integrated circuits and systems, 1009) 1992.
Meng Han, Yan Zeng, Jilin Zhang, Yongjian Ren, Meiting Xue, and Mingyao Zhou. A novel device placement approach based on position-awee aware subgraph neural networks. Neurocomputing, 582: 127501, 2024.
Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Ganger, and Phil Gibbons. Pipedream: Fast and efficient pipeline parallel dnn training. arXiv preprint arXiv:1806.03377 03377, 2018.
Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
David S Johnson, Cecilia R Aragon, Lyle A McGeoch, and Catherine Schevon. Optimization by simulated annealing: An experimental evaluation; part i, graph partitioning. Operations research, 37(6))865-897) 1989.
George Karypis. Metis: Unstructured graph partitioning and sparse matrix ordering system. Technical report, 1997.
Brian W Keriigha and Shen Lin. An efficient heuristic procedure for partitioning graphs. The Bell system technical journal, 49(2):291-300 1970.
Elias Khalil, Hanjun Dai, Yuyu Zhang, Bistra Dilkina, and Le Song. Learning combinatorial optimization algorithms over graphs. Advances in neural information processing systems, 30, 2017.
Scott Kirkpatrick, C Daniel Gelatt Jr, and Mario P Vecchi. Optimization by simulated annealing. Science, (20(4598))671-663) 1983.
Leonard Kleinrock. A conservation law for wide class of queueing disciplines. Naval Research Logistics Quarterly, 1812911 1965.
Wouter Kool, Herke van Hoof, and Max Welling. Attention, learn to solve routing problems! In International Conference Learning Representations (ICI.), 2019.
12

---

Published as a conference paper at ICLR 2026
Yu- Yu-Kwong Kwok and Ishfaq Ahmad. Static scheduling algorithms for allocating directed task graphs to multiprocessors. ACM Computing Surveys (CSUR), 31(4):406 471, 1999.
Laterre, Fu, MK Jabri, AS Cohen, D Kas, K Hajjar, TS Dahl, A Kerkeni, and Beguir. Ranked reward: Enabling self-play reinforcement learning for combinatorial optimization. Advances in Neural Information Processing Systems, 2018.
D Li, C Ren, Z Gu. Y Wang, and FCM Lau. Solving packing problems by conditional query learning. In International Conference on Learning Representations, 2020a.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. PyTorch distributed: Experiences on accelcrating data parallel training. Proceedings of the VLDB Endowment, (312))3005530018, 2020b.
Shigang Li, Tal Ben-Nun, Giorgi Nadiradze, Salvatore Di Girolamo, Nikoli Dryden, Dan Alistarh, and Torsten Hoefler. Breaking (global) barriers in parallel stochastic optimization with wvit-avoiding group averaging. IEEE Transactions on Parallel and Distributed Systems, 32(7):1725-1739, 2020c.
Xutong Liu, Siwei Wang, Jinhang Zuo. Han Zhong, Xuchuang Wang, Zhiyong Wang, Shuai Li, Mohammad Hajiesmaili, John CS Lui, and Wei Chen. Combinatorial multivariant multi-armed bandits with applications to episodic reinforcement learning and beyond. In Forty-first International Conference on Machine Learning, 2024.
Hao Lu, Xingwen Zhang, and Shuang Yang. A learning-based iterative method for solving vehicle routing problems. In International Conference on Learning Representations, 2019.
Wenyan Lu, Guihai Yan, Jiajun Li, Shijun Gong, Yinhe Han, and Xiaowei Li. Flexflow: A flexible dataflow accelerator architecture for convolutional neural networks. In 2017 IEEE international symposium on high performance computer architecture (HPCA), pp. 553-564. IEEE, 2017.
Sahil Manchanda, Akash Mittal, Anuj Dhawan, Sourav Medya, Sayan Ranu, and Ambuj Singh. Gcomb: Learning budget-constrainee combinatorial algorithms over billion-sized graphs. Advances in Neural Information Processing Systems, 33:20000-2000 -00011 2020.
Stefano Markidis, Steven Wei Der Chien, Erwin Laure, Ivy Bo Peng, and Jeffrey S Vetter. Nvidia tensor core programmability, performance & precision. In 2018 IEEE international parallel and distributed processing symposium workshops (IPDPSW), pp. 522-531. IEEE, 2018.
Nina Mazyavkina, Sergey Sviridov, Sergei Ivanov, and Evgeny Burnaev. Reinforcement learning for combinatorial optimization: A survey. Computers & Operations Research, 134:10540M 2021.
Azalia Mirhoseini, Hicu Pham, Quoc V Le, Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, Naveen Kumar, Mohammad Norouzi, Samy Bengio, and Jeff Dean. Device placement optimization with reinforcement learning. In International conference on machine learning, pp. 2430-2439. PMLR, 2017.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. Efficient large-scale language model training on gpu clusters using megatron-Im. Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-15, 2021.
Mohammadreza Nazari, Afshin Oroojlooy, Lawrence Snyder, and Martin Takác. Reinforcement learning for solving the vehicle routing problem. Advances in neural information processing systems, 31, 2018.
Aditya Paliwal, Felix Gimeno, Vinod Nair, Yujia Li, Miles Lubin, Pushmeet Kohli, and Oriol Vinyals. Reinforced genetic algorithm learning for optimizing computation graphs arXiv preprint 1rXiv:1905.02494. 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, segghrerformence deep learning library. Advances in neural information processing systems, 32, 2019.
13

---

Published as a conference paper at ICLR 2026
François Pellegrini. A parallelisable multi-level banded diffusion scheme for computing balanced partitions with smooth boundaries. In Euro-Par 2007 Parallel Processing: 13th International Euro-Par Conference, Rennes, France, August 28-31, 2007. Proceedings 13, pp. 195-204. Springer, 2007.
Sujan Kumar Saha, Yecheng Xiang, and Hyoseung Kim. Stgm: Spatio-temporal gpu management for real-time tasks. In 2019 IEEE 25th International Conference on Embedded and Real-Time Computing Systems and Applications (RTCSA). pp. 1-6. IEEE, 2019.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference Learning Representations (ICLR), 2017.
Shaohuai Shi, Xinglin Pan, Qiang Wang, Chengjian Liu, Xiaozhe Ren, Zhongzhe Hu, Yu Yang, Bo Li, and Xiaowen Chu. Schemoe: An extensible mixture-of-expert distributed training system with tasks scheduling. In Proceedings of the Nineteenth European Conference Computer Systems, pp. 236- 249,2499922022 2024.
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Leslie G Valiant. A bridging model for parallel computation. Communications of the ACM, 33(8): 103-111, 1990.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Cliff Woolley. Neel: Accelerated multi-gpu collective communications, 2015.
Vignesh Yaadav. Exploring and building the llama 3 architecture: A deep dive into components, coding, and inference techniques, 2024.
Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric Xing, et al. Alpa: Automating inter-and {Intra- Operator} parallelism for distributed deep learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pp. 559-578, 2022.
Yanqi Zhou, Sudip Roy, Amirali Abdolrashidi, Daniel Wong, Peter C Ma, Qiumin Xu, Ming Zhong, Hanxiao Liu, Anna Goldie, Azalia Mirhoseini, et al. Gdp: Generalized device placement for dataflow graphs. arXiv preprint 1rx.v.1110001578, 2019.
14

---

Published as a conference paper at ICLR 2026
We provide brief roadmap of the appendix below:
Appendix A: Device assignment analysis with system profiling and visualizations in graphs across methods.
Appendix B: Enumerative Assignment Algorithm (Algorithm 4) for our baseline ENUMER- ATIVE OPTIMIZER
Appendix C: Implementation of the system we used in Stage III training and for evaluat- ing execution time for assignments reported in tables.
Appendix D: Detailed configurations for computation graphs in experimental setups.
Appendix E: Detailed Graph Neural Network features used in dual-policy networks.
Appendix F: Configurations for synchronous system runs in Table 1.
Appendix G: More ablation studies simulator (Stage training and DOPPLER-SIM), 2) reinforcement learning trainings across random seeds, 3) efficiency performance trade-off for message passing on Doppler, 4) pretrain placeto with Stage
Appendix H: Experiments on various hardware configurations, including various GPU memory sizes and different numbers and types of GPUs.
Appendix I: Discussion on scaling to much larger dataflow graphs and larger numbers of GPUs.
Appendix J: Doppler's transfer ability across hardware with analysis.
15

---

Published as a conference paper at ICLR 2026
A
MORE DEVICE ASSIGNMENT ANALYSIS AND VISUALIZATIONS
A.1
COMPUTATION NODE DETAILS
input: input tensors
matmul: matrix multiplications on two matrices
input elemwise: elementwise operations (e.g., RcLU) on an input tensor
straight elemwise: elementwise operations (e.g., ReLU) with two inputs having the same dimensions
beast elemwise: takes two inputs of different shapes (e.g.. a matrix and vector), performing elementwise operation (e.g.. ReLU) with one element of the vector applied to entire row of the matrix
max reduction: reduce one dimension by finding the maximum.
min reduction: reduce one dimension by finding the minimum.
sum reduction: reduce one dimension by finding the sum along that dimension.
product reduction: reduce dimension by finding the product along that dimension.
formation: placeholder operation that forces aggregations in meta-ops groups (defined in Appendix B) to form single tensor.
complexer: a conversion between floating-poin and complex tensors.
fill: an operation to create tensors with all elements set to the same scalar, or to assign all
lower upper diagonal elements with provided scalar values.
squeezer: adding or removing singleton dimensions of tensor.
selec: an operation to copy subset of several input tensors into an output tensor- generalization tensor subset and tensor concatenation.
A.2 ASSIGNMENT PROFILING
This section examines the performance results from Table 2, where the DOPPLER algorithm achieved a lower runtime compared to CRITICAL PATH, PLACETO, GDP, and the expert-designed ENUMERA- TIVEOPTIMIZER.
During the development of the ENUMERATIVEOFTIIIZER algorithm, it became clear that for meta- ops (described in Appendix B) containing many computations, the devices should be fully utilized and load-balanced. Put more succinctly, it is expected that good assignments should minimize data transfers while maximizing computational resource utilization. In the following, provide detailed analysis into CHAINMM and FFNN throught assignment visualizations and system profiling.
A.2.1 CHAINMM
The assignments in Fig. 7 for DOPPLER show all four devices (four colors of nodes) being used, whereas for ENUMERATIVEOPTIMIZE computations. The corresponding device utilization plots are shown Fig. 9 and Fig. 10, respectively. It appears that, indeed, ENUMERATIVEOPTIN toward the end of the computation (after 80ms). In contrast, DOPPLER performs well from the beginning to the end, shown in Fig. 7.
in Fig. 8, only two of the devices are used for the latter
does not fully utilize available compute resources
Jdbnn slad
Figure 7: Assignment found by DOPPLER for ChainMM
16

---

Published as a conference paper at ICLR 2026
Figure 8: Assignment found by ENUMERATIIVEPPTIIIZE for ChainMM
Kernel
G
Memory
Kerne
GPU
Memory
Kernel
GPU
Memory
Kemel
GPU 3
Memory
20
40
60
80
100
Time (ms)
Figure 9: Device and transfer utilization for DOPPLER, CHAINMM.
Kerel
GPU
Memory
Kerne
GPU 1
Memory
Kernel
GPU
Memory
Kernel
GPU
Memory
20
40
60
80
100
120
Time (ms)
Figure 10: Device and transfer utilization for ENUMERATIVEOPTMIEEE, CHAINMM.
17

---

Published as a conference paper at ICLR 2026
A.2.2 FFNN
The assignments for FENN with DOPPLER PLACETO are shown Fig. and Fig. 12; corresponding device utilization plots are shown in Fig. 13 and Fig. 14. In the DOPPLER assignments, subsequent vertices typically share the same device assignment. In turn, this leads to lower amount of data transfer. For the PLACETO assignments, however, subsequent vertices do not tend to have the same device assignment, possibly leading to a large amount of data transfer. the device utilization figures, this is indeed boree out: the PLACETO execution is three times slower, with most of the time spent on data transfers across GPUs.
Figure 11: Assignments found by DOPPLER for FFNN.
Figure 12: Assignments found by PLACETO for FFNN.
Kerne
GPU
Memory
Kernel
GPU 1
Memory
Kemel
GPU
Memory
Kernel
GPU
Memory
10
20
30
40
Time (ms)
Figure 13: Device and transfer utilization for DOPPLER, FFNN.
18

---

Published as a conference paper at ICLR 2026
Kerne
GPU 0
Memory
Kernel
GPU
Memory
Kernel
GPU 2
Memory
Kerel
3PU
Memory
20
40
60
80
100
120
Time (ms)
Figure 14: Device and transfer utilization for PLACETO, FFNN.
19

---

Published as a conference paper at ICLR 2026
A.3 LLAMA-BLOCK ASSIGNMENT ANALYSIS WITH RESPECT NODE- SEDE-SELECTION ORDER
We showed the assignments with respect to the selected node order in Fig. 15, Fig. 16, Fig. 17, Fig. 18, Fig. 19. We observed two interesting patterns that SEL, follows.
The first pattern is that SELg performs level-wise node assignment to maintain load balancing. In Fig. 15, Fig. 16, and Fig. 17, each matrix multiplication, tensor formation, and element-wise input operation is selected sequentially and distributed evenly across the four GPUs. This strategy keeps all GPUs fully utilized during this phase of execution. At this stage, the workload is computation- bound, meaning many operations are ready to execute in parallel. As result, maximizing GPU utilization by distributing computations evenly is more important than minimizing communica- tion overhead-2 trade-off that competing baselines fail to manage effectively.
The second pattern that SEL, prioritizes selecting nodes for which it has high confidence in determining optimal device placement. In Fig. 18 and Fig. 19, the system becomes network- bound, with limited parallel computations available and communication bandwidth emerging as the primary bottleneck. In this regime, it is more beneficial to select nodes along the same execution path and assign them to the same GPU to reduce cross-device communication, rather than distributing independent nodes across different GPUs.
20

---

Published as a conference paper at ICLR 2026
Figure 15: Assignment found by DOPPLER for LLAMA-BLOCK at Step 78
21

---

Published as a conference paper at ICLR 2026
Figure 16: Assignment found by DOPPLER for LLAMA-BLOCK at Step 86
22

---

Published as a conference paper at ICLR 2026
Figure 17: Assignment found by DOPPLER for LLAMA-BLOCK at Step 94
23

---

Published as a conference paper at ICLR 2026
Figure 18: Assignment found by DOPPLER for LLAMA-BLOCK at Step 146
24

---

Published as a conference paper at ICLR 2026
Figure 19: Assignment found by DOPPLER for LLAMA-BLOCK at Step 153
25

---

Published as a conference paper at ICLR 2026
A.4 LLAMA-LAYER ASSIGNMENT VISUALIZATIONS
We show the best assignment we found for five methods on the LLAMA-LAYER computation graph: DOPPLER (Fig. 20), EXUMERATIVOFTTIZZR (Fig. 21), PLACETO (Fig. 22), CRITICAL PATH (Fig. 23), and GDP (Fig. 24).
XEE veyhert agst 1000
Figure 20: Assignment found by DOPPLER for LLAMA-LAYER
26

---

Published as a conference paper at ICLR 2026
Figure 21: Assignment found by ENUMERATIVEEPTIMIEER for LLAMA-LAYER
27

---

Published as a conference paper at ICLR 2026
dieedee
demite
Kanc domenee
lice donetol
Nnstcomnoz
Koeeke
decelal
Figure 22: Assignment found by PLACETO for LLAMA-LAYER
28

---

Published as a conference paper at ICLR 2026
tace coceo
Figure 23: Assignment found by CRITICAL PATH for LLAMA-LAYER
29

---

Published as a conference paper at ICLR 2026
tace aoceo EaE
Figure 24: Assignment found by GDP for LLAMA-LAYER
30

---

Published as a conference paper at ICLR 2026
B ENUMERATIVE ASSIGNMENT ALGORITHM (ALGORITHM 4)
TL;DR. The Enumerative Assignment Algorithm is greedy, sharded-ops--y---hardeee device placement strategy that traverses the graph level-by-leve,, exhaustively trying shard assignments to minimize estimated communication cost, without global lookahead or learning,
In this section, describe our enumerative assignment algorithm, which a level-by-level, exhaustive enumeration in an attempt to find an assignment A for the vertices in a graph G that minimizes ExecTimeeAA This algorithm uses a greedy approach that first groups vertices based on the graph structure and then attempts a subset of possible assignments for the operations within each group, selecting the assignment with the minimum estimated cost.
This algorithm requires that the vertices in the graph & have been organized into list of "meta- ops" called M. our input dataflow graph has been created by sharding a compute graph using framework such as Alpa Zheng al. (2022), each operation in the input dataflow graph G is descended from some operation that has been sharded. For example, consider Fig. All eight MMul 1. ops at the lowest level of the graph, well as the four MAdd ops at the next level, were created by sharding X Y. We group all twelve of these operations and term them a "meta-op." Further, we topologically order these meta-ops that if m. comes before m2 in M, it means that none of the vertices in m2 can be reached from any vertex in mi by traversing E.
Note that each meta-op has two subsets- one of which may be empty: set of computationally expensive operations (such as the MMul ops) that result directly from sharding the original operation, and set of less expensive operations needed aggregate and/or recompose the results of the first set of operations. For a meta-op m, we call these m. rd0ps and m. reduce@ps. Note that the original operation is always sharded so that if there are n devices, there are n items in m.shardOps and load balancing of these shards is crucial. Thus, our tactic is to always partition shshhardpps across the n devices and never assign two operations in m.shardops if the meta-op is sharded into n parts, there will always be at most n items in m.reduceOps. Therefore, we always partition them across (possibly a subset of) the devices.
the same device. Likewise,
Given this, our algorithm, ENUMER ATIVEOPTIMIEERR proceeds through the list M of meta-ops in order. For each m, it exhaustively tries all assignments of m.sha rdOps. Each assignment costed by computing the time required to transfer all of the items in m.shardOps to where they will be consumed. These times are estimated using statistics gathered by testing transfers on the actual hardware. Once m. m.shardOps is placed, m.reduceOps is placed using the same cost model. Because of the ordering of the meta-ops in M, and because we process m.shar m.shardOps before m.reduceOps, we always know the assignment of the inputs to m.shardOps m.reduceOps before placing them, and thus it is easy to compute the cost.
This algorithm is greedy in the sense that it processes meta-ops one time, from start finish, using the topological ordering. Thus, if the cost model is correct, it will be optimal only as long as each m.shardOps and merreuueOpp executed. Obviously, this is not the case in reality with dynamic scheduler, but one might expect the algorithm to produce good assignment in practice.
run in lock-step, with barrier before each is
31

---

Published as a conference paper at ICLR 2026
Algorithm 4 EnumerativeOOpimizzer
Input: Sorted list of meta-ops M containing all vertices in g Output: Assignment A
A () % assignment is empty
% loop thru meta-ops
for E .M do
% First deal with the shared op
A < get BestBestAss Masssgggmmm... (gw.nnsaarddpsss.. A)
% now place the reductions
A wwt.eetssssigmmeeeeeeeee/eoe A)
end for return A
% computes the best assignment for set of vertices that are expected to in parallel all devices
subroutine get gn(vertices, A) bestCost < oo; best.n <- null
% loop through all possible device assignments
for D) E allFerms(()) do whichDev + 0; cost 0 0
% loop through the ops created by sharding this meta-op for vertices do
% loop through inputs to this op and add network cost for en E V s.t. (v1, v) E & do
cost cost Net (me dus D DwhiccDev)
end for whichDev +- whichDev + end for if cost bestCost then
bestCost + cost best.ssiin D
end if end for
% we have the best assignment for this meta op, so record it whichDev t- 0
for E vertices do
ag + -we append to A
whichDev + whichDev +
end for return
32

---

Published as a conference paper at ICLR 2026
C SYSTEM IMPLEMENTATION (DOPPLER-EINSUMMAN)
The underlying system runtime Bourgeois et al. 2 Ding et al. (024) that executes our graphs is written in C++. The dataflow graph 9 is executed asynchronously by single-threaded event loop whose job is to monitor when the dependencies implicit in the graph are satisfied. When the inputs to vertex found to be available, the event loop checks whether the resources necessary execute the vertex are also available (a "resource" may be an open GPU stream or open communication channel). If resources are available, the event loop may choose execute the vertex. We use the "event loop" because the loop requests the necessary resources to execute the vertices in G in response to "events." " An event oceurs whenever a graph dependency satisfied or when a previously used resource becomes available. The main event loop is notified of this via an asynchronous callback. In our implementation, we use CUDA version 11..00 and cuTensor version 2.0.1.
D COMPUTATION GRAPHS USED IN THE EXPERIMENTS
This section describes the computation graphs used in our experiments, which are designed to capture range of workload characteristics encountered in modern ML systems. We include (i) ChainMM, a matrix-mulipliccttinn graph with long dependency chains and parallel substructures; (ii) FFNN, feed-forward neural network modeling dense linear layers and nonlinear activations typical of deep learning workloads; and (iii) Llama- block Llama-layer, which represent transformer-styi computation graphs from large language models. Together, these graphs span diverse operator mixes, graph sizes, and dependency patterns, enabling a comprehensive evaluation DOPPLER across synthetic, neural, and large-model workloads.
D.1 CHAINMM
Inputs: Five matrices A, D.C.D.D.. C, EE
R. 10000x 10000
Computation: (A x (C
Graph size: 112 nodes.
D.2 FFNN
Inputs:
Input batch matrix: X E TR22"2" x2"
First-layer weights: W(1) R22
First-layer bias: ((1) E 2
Output-layer weights: W x22
Output-layer bias: ((2) R
Computation:
Hidden layer:
= ReLU (xw()+b0)) Hl x222
Output layer:
Y Softmax (HW(2) + +b(2). Y x22
Activations: ReLU for the hidden layer; Softmax for the output layer. Graph size: 192 nodes.
D.3 LLAMA-BLOCK AND LLAMA-LAYER
Model configuration:
Number of parameters: 7B.
Embedding dimension: 4096.
33

---

Published as a conference paper at ICLR 2026
Maximum sequence length: 4096.
Vocabulary size: 32,000 tokens.
Batch size: 1.
Number of layers: 1.
Graph size: 215 nodes.
Computation: The computation graph follows the standard Llama transformer block and layer structure, as illustrated in Fig. 25.
Linar
PVV5 Norr
SUGLU
Self Alte viee SenttreeeeeGroeeellllae (e Oselyy--erraaa Aterloo) Afterton)
XV Cache
Poutoonal Encodinge
Crbbaidings
LLaMA
Figure 25: Llama-block and Llama-layer architecture. Figure from Yaadav 2024
34

---

Published as a conference paper at ICLR 2026
E GRAPH FEATURES Xg AND DEVICE FEATURES XD
Given a computation graph 9 = (V,8) with V {ui, U2 u2,..., Un and = {61, define the following:
we
Computation cost for v, We use the floating toating--poitt operations per second of as the computation cost.
Communication cost for eijj = (vi, vy). The number of bytes of the output tensor of U; multiplied by a communication factor. In our case, to 4. We benchmark the execution time of the simulator versus the real execution engine. We tried communication factor values from 10 and found 4 to be the closest match for the simulator with respect to the real execution engine.
the communication factor equal
E.1 STATIC GRAPH FEATURES Xg
The graph features matrix Xg is an n x 5 matrix, where each row contains the following five features:
Computation cost for Uz-
Sum of communication cost from predecessor nodes to Vi. The sum of communication costs for all Gi such that (vi, Uy) E &.
Sum of communication cost from 2y to descendant nodes. The sum of communication costs for all Cj such that (Ny, 0%) E E.
t-level cost of vy. The of all computation and communication costs t-level path for Vy. A t-level path is defined in Section 4.2.
b-level cost of uz. The sum of all computation and communication costs on b-level path for Vj. A b-level path is defined in Section 4.2.
E.2 DYNAMIC DEVICE FRATURES X FOR DEVICE d AT TIME STEP GIVEN NODE V
The device features matrix Xp is a (D] x 5 matrix, where each row contains the following five features:
Total node computation cost. The sum of the computation costs for all nodes that have been assigned to device d at time step f.
Total predecessor node computation cost. The sum of computation costs for all predeces- sor nodes target node time step that are currently assigned device d.
Minimum start time of all inputs. The earliest time start executing predecessor node of u on device d at time step f.
Maximum end time all inputs. The latest time for all predecessor nodes of v to finish device d at time step t.
Earliest start time to compute v. The earliest time for device d to finish receiving inputs and start executing v.
F SYNCHRONOUS SYSTEM CONFIGURATION
Table I, we run CHAINMM using ScalaPack and FFNN using Pytorch Lightning with tensor parallel on 4 NVIDIA Tesla P100 GPUs with 16GB memory each.
35

---

Published as a conference paper at ICLR 2026
G MORE ABLATION STUDIES
We conduct additional ablation studies (1) the simulator implementation, (2) random seeds used in training the dual policy, (3) the number of message passsing rounds per episode, and (4) the inclusion of pre-training stages for PIACCETO. We aim evaluate the performance and training efficiency of our three-stage training approach with a dual-policy design compared with alternatives.
G.1 SIMULATOR ABLATION STUDIES
We compare the simulated and real system execution times for the same device assignments in the following plots. We hypothesize that the simulator serves as a cost-effective way to approximate system execution times without sacrificing significant precision.
On the left of Fig. 26, we show the running times for both the simulator and the real system when training the dual policy on the Chai nMM computation graph via imitation learning. We observe that the simulator tends to overestimate the running time compared to the real system and has some difficulty differentiating between assignments with similar running times. However, the simulator provides approximate running times that follow the same trend as the real system (e.g.. high-quality assignments tend have shorter running times, and low-quality ones exhibit longer running times).
On the right of Fig. 26, we present a statistical analysis comparing the performance of the simulator with the real system under the same setup. We find Pearson coefficient 0.79 and Spearman coefficient of 0.69. This ablation study demonstrates that the simulator offers cost-effective way to approximate system execution times.
Real System vs. Simulator Running Time in ChainMM Training
Real System vs. Simulacor Running Time (Pearson: 0.79, Spearman: 0.69)
Simulation Rumning Time
Real System Running Time
Data
Simulatoon Humniro time (millted)
(Pearson)
Running time (mil sec
Linear Fit (Speerman)
20e
120
Episode ID
Roal System Running Time (millisoc)
Figure 26: (left) A line chart showing a comparison between the simulator running time and the real system rumning time throughout training. (right) A scatter plot showing the simulator running time versus the real system running time with Pearson and Spearman fitting lines.
G.2 EXPERIMENT WITH RANDOM SEEDS
In this study, we aim to test the hypothesis that the best assignment found by our approach, DOPPLER, remains consistent across different random seeds during training. Due to the cost of training, we are unable to run experiments multiple times with different random seeds for all computational graphs and methods. Therefore, we conduct five training runs of DOPPLER-SYS with different seeds on the Cha 1nMM computation graph to test this hypothesis. This experimental setup-including including the computation graph, dual dual-policy architectures, and training hyperparameters is identical across runs, differing only in the random seeds. For each training run, evaluate the best-found assignment over 10 system executions and report the mean and standard deviation Table 5. The results show that DOPPLER achieves consistent performance across different random seeds.
36

---

Published as a conference paper at ICLR 2026
Model
Runl
Run2
Run3
Run4
Run5
ChainMM
123.2 - 3.7
119.6 + 2.2 122.7 + 2.1
123.9 + 2.5
121.7 + 0.9
Table 5: Experiment running DOPPLER using different random seeds on CHAINMM computation graph. We test the best assignment found at the end of the training across different seeds with 10 system runs and report the mean and standard deviation of the system running time (in milliseconds) for each assignment.
G.3 MESSAGE-PASSING PASSING ABLATION STUDIES
To enhance training efficiency for large computation graphs, we apply message passing on the graph once per MDP episode instead of once per MDP step. When message passing is applied once per MDP step, the number of pressage-passingg rounds per episode equals the number of nodes in the graph. We hypothesize that this modification has negligible impact on DOPPLER's convergence but significantly reduces training time-proportionall to the number of nodes the computation graph. We conduct this ablation study on the Cha nMM computation graph using the simulator to save time, since per-step message passing incurs prohibitively high training costs.
DOPPLER-SIN
DOPPLER SIw---In--ee------ee----
Best assignment
122.5 + 4.0
121.7 = 3.2
Best assignment @ Number of episodes
3425
963
Best assignment @ Number of message passing
3425
107,856
Run time reduction
0.7%
Extra message-passing
3049. %
Table Running time (in milliseconds) the best device assignment found at the end, along with the number of message-passingg passing steps conducted until finding the best assignment. The reported time for the best assignment includes both the average and standard deviation the system running time over 10 system rounds. DOPPLER-SIM refers to performing message passing the computation graph per MDP episode, while per MDP step within each episode.
mpnn--er-steep denotes conducting wessaag-passing
In each MDP step within an episode, we assign a device to the currently selected node. Therefore, SAARR------mm--pe----e-e
for the equals the number of nodes in the computation graph. The ChainMM computation graph consists of 112 nodes. Table 6 shows that the best assignmentS-reported the first row by their running times- were found at episode 3425 for DOPPLER-SIM and at episode 963 for Although the best assignment was found in fewer episodes with DOPPLER SEERR-SI--Mma SIM-mpnn- peonn-per-stpp.. completing 963 episodes took significantly more wall-clock time than 3425 episodes for DOPPLER- SIM, because peessage-passing took the majority time during training.
approach, the number of meessagepasssngg rounds per episode
SMPPLER-SIM-Mpn--e
We evaluate the efficiency of the two approaches based on the number of message-passinn passsig rounds required to find the best assignment. DOPPLER-SIM, performs 3425 message- per episode), while DOPPLER-SIM-mpnne message-passing operations (one per MDP step, or equivalently, per node in the computation graph). DOPPLER-SIM-mih S----------------------- ser-snpperteppppeeee achieves a 0.7% reduction in runtime for the best assignment compared to DOPPLER-SDD SIM, but at the cost of 3049.19 % more message- paassagg..ssing..... Therefore, these results support hypothesis that the modified approach has greatly reduced the training time with negligible impact on the performance.
operations (one conducts 963 episodes x 112 nodes = 107,858
G.4 PLACETO ABLATION STUDIES
We conduct an ablation study on policy design for learning device assignments in both DOPPLER and PLACETO, explicitly including pre-training stages for PLACETO to isolate the effect of the
37

---

Published as a conference paper at ICLR 2026
underlying policy design from the benefits of pre-training. We hypothesize that the dual-policy design DOPPLER outperforms PLACETO regardless of the inclusion of training stages. To test our hypothesis, pre-train the PLACETO policy using imitation learning and compare it with DOPPLER-SIM. which is trained using the imitation learning stage (Stage I) and the simulation-base RL stage (Stage II).
PLACETO-pretrait 99.0 + 5.7
PLACETO 126.3 + 5.8
DOPPLER-SID 49.9 1.1
DOPPLER-SYS 47.4 0.7
Best Assignment
Table 7: The mean and standard deviation of the running time (in milliseconds) for the best assignment found for DOPPLER, compared to PLACETO and its pre-training version, PLACETO-pretrain, over 10 system runs. The results indicate that even with the pre-training stage, PLACETO-pretrait performs worse than DOPPLER-SIM,
Table 7 shows that DOPPLER discovers more effective device assignments than PLACETO. This result isolates the impact of the training stages and supports the claim that DOPPLER'S 's dual-policy design outperforms the policy design in PLACETO.
H DOPPLER'S S EXPERIMENTS ON DIFFERENT HARDWARE CONFIGURATIONS
The experiments far demonstrate that DOPPLER outperforms the alternatives four Tesla P100 GPUs, each with 16GB of memory. We hypothesize that DOPPLER can find better device assignments than alternatives across different hardware configurations. In this section, we conduct experiments with DOPPLER and other methods under (I) varying GPU memory sizes on four P100 GPUs and (2) different numbers and types of GPUs. Each experimental setup described in the following two subsections.
H.1 EXPERIMENTS WITH RESTRICTED GPU MEMORY
We aim to test the hypothesis that DOPPLER can adapt to hardware setups with restricted GPU memory. Table shows results four NVIDIA P100 GPUs, each restricted use 8GB out of their 16GB total memory. DOPPLER learns to adapt to memory constraints, achieving to 49.6% and 18.6% runtime reductions compared to the best baseline and ENUMERATIVEOPTIIEER respectively, while heuristics fail to adapt due to dynamic memory allocations in WC systems. These results confirm that DOPPLER can adapt to restricted memory settings.
GPUs
RUNTIME REDUCTION
MUDEL
1 GPU
CRITICaL PATH
PLACETO
ENUMÖPT:
DOPPLER-SYS
BASELINE
Exu MOP.
CRAINMM
439.8 2: 4.6
310 4.9
243.5 4: 5.9
133.5 + 10.4
122.6 2.2
49.6%
8.2%
FFNN
148.2 19.1
225.8 : 19.1
126.8 5.7
19.2 # 0.9
46.0 = 0.8
63.7%
6.5%
LLAMA-BLOCK
463. 4.88
216.8 + 4.6
433.5 + 6.2
233.8 +: 8.1
190.2 = 11.2
12.3%
18.6%
LLAMA-LAYER
462.6 6... 9.4
292.5 +: 5.1
302.1 + 20.2
172.8 d: 4.3
154.0 + 3.7
47.4%
10.9%
Table 8: Real engine execution times (in milliseconds) for assignments identified by our approaches (EnumerativeOptimizze Doppler-SYS) using 4 GPUs with access 8G out of 16G GPU memory for each GPU compared against those produced using GPU and by two baselines (CRITICAL PATH and PLACETO). The results show that Doppler-SYS outperforms all baselines across all settings.
H.2 EXPERIMENTS ON DIFFERENT NUMBERS AND TYPES OF GPUs
We aim to test the hypothesis that DOPPLER can find better device assignments regardless of the number and type of GPUs in the server. Table presents results from running various computation graph architectures
eight NVIDIA V100 GPUs, each with 32GB of memory. The setup consists of two device meshes, each fully interconnected via NVLink, with a total of four additional NVLinks spanning between the meshes. We observed that DOPPLER effectively leverages NVLink minimize
38

---

Published as a conference paper at ICLR 2026
mnter-mesh data transfer. We show that DOPPLER achieves up 67.7% runtime reduction compared to the existing baseline and up to 19.3% compared to
GPUs
RUNTIME REDUCTION
MODEL
1 GPU
CRITICAL PATH
ENUMOPT.
DOPPLER-SYS
BASELINE
ENUMÖPT.
CHAINMM
83.5
69.6
33.5 2.5
32.1
53.9%
4.1%
FFNN
51.4 1.8
50.0 =: 6.0
2..
16.2 = 2.3
67.7%
19.3%
LLAMA-BLOCK
154.4 + 6.3
6.0
109.6 + 4.2
109.7 3.0
6.7%
-0.1%
LLAMA-LAYER
105.0 4.8
105.4
97.5 1.1
90.6 =
13.7%
7.1%
Table 9: Real engine execution times (in milliseconds) for assignments identified by our approaches (EnumerativeOpimizze, Doppler-SYS) using 8 V100 GPUs compared against those produced using 1 GPU and by one baseline (CRITICAL PATH). The results show that Doppler-SY3 outperforms the alternatives in 3 out of 4 settings.
I DISCUSSION ON SCALING TO MUCH LARGER DATAFLOW GRAPHS
With respect to much larger graphs, we expect the linear scale-up to continue, but in practice, massive graphs are unlikely to be an issue. While we cannot know what companies such as OpenAl are doing, in practice, inference is likely not run more than few dozen GPUs (graphs grow linearly size with increasing GPU counts). Training, while requires thousands of GPUs, is performed by relatively independent machines through data parallelism and pipelining each machine gets an identical transformer or MoE layer and subset of the data so it runs the same computation as all of the other machines. In this case, one would likely train single massive graph. Rather, each repeated block or layer would be used to train a separate dual policy in parallel and be given the same assignment each machine with repeated structure throughout the cluster (assuming uniform hardware), with the runtimes collected across the cluster used to compute reward.
J DOPPLER'S TRANSFER ABILITY ACROSS HARDWARE AND DETAILED ANALYSIS
We conducted transfer learning study in which DOPPLER is trained on computation graph running on four P100 GPUs with full NVLink connectivity (i.e., every GPU is directly connected to the others), and then evaluated on system with eight V100 GPUs that have only partial NVLLink connectivity. In the eight-GPU system, the V100 GPUs are organized into two groups of four (GPU 0-3 and GPU 4-7). Within each group, the GPUs are fully connected via NVLink. However, connectivity between the two groups is limited: only four NVLink connections link the groups in total. This creates a hierarchical topology with high-bandwidth communication within each 4-GPU group and more restricted bandwidth across groups
In the assignment found by DOPPLER, measured how the number of data transfers-and and the ratios shown in parentheses- across the two 4-GPU groups, within the same 4-GPU group, and within the same single GPU change from zero-shot 2h-shot for FFNN eight GPUs, reported in Table 10.
across
Table 11 shows the real engine execution times (in milliseconds) for the assignment identified by DOPPLER under different few-shot settings (Zero-shot, 2K-shot) compared against the baseline for the transfer learning study above. After 2K episodes, DOPPLER finds better assignments for both ChainMM and FFNN compared training solely eight GPUs without generalization (8K episodes) for runtime reduction of 19.0% (ChainMM) and 11.1% % (FFNN).
ACROSS GROUPS
SAME GROUP
SAME GPU
ZERO-SHOT 2K EPISODES
22 (10.6%)
14 (6.7%)
172 (82.7 %)
7 (3.4%)
4 (1.9%)
197 (94.7 %)
Table 10: Data transfers across two 4-GPU groups, within the same 4-GPU group, and within the same single GPU changes from Zero-shot to 2K-shot for FFNN eight GPUs.
39

---

Published as a conference paper at ICLR 2026
COMPUTE
HARDWARE
HARDWARE
ZERD-S0OOT
DOPPLER-SYS
Opt.
P100
V1OD
$1.2
(3.5
FENN
4 P100 GPUS
V GPUS
23.1
(4.4
(2.3)
(6.0)
20.0
Table 11: Real engine execution times (in milliseconds) for the assignment identified by DOPPLER under different few-shot settings (Zero-shot, 2K-shot) compared against the baseline.
40


Reviews:
[{"review_id": "dnOOuWGY6O", "reviewer": "Reviewer_qMoz", "summary": "This study developed a reinforcement learning based approach (DOPPLER) to enable efficient assign operations in a dataflow graph. Experiments shows that the trained RL model could reduce the execution time in neural network operations with different network architectures.", "questions": "Please see weakness section.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 1, "presentation": 1, "contribution": 2, "strengths": "The problem has a clear motivation, illustrated by a dataflow graph, in the introduction section for matrix multiplication in deep learning operations.", "weaknesses": "#1 It is not easy for readers to follow the logic of this paper. The paper fails to clearly describe the problem and methods.\n\n#2 The abstract is too short to demonstrate the significance of DOPPLER. Basically, there are no quantitative results in abstract\n\n#3 “Simulator” for system execution was used through the paper without any description. Please clarify. \n\n#4 Since section 3 was titled as problem definition, is section 2 just a background? More specifically, is Algorithm 1 an existing algorithm or the problem you aim to study in this paper? No citations were provided in Section 2.\n\n#5 The paper does not provide sufficient technical details but rather forced everything to the appendix.\n\n#6 GNN is too broad. What specific architecture did you use? Clarification and citation are needed in the main text.\n\n#7 The rationale of three-stage training was not well justified. Any challenges or preliminary experiments that could support your proposed training mechanism?\n\n#8 What reinforcement learning algorithm did you use?\n\n#9 How does “RL for combinatorial optimization” in the related work section contribute to underpin the background of your work? E.g., you mentioned many works for TSP but does it help here?\n\n#10 The “combinatorial structure” in line 197 was not clearly introduced in the main text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CnroGu1k9e", "reviewer": "Reviewer_kPCd", "summary": "The authors introduce DOPPLER, a method to efficiently assign computational operations in dataflow graphs to GPU devices in work-conserving (asynchronous) systems. The key innovation is a dual-policy learning framework that separates the problem into selecting which operation to assign next and determining which device should execute it. DOPPLER employs a three-stage training approach combining imitation learning, simulation-based RL, and continuous real-system RL during deployment, to achieve significant execution time reductions over existing baselines using clever design choices to ensure both good performance and scalability.", "questions": "- How does this work relate to the substantial body of research on learning-augmented scheduling? While your setting differs (multi GPU dataflow graphs vs. traditional job scheduling), the core problem structure appears similar:\n\t- Both involve sequential assignment decisions under uncertainty\n\t- Both use predictions/learning to schedule DAGs\n\t- Both must balance load and minimize communication/movement costs\n\tThe learning augmented scheduling literature has developed concepts like consistency (performance with accurate predictions), robustness (worst case guarantees), and prediction error bounds. Could these frameworks provide theoretical grounding for DOPPLER? For example, can you characterize DOPPLER's consistency and robustness formally, or provide approximation guarantees based on simulator accuracy?\n- How does your work scale to larger graphs? Is this a realistic setting for dataflow? And does your linear scalability still hold?\n- How does this work for dynamic graphs (that is when the graph changes dynamically) - is this a possible setting for your problem?\n- Why is an RL method needed for this problem versus approaches like Bayesian Optimization or other direct optimization methods?\n- Hyperparameter selection? To deploy your algorithm in practice for heterogeneous workloads and hardware, how should one set hyperparameters?\n- Why those specific neural network architectures? This is not a criticism - rather I'm seeking to understand why you made those design choices.\n\nIf some of my more important questions are addressed, I'm very inclined to increase my score. I think this will be a very impactful addition to the community speed up compute especially for ML jobs.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Very significant practical impact in performance gains over best baselines towards the dataflow graph assignment problem.\n- Well clever design choices to enable both scalable training and improved performance \n\t- Pretraining techniques to improve performance during the actual RL phase.\n\t- Clever message passing implementation to reduce training time with very negligible performance impact\n- Strong generalization across hardware architectures", "weaknesses": "- Evaluation seems limited to relatively small graphs (~200 nodes). Real life workloads could be much bigger that could affect the linear scalability claim. In my own personal experience, scheduling larger graphs sometimes introduces additional complexity that sometimes is not captured on smaller graphs\n- Dynamic Dataflow - I could imagine that dataflow may change dynamically during execution - especially with larger computational graphs. Is this a setting your algorithm can handle because from the writing this doesn't seem to be considered.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kvhfpF2eEH", "reviewer": "Reviewer_hjpH", "summary": "### **Review Summary**\n\nThis paper proposes DOPPLER, a reinforcement learning (RL) framework for the device placement of dataflow graphs in an asynchronous, \"work-conserving\" (WC) system. The authors argue that existing bulk-synchronous systems under-utilise resources and that prior RL solutions use a single, inflexible placement policy. DOPPLER's core contributions are a **dual-policy architecture**—a selection policy (SEL) to choose the next node and a placement policy (PLC) to assign it—and a **three-stage training process** (Imitation Learning, RL in simulation, and RL on the real system). The method is shown to outperform heuristic and RL baselines in execution time across several ML workloads.\n\nWhile the problem is highly relevant and the empirical results are strong, the paper's core contributions are not well-justified. The necessity of the dual-policy (SEL+PLC) over a single placement policy is poorly motivated and its benefits are conflated with other implementation choices. Furthermore, the 3-stage training pipeline relies on an online, real-world RL stage, which is impractical for the very production environments the paper targets, undermining the \"cost-effective\" claims.", "questions": "none", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "### **Strengths**\n\n1.  **Important Problem:** The paper correctly identifies a key inefficiency in modern ML systems: bulk-synchronous execution (e.g., all-reduce) leads to device idle time [cite: 69-72]. It convincingly argues that asynchronous, work-conserving (WC) systems offer significant speedups (Table 1) [cite: 161-163], but that optimizing for them is a more complex temporal problem [cite: 169-171].\n\n2.  **Comprehensive Training Framework:** The three-stage training paradigm (Stage I: Imitation Learning, Stage II: Simulation-based RL, Stage III: Real-system RL) is a thorough approach [cite: 177-180, 487]. Figure 3 demonstrates that this curriculum is highly effective, using imitation and simulation for a strong warm start before the final, high-performance real-system tuning [cite: 731-734].\n\n3.  **Strong Empirical Performance:** The final model, DOPPLER-SYS (which uses all 3 stages), demonstrates state-of-the-art performance. It achieves significant execution time reductions (up to 52.7%) compared to all baselines, including prior RL methods (PLACETO, GDP) and a strong, purpose-built heuristic (ENUMERATIVE OPTIMIZER) [cite: 619-621, 642-644].\n\n4.  **Pragmatic Scalability:** The paper includes a critical implementation detail: running GNN message passing only *once per episode* (per graph) rather than *once per step* (per node) [cite: 481-485]. Appendix G.3 correctly identifies this as a massive efficiency gain (a 30x reduction in message-passing operations for a <1% performance trade-off), making the approach scalable and computationally viable [cite: 2331-2334, 2399-2400].", "weaknesses": "### **Weaknesses and Questions**\n\n1.  **Weak Motivation for Dual Policy:** The paper's central contribution is the \"dual-policy\" (SEL+PLC) architecture [cite: 61, 174-176, 182]. However, the *reason* for factoring the policy this way is poorly justified. The paper never compares this two-policy agent against the most obvious and standard baseline: a *single* policy that directly outputs a `(node, device)` pair. The ablation in Table 3 [cite: 646-649] only compares DOPPLER to variants where one of the two policies is replaced by a heuristic; this does not prove that the dual-policy factorization itself is beneficial.\n\n2.  **Misleading Scalability Claims:** The paper's impressive scaling results (Fig. 4) [cite: 737-739] are presented as a benefit of DOPPLER. However, the text (and Appendix G.3) reveals this speedup is due to an implementation choice (GNN-per-episode) that is *orthogonal* to the dual-policy architecture [cite: 481-485, 2331-2400]. The paper criticizes PLACETO for being \"prohibitive\" and \"inefficient\" because it runs the GNN per-step [cite: 482-483], but this is an apples-to-oranges comparison. A fair comparison would require re-implementing all baselines with the same \"GNN-per-episode\" optimization. As such, the scalability results do not support the *dual-policy architecture* itself.\n\n3.  **Impracticality of Stage III (Real System RL):** The paper's best results rely on `DOPPLER-SYS`, which uses Stage III, an online RL stage on the \"real hardware\" [cite: 180, 552-554]. The paper claims this is \"for free\"[cite: 554], which is highly misleading. This stage requires deploying an *exploratory* (i.e., non-optimal) policy in a production environment, forcing users to suffer through the high-variance, low-performance assignments generated during exploration. Figure 3 (\"0k-0k-8k\" curve) shows exactly how slow and unstable this \"from-scratch\" online learning is [cite: 732-734]. This reliance on an impractical training step undermines the paper's applicability.\n\n4.  **Simulator-Reality Gap:** The only reason the impractical Stage III is necessary is that the simulator in Stage II is not accurate enough. This is shown by the consistent performance gap between `DOPPLER-SIM` (sim-trained) and `DOPPLER-SYS` (real-trained) in Table 2[cite: 619, 643]. The paper explicitly confirms this gap in Appendix G.1, stating the simulator \"tends to overestimate\" and has \"trouble differentiating between assignments\" [cite: 2220-2221]. The 3-stage framework is thus not just a feature but a *crutch* to overcome a flawed simulator.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aYqG1z4wtY", "reviewer": "Reviewer_navF", "summary": "This paper shows DOPPLER, a dual-policy reinforcement learning framework for optimizing device assignment in asynchronous. \nTraditional synchronous scheduling enforces global barriers and that leads to idle GPUs. \nThe presented method addresses device placement in an asynchronous dataflow environment where operations can be executed dynamically. The framework does a three-stage pipeline: (1) imitation learning using heuristic supervision, (2) reinforcement learning in simulation, and (3) online fine-tuning during real deployment. Experimental results on dataflow graphs from matrix chains, feedforward networks, and Llama layers show up to 52.7% reduction in execution time relative to state-of-the-art baselines. The method also generalizes to new hardware and architectures with few-shot adaptation.", "questions": "A few questions about the paper:\n\nIn the Stage III online phase, how is policy stability ensured when runtime-derived rewards could be noisy, is there a mechanism to prevent noisy updates?\n\nWhat is the observed correlation between simulated runtime rewards and real execution times? This could validate the Stage II training phase and strengthen the paper\n\nWould integrating model-parallel or operator-fusion heuristics into DOPPLER’s action space further improve placement quality?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The strengths are:\n\nThe paper is well written and organized. It also shows ablation studies and pseudocode.\n\nThe work addresses device placement under asynchronous execution where small per-query improvements scale to substantial cost reductions at deployment scale. A dual-policy decomposition separates operation selection and device placement. Then they introduce a three-stage training pipeline that bridges imitation learning, simulation-based RL, and real-system adaptation.\n\nThe framework’s shows transferability and few-shot adaptability to new hardware that is relevant for real-world heterogeneous distributed inference and training systems.", "weaknesses": "The weaknesses are:\n\nThe paper is missing explanation or studies for convergence or sample complexity analysis of the dual-policy setup. Adding those could strengthen the understanding of stability during online adaptation. \n\nThe experiments focus on up to 8 GPUs. It remains unclear how DOPPLER scales to large GPU clusters (256 GPUs) or heterogeneous systems with varying compute/memory characteristics.\n\nThe paper presents a 3 stage training, which seems rather complex. The incremental benefit of Stage III appears limited, and simpler two-stage alternatives could potentially achieve similar performance with reduced engineering overhead. The paper mentions a simulator used in Stage II but offers limited quantitative validation of how closely its performance correlates with real hardware results. This impacts the credibility of simulation-based pretraining.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)