ai short=AI, long=artificial intelligence, mtnlg short=MT-NLG, long=Megatron-
Turing natural language generation, dnn short=DNN, long=deep neural network,
ocs short=OCS, long=optical circuit switched, eps short=EPS, long=electronic
packet switched, rl short=RL, long=reinforcement learning, gnn short=GNN,
long=graph neural network, jct short=JCT, long=job completion time, mdp
short=MDP, long=Markov decision process, ppo short=PPO, long=proximal policy
optimisation, dag short=DAG, long=directed acyclic graph, mpi short=MPI,
long=message passing interface, ml short=ML, long=machine learning, co
short=CO, long=combinatorial optimisation, ilp short=ILP, long=integer linear
programme, dqn short=DQN, long=deep Q-learning,

Demonstrate that deciding how much to partition distributed jobs is a key
factor in determining overall system throughput.

Demonstrate that optimising for only the job completion time leads to high
blocking rates and poor throughput in dynamic job arrival scenarios.

Introduce a new partitioning algorithm which leverages reinforcement learning,
a graph neural network, and a novel formulation of the user-defined job
completion time specification to automatically learn to partition jobs such
that the blocking rate is minimised and user requirements are met.

Demonstrate the proposed algorithm out-performing baselines on a state-of-the-
art optical network architecture running five real deep learning computation
graphs.

# Partitioning Distributed Compute Jobs with Reinforcement Learning and Graph
Neural Networks

Christopher W. F. Parsonson111UCL Zacharaya Shabka22footnotemark: 2 Alessandro
Ottino33footnotemark: 3 Georgios Zervas44footnotemark: 4

###### Abstract

From natural language processing to genome sequencing, large-scale machine
learning models are bringing advances to a broad range of fields. Many of
these models are too large to be trained on a single machine, and instead must
be distributed across multiple devices. This has motivated the research of new
compute and network systems capable of handling such tasks. In particular,
recent work has focused on developing management schemes which decide how to
allocate distributed resources such that some overall objective, such as
minimising the job completion time (JCT), is optimised. However, such studies
omit explicit consideration of how much a job should be distributed, usually
assuming that maximum distribution is desirable. In this work, we show that
maximum parallelisation is sub-optimal in relation to user-critical metrics
such as throughput and blocking rate. To address this, we propose PAC-ML
(partitioning for asynchronous computing with machine learning). PAC-ML
leverages a graph neural network and reinforcement learning to learn how much
to partition computation graphs such that the number of jobs which meet
arbitrary user-defined JCT requirements is maximised. In experiments with five
real deep learning computation graphs on a recently proposed optical
architecture across four user-defined JCT requirement distributions, we
demonstrate PAC-ML achieving up to $56.2\%$ lower blocking rates in dynamic
job arrival settings than the canonical maximum parallelisation strategy used
by most prior works.

###### keywords:

Deep Learning , Reinforcement Learning , Graph Neural Networks , Distributed
Asynchronous Computing , Job Partitioning , Optical Networks

††journal: Journal of Parallel and Distributed Computing

##  1 Introduction

The last decade has seen an exponential increase in the amount of compute
demanded by big data jobs such as ai and genome processing, with resource
requirements doubling every $3.4$ months since 2012; $50\times$ faster than
Moore’s Law (OpenAI, 2018). This trend is showing no sign of slowing down. The
fundamental relationship between neural network accuracy and scale (Kaplan et
al., 2020) provides a strong incentive for practitioners seeking performance
improvement to further increase their resource requirements. Moreover, brain-
scale ai will require at least as many parameters as the $\approx$1 000
trillion synapses present in the human brain (Furber, 2016); several orders of
magnitude more than the largest models used today.

The compute time and memory requirements of state-of-the-art big data
applications already far outstrip the capabilities of any single hardware
device. For example, one of the current largest dnn, mtnlg (Smith et al.,
2022), contains 530 billion parameters. These parameters alone occupy
$\approx$1 000 GB, exceeding the capacity of the largest A100 GPU by over an
order of magnitude, and the parameter loss gradients tracked during training
occupy several times more. Even if the model could be fitted onto a single
device, the training time would be $\approx$900 years555Assuming it takes 8
V100 GPUs $36$ years to train a 175 billion parameter model (NVIDIA, 2022) and
extrapolating.. To address these compute time and memory demands, rather than
using a single device, big data jobs must be distributed and parallelised
across a cluster of machines. For example, the Selene supercomputing cluster
(NVIDIA, 2020) consists of 358 400 A100 GPU tensor cores, bringing the mtnlg
training time from 900 years down to the order of days666Assuming a linear
parallelisation speedup and $0$ communication overhead..

However, parallelising jobs across ever-more machines brings its own
challenges. With any parallelisation strategy, at some point the output of
each ‘worker’ (a single device processing at least part of a job) must be
collected and synchronised to get the overall result of the parallelised
computation. This synchronisation requires _communication_ between the
workers. As the number of workers used to execute a job is increased, the per-
worker computation demands decrease, but the overall communication overhead
between workers grows (see Figure 1). This shifts the performance bottleneck
away from the workers themselves and into the network connecting them, and
brings additional challenges with managing varying traffic characteristics for
different job types and parallelisation strategies (Wang et al., 2022;
Parsonson et al., 2022a; Benjamin et al., 2021, 2022).

Figure 1: How the network overhead of six distributed deep learning jobs
(encompassing object tracking, recommendation, natural language processing,
and image recognition) increases with the number of workers used in Meta’s GPU
cluster (Wang et al., 2022).

To address the communication bottleneck in distributed computing, recent works
have sought to develop optical clusters (Benjamin et al., 2020; Ballani et
al., 2020; Khani et al., 2021; Wang et al., 2022; Ottino et al., 2022);
machines interconnected by optical switches (Parsonson et al., 2020; Gerard et
al., 2020, 2021). Compared to their electronic counterparts, optically
switched networks offer orders of magnitude improvements in scalability,
bandwidth, latency, and power consumption (Ballani et al., 2020; Zervas et
al., 2018; Mishra et al., 2021) (see Section 3).

Optical clusters are typically operated under the ocs paradigm due to its non-
blocking circuit configurations with high capacity and scalability (Raja et
al., 2021). ocs networks are fundamentally different from the eps
architectures used by most current clusters, resulting in entirely new
communication patterns and resource demand characteristics. Consequently, new
compute and network resource management schemes are needed in order to
optimally allocate jobs and maximise performance.

Of the many resource management tasks which must be performed in a compute
cluster, _job partitioning_ (how to split a job up across how many devices) is
key to overall performance. More partitioning can lead to lower compute times.
However, it may also increase network overhead and occupancy of cluster
resources, possibly leading to future jobs being blocked upon arrival and
consequently lower overall cluster throughput. Prior works such as SiP-ML
(Khani et al., 2021) have introduced simple partitioning heuristics for
optical networks which have notably improved cluster performance. However,
they have not been designed under the more realistic setting of dynamic and
stochastic job arrivals, have not considered the state of the cluster in a
‘network-aware’ manner when making partitioning decisions, and have been
crafted to optimise for the sub-optimal objective of minimising jct.

In this work, we first argue that simply minimising the jct is a naive
objective because it brazenly encourages more parallelisation of a job request
without considering the effect this has on the ability of a cluster to service
subsequent jobs. We then introduce a new more subtle formulation of the
optimisation metric, the _user-defined blocking rate_ , which more aptly
encompasses the desires of cluster users. Next, we propose a simple
modification of the quantised SiP-ML partitioner which, rather than maximally
parallelising all jobs, minimally parallelises them such that they meet the
user-defined maximum acceptable completion time. Then, we propose a novel
network-aware partitioning strategy (see Figure 5 and Section 5) called PAC-ML
(partitioning for asynchronous computing with machine learning) which utilises
rl and a gnn to flexibly meet the demands of the user in an arbitrary manner
given the current state of the cluster network. Finally, we demonstrate our
method in simulation on the recently propsed RAMP optical architecture (Ottino
et al., 2022), achieving up to $56.2\%$ lower blocking rates than the best
heuristic baseline. We show that different user-defined demand environments
require different partitioning strategies for optimal results, and that a key
advantage of PAC-ML is that it is able to discover performant strategies
automatically without the need for handcrafted heuristics or environment-
specific tuning.

##  2 Related Work

Recent years have seen a surge of interest in developing methods to distribute
ml tasks across multiple devices (Ben-Nun and Hoefler, 2019; Mayer and
Jacobsen, 2020). One approach has been to optimise the physical plane of the
distributed cluster such as its compute and network devices and architectures
(Parsonson et al., 2020; Khani et al., 2021; Wang et al., 2022; Ottino et al.,
2022). In this work, we instead focus on optimising the virtual plane, which
determines how physical layer resources are allocated to execute a job. We
divide the virtual plane into three sub-components: Job $(1)$ partitioning
(how many devices to use); $(2)$ placement (which devices to use); and $(3)$
scheduling (in which order to use the devices). Many prior virtual plane works
have considered $(2)$ and $(3)$ (how to distribute), whereas we focus on $(1)$
(how much to distribute). However, in this section we comment on recent
progress across all these fields, since we leverage this progress throughout
our work.

ML for discrete optimisation. Many co problems turn out to be NP-hard,
rendering exhaustive search techniques intractable for practical application
(Bengio et al., 2021). Consequently, practitioners rely on either approximate
algorithms, which give restricted performance guarantees and poor scalability
(Williamson and Shmoys, 2011), or heuristics, which have limited solution
efficacy (Halim and Ismail, 2019). Since the first application of neural
networks to co by Hopfield and Tank (1985), the last decade has seen a
resurgence in ML-for-CO (Bello* et al., 2017; Dai et al., 2017; Barrett et
al., 2019; Gasse et al., 2019; Barrett et al., 2022; Parsonson et al., 2022b).
The advantages of ML-for-CO over approximation algorithms and heuristics
include handling complex problems at scale, learning either without external
input and achieving super-human performance or imitating strong but
computationally expensive solvers, and (after training) leveraging the fast
inference time of a dnn forward pass to rapidly generate solutions. Since
almost all cluster resource management tasks can be reduced to canonical co
problems (Bengio et al., 2021), many state-of-the-art resource management
methods utilise recent advances in ml-for-co.

Job placement. Mirhoseini et al. (2017) were the first to apply ml to the task
of deciding which operations in a computation graph to place on which devices
in a cluster. They used a sequence-to-sequence model consisting of an LSTM dnn
with an attention mechanism trained with the simple REINFORCE policy gradient
rl algorithm (Williams, 1992) such that the jct of a deep learning job was
minimised, outperforming handcrafted heuristics when training the Inception-V3
computer vision and LSTM natural language processing models. Gao et al. (2018)
furthered this work by replacing REINFORCE with the more advanced ppo rl
algorithm (Schulman et al., 2017) with lower variance and reduced training
hardware demands. They demonstrated their method beating Mirhoseini et al.
(2017) on the CIFAR-10 image recognition benchmark in terms of jct. Mirhoseini
et al. (2018a) proposed a novel hierarchical model which decomposed the job
placement task into a joint group-and-place problem, reducing the jct of
Inception-V3, ResNet, LSTM, and NMT models by up to $60\%$ relative to the
state-of-the-art.

All works up to this point used dnn architectures restricted to Euclidean-
structured input data. Consequently, in order to handle non-Euclidean graph-
structured data such as computation graphs and cluster networks, they had to
be re-trained each time a new graph structure was considered. Addanki et al.
(2019a) were the first to instead leverage a gnn, as well as the grouping
scheme of Mirhoseini et al. (2018a), to learn to generalise across different
job types with varying computation graph structures, demonstrating device
placement schemes which were on par with or better than prior approaches on
Inception-V4, NASNet, and NMT after $6.1\times$ fewer training steps. Khadka
et al. (2021) furthered the use of gnn for job placement by combining gnn, rl,
and population-based evolutionary search with the hierarchical group-and-place
scheme of Mirhoseini et al. (2018a). Concretely, they replaced the manually-
designed operation grouping heuristic with a learned policy capable of
superior scaling and jct performance.

Job scheduling. Bao et al. (2018) addressed the job scheduling problem (the
order in which to execute operations placed across a set of devices) using a
primal-dual framework for online job scheduling. They represented the problem
as an ilp which their proposed algorithm could solve in polynomial time in an
online fashion such that the cluster resources were maximally utilised and the
jct minimised. Li et al. (2021) proposed a placement-aware scheme which
leveraged the pre-determined device placement allocation to decide on a job
schedule which could reduce the average jct by up to $25\%$ relative to other
scheduling methods. Paliwal et al. (2020) went further by utilising an rl-
trained gnn and a genetic algorithm to jointly optimise both job placement and
scheduling, demonstrating both lower jct and peak memory usage than other
strategies when distributing TensorFlow computation graphs across a cluster.

Job partitioning. To the best of our knowledge, Khani et al. (2021) are the
only ones to have explicitly considered the question of how much to distribute
a computation graph in the context of an optical network. Like other works,
they assumed that a maximum parallelisation strategy (i.e. partition the job
across as many workers as possible) is a desirable objective, and then focused
on how best to design the physical layer such that the jct could be minimised.

All works discussed in this section have assumed that the jct is the key
objective to minimise. Consequently, where the question of partitioning is
considered, prior works have assumed that more parallelisation is desirable.
However, we posit that user-critical metrics such as throughput and blocking
rate are compromised by prioritising optimisation of the jct in a cluster
setting with dynamic job arrivals. To address this shortcoming, we propose a
new ml-based resource management scheme which explicitly addresses the
partitioning question. Concretely, our work leverages the emergent trend from
these other virtual plane fields, namely utilising an rl-trained gnn, to
decide how much to partition different jobs in a dynamic setting with
arbitrary user-defined completion time requirements.

##  3 Background

Figure 2: Diagram showing a dnn job DAG being partitioned. Top: A forward pass
DAG where each node has an associated partition degree (how many times it will
be divided when partitioned). Bottom: A partitioned DAG with forward and
backward passes handled consecutively. Green edges in the graph represent data
flow (i.e. output to input) between consecutive operations in the forward
pass. Orange edges represent gradient exchanges processed in the backward pass
(backpropagation). Blue edges represent full connectivity collective
operations to synchronise weight updates across partitioned components of an
operation. Note that, for brevity, the top unpartitioned DAG only shows the
forward pass (since, before partitioning, the graph structure is identical to
the backward pass), whereas the bottom partitioned DAG shows both the forward
and backwards passes (since, after partitioning, the graph structures are
different).

###  3.1 Parallelisation

Types of parallelism. Parallelisation is the process of distributing a
computational job across multiple devices. This is done in order to reduce the
time and/or physical memory needed to complete the job. There are three main
types of deep learning parallelism; data parallelism, model parallelism, and
hybrid parallelism (see Appendix 9.1.1 for extended background information on
these methods). Although today the most common method for dnn training
parallelisation is data parallelism for its simplicity and limited network
overhead, we focus on the less common but more desirable model parallelism
paradigm for its strong scaling capabilities (Khani et al., 2021). Our
proposed partitioning methods are applicable to hybrid and pipeline
parallelism, but these require additional simulation complexity and are
therefore beyond the scope of this manuscript.

Computational jobs. A computational job is a dag whose nodes are operations
and edges are dependencies. Operations are computational tasks (e.g. some
mathematical reduction, a database query, etc.). Dependencies are either
control dependencies, where the child operation can only begin once the parent
operation has been completed, or data dependencies, where at least one tensor
is output from the parent operation and required as input to the child
operation. In the context of dnn, a job dag is a sequence of forward pass,
backward pass, and parameter update operations which need to be performed on
data exchanged between operations. Whether or not this data passes through a
communication network is determined by how the operations are partitioned,
placed across a cluster of workers, and parallelised.

Job partitioning. Job partitioning refers to the process of splitting the
operations of a job dag into $u$ (the partition degree) smaller sub-operations
which can in turn be placed across $u$ workers, thus reducing their run time
and memory requirements. Partitioning is used in the model, hybrid, and
pipeline parallelisim paradigms. More partitioning can decrease compute time
and memory requirements, but requires more inter-worker communication, complex
intra-worker operation scheduling, and greater resource utilisation, therefore
potentially increasing overall completion time, cluster complexity, and
subsequent job blocking rates. Figure 2 visualises how an initial DAG for some
arbitrary neural network architecture, where each operation has a partitioning
degree, can be re-represented in terms of its partitioned form. Both forward
and backward passes are explicitly represented since inter-operation
information dependencies (i.e. the edges in the graph) are not the same in
each pass.

###  3.2 Optical Networking

Most current cluster networks use optic fibre communication links, but the
switch devices which interconnect the network are usually electronic.

Limitations of electronic networking. Electronic networks have poor
scalability, bandwidth, latency, and power consumption. Concretely, since the
per-port bandwidth is limited and the power consumption required to cool
active electronic devices is expensive, the bisection bandwidth achievable in
an electronic network is restricted, thus hampering scalability. Consequently,
although the compute power of DCN server nodes, as measured by FLOP/s, has
increased by a factor of $65$ over the last $18$ years, the bandwidth of the
DCN network facilitating communication between these nodes has only increased
by a factor of $4.8$, resulting in an 8-factor decrease in bytes communicated
per FLOP (Bergman, 2018). This has created a performance bottleneck not in the
server nodes themselves, but rather in the network connecting them. This issue
is especially compounded when striving for strong scaling via model
parallelism with distributed computing, and with the trend towards larger
models with ever more parameters as described in Section 1.

Optical circuit switched networks. Cluster networks with optical switches have
the potential to offer significant improvements in performance (due to larger
bandwidth and lower switching latency) and energy efficiency (due to the lack
of optical-electronic-optical conversion overhead), as well as the capability
to scale to next-generation large-scale distributed compute jobs with exascale
bandwidth and compute (Ottino et al., 2022). ocs networks in particular offer
a promising avenue with which to realise commercial optical networks due to
their non-blocking circuit configurations with high capacity and scalability
and low deterministic switching latency. In contrast to optical packet
switched networks, ocs networks are simpler to implement and they eliminate
the need for in-switch buffering or queuing and addressing.

Figure 3: The mean network overhead of the $6$ distributed deep learning jobs
reported by (Wang et al., 2022) in Meta’s GPU cluster compared to that of RAMP
as reported by Ottino et al. (2022) on the $5$ jobs considered in our work.
Note that this is an approximate comparison, and that the important takeaway
is that RAMP retains low network overheads as jobs become increasingly
distributed.

RAMP. RAMP is a state-of-the-art ocs architecture designed specifically for
cloud data centres and distributed deep learning systems (Ottino et al.,
2022). RAMP networks are parameterised by $N_{C}$ communication groups,
$N_{R}$ racks per communication group, and $N_{S}$ servers per rack, resulting
in a $N_{W}=N_{C}\times N_{R}\times N_{S}$ worker cluster with a colloquially
termed ‘RAMP shape’ defined by tuple $\langle N_{C},N_{R},N_{S}\rangle$. At
its core, RAMP proposes a novel set of mpi for performing the synchronisation
steps (AllReduce, AllGather, etc.) required by distributed dnn training jobs.
These will be referred to as collective operations. These mpi are designed to
take full advantage of the high bandwidth provided by optical network
architectures. Consequently, as shown in Figure 3, the network overhead of
RAMP remains remarkably low as the number of workers used to execute a job
increase (see Section 6 for experimental details). The RAMP authors showed
that this low network overhead enables unprecedented scalability with up to
$65\,536$ worker nodes capable of training $O(\text{trillion})$ parameter dnn
models.

RAMP placement rules. As detailed in Ottino et al. (2022), a group of workers
in a RAMP shape can only undergo collective operations if they are selected
with respect to certain rules, loosely termed here ‘symmetry’ rules. For shape
$\langle N_{C},N_{R},N_{S}\rangle$, these rules are as follows: $(1)$ $N_{S}$
workers per rack spread over $N_{R}$ racks requires that the set of workers on
each rack span $N_{R}$ distinct communication groups. These $N_{R}$ distinct
communication groups do not have to be the same set across racks. $(2)$
$N_{S}$ workers on $N_{R}=1$ rack must span $N_{S}$ communication groups.
$(3)$ $N_{S}$ workers spread over $N_{R}$ racks ($N_{S}=1$ worker per rack)
must span $N_{S}$ distinct communication groups.

In our simulations, we use a simple first-fit operation placement heuristic
which conforms to these rules (refer to Appendix 9.4.4 for further details).

###  3.3 Reinforcement Learning

rl is the study of optimal decision making in natural and artificial systems
(Sutton and Barto, 2018). In the general rl setting shown in Figure 5, an
agent interacts with an environment at each sequential time step $t$. The
environment can be described by tuple $\langle T,R\rangle$, where $T$ is a
state transition probability matrix defining the transition probabilities from
all states $s$ to all successor states $s^{\prime}$ taking action $u$ where
$T^{u}_{ss^{\prime}}=\mathbb{P}(S^{t+1}=s^{\prime}|S^{t}=s,U^{t}=u)$, and $R$
is a scalar reward function giving the expected immediate (next state) reward
given current state $s$ and chosen action $u$ where
$R^{u}_{s}=\mathbb{E}(R^{t+1}|S^{t}=s,U^{t}=u)$.

Markov decision process. The environment is usually assumed to have the Markov
property whereby $\mathbb{P}(s^{t+1}|s^{t})=\mathbb{P}(s^{t+1}|h^{t})$; that
is to say that the probability of the next state being $s^{t+1}$ given the
current state $s^{t}$ is the same as the equivalent probability given all
previous states in history $h^{t}=\\{s^{1},...,s^{t}\\}$. As such, this RL
setting is usually assumed to be a mdp described by tuple $\langle
S,U,T,R,\gamma\rangle$ where $S$ is a finite set of possible environment
states, $U$ is either a discrete (finite) or continuous (infinite) set of
possible actions, and $\gamma\in[0,1]$ is a discount factor specifying the
factor by which to multiply future expected rewards to discount their present
value. Since Markov states are stochastic, future rewards are never fully
certain and are therefore expressed as an expectation.

Agent goal. The agent’s goal is to learn to maximise its expected total
discounted future reward, termed the ‘value’ or ‘return’
$G^{t}=\sum_{k=0}^{\infty}\gamma^{k}R^{t+k+1}$, over the course of an episode
(a sequence of decision steps which may or may not terminate at some point).
To do so, the agent can use model-free rl to avoid explicitly modelling the
environment by only using its policy function and/or its value function to
make decisions. The policy function $\pi$ maps an observed state $s^{t}$ to a
corresponding action $u^{t}$ such that some estimated score objective is
maximised. The value function estimates the expected return $G_{t}$ from being
in state $s^{t}$ and following policy $\pi$ (the state value function $v$) or
from being in state $s^{t}$, taking action $u^{t}$, and following policy $\pi$
(the action value function $q$). Crucially, value and policy functions can be
approximated and learned with dnn, enabling rl to be scaled to large problem
instances (see Appendix 9.1.2 for extended background information on dnn).

Advantages of RL. Using traditional rl has several advantages over heuristics
and other ml paradigms such as supervised learning. First, no external data
from human-designed or computationally expensive heuristics is required,
enabling an agent to learn super-human policies without potentially sub-
optimal initial biases towards a certain strategy or a costly expert example
collection-and-labelling phase (Silver et al., 2016). Second, a dnn with a
finite number of layers and neurons will have its expressivity constrained
(Dong et al., 2020), restricting the complexity of the set of functions it is
capable of approximating. Because the objective of an rl agent is to maximise
its expected future return which, under the assumption that a suitable reward
function has been crafted, is equivalent to maximising performance on a given
task, rl agents are able to maximise task performance given dnn expressivity
constraints. Third, since rl agents maximise future return, they are capable
of learning sophisticated non-myopic policies which sacrifice short-term
reward in exchange for higher long-term return (Sutton and Barto, 2018).

##  4 User-Defined Blocking Rate

To motivate our work, we first explore the key metrics to consider when
evaluating a job partitioning strategy with the help of an experiment on $32$
GPU workers, and then introduce a new formulation of the user-defined blocking
rate. All experimental details are given in Section 6.

Figure 4: (a-b) Demonstration of how more partitioning can lead to a lower jct
than no partitioning (i.e. sequentially running the job on a single device),
but this may be at the cost of a higher blocking rate since more cluster
resources are occupied when subsequent jobs arrive. (c-d) Demonstration of how
optimising for the cluster throughput leads to an unfair bias towards more
partitioning, because more parallelism creates more work for the cluster and
therefore artificially increases cluster throughput even though, from the
perspective of the user, the original offered throughput may be lower.

The inadequacy of optimising the job completion time. As discussed in Section
2, most prior works researching management schemes for distributed computing
aim to minimise jct; the time taken to complete a given job. If a job $j$
begins running at wall clock time $t_{wc,j}^{\text{start}}$ and is completed
at time $t_{wc,j}^{\text{end}}$, researchers usually record the completion
time as $\text{JCT}_{j}=t_{wc,j}^{\text{end}}-t_{wc,j}^{\text{start}}$.
Consequently, most systems maximise the degree to which they parallelise jobs
in order to minimise jct. While it is true that end users undoubtedly want
this jct metric to be minimised, it fails to quantify when a job was blocked,
which occurs when no cluster resources were available to service it. While
more parallelism will often lead to a lower jct for a given job, it will also
use up more of the cluster’s compute and network resources, potentially
blocking future job arrivals (see Figure 4). Therefore in practice, end-users
wish to minimise both the jct and the overall blocking rate (the fraction of
jobs blocked over a given time period). While maximum parallelisation will
lead to a minimised jct, we posit that a balance between these two extreme
parallelisation strategies can more aptly optimise for both the jct and
blocking rate.

Alternative optimisation objectives. One metric which encapsulates both the
jct and blocking rate is throughput; the information processed per unit time.
There are two issues with using throughput as an optimisation objective. (1)
Operators must be careful how they measure the throughput to be optimised. If
they measure the cluster throughput (the total cluster information processed
per unit time), they will be biased towards more parallelisation, because when
a job is partitioned and parallelised, the edge dependencies coming in to and
out of the partitioned operation node(s) must be replicated (see Figure 2).
This artificially creates more information for the cluster to process even
though, from the end users’ perspective, the total information processed of
their original demand is the same. Therefore, the offered throughput (the
total original demand information (i.e. before partitioning was applied)
processed per unit time) is a more suitable throughput metric to optimise.
Figure 4 shows an example of how a ‘maximum partitioning’ strategy, such as
that used by SiP-ML (Khani et al., 2021), can have superior cluster throughput
when compared to a ‘no partitioning’ strategy (sequentially running the job on
a single device) despite having lower offered throughput. However, offered
throughput is still an inadequate optimisation metric, because (2) in
practice, different jobs being serviced by the cluster originating from
different client users have different priorities and job completion time
requirements. For example, two identical machine learning training jobs might
be submitted to the cluster, but one from a user who intends to deploy the
model commercially and requires it to be completed overnight, and the other
from a user who is employing the model for research and has less stringent
completion time requirements. Ideally, operators would direct their clusters
to meet flexible user-defined per-job completion time requirements.

The user-defined blocking rate. To enable users to dynamically determine the
completion time on a per-job basis whilst also maximising the number of job
demands satisfied, we introduce a new formulation of the user-defined blocking
rate objective for the partitioning algorithm to optimise. Given a job which,
if executed sequentially on one device, would be completed in
$\text{JCT}_{j}^{\text{seq}}$, we define the maximum acceptable jct as
$\text{JCT}_{j}^{\text{acc}}=\beta\cdot\text{JCT}_{j}^{\text{seq}}$, where
$\\{\beta\in\mathbb{R}:0<\beta\leq 1\\}$. Here, $\beta$ is a parameter chosen
by the user which determines how quickly the job must be completed. If
$\text{JCT}_{j}>\beta\cdot\text{JCT}_{j}^{\text{seq}}$, then the cluster will
have failed to complete the job within the required time and the job will be
recorded as having been blocked. The user-defined blocking rate is therefore
the fraction of jobs which failed to meet the
$\text{JCT}_{j}\leq\beta\cdot\text{JCT}_{j}^{\text{seq}}$ requirement over a
given period of time. Note that rather than brazenly optimising for either the
jct or the blocking rate, the user-defined blocking rate enables the cluster
operator to instead dynamically specify their desired completion time on a
per-job basis, and the performance of the cluster is evaluated according to
how well it was able to meet the requirements of the user. Furthermore, the
$\beta$ parameter corresponds to the speed-up factor being requested by the
user and, since $\\{\beta\in\mathbb{R}:0<\beta\leq 1\\}$, can be given
directly as input to a dnn.

##  5 PAC-ML Partitioning Methodology

Figure 5: An overview of our PAC-ML approach transitioning from step
$t\rightarrow t+1$. At each time step $t$ when there is a new job to be placed
on the cluster, we: (i) Use a gnn to generate an embedded representation of
the node and edge features in the job’s computation graph, and a standard
feedforward dnn to do the same for the global job and cluster features; (ii)
concatenate the outputs of (i) and use another feedforward dnn to generate a
logit for each action $u^{t}\in U^{t}$; (iii) pass the chosen action $u^{t}$
to the environment and partition the job accordingly; (iv) apply any internal
environment allocation heuristics (operation and dependency placement and
scheduling, etc.) to attempt to host the job on the cluster; (v) if accepted
onto the cluster, perform a lookahead to evaluate the job’s completion time;
(vi) fast-forward the environment’s wall clock time $t_{\text{wc}}$ to when
the next job arrives, and return the corresponding reward $r^{t+1}$ and
updated state $s^{t+1}$.

rl agents can learn general policies without the need for human guidance. An
rl job partitioner therefore has the potential to take an arbitrary maximum
acceptable jct provided by the user and automatically decide how much to
distribute the job such that, over a period of time, the number of jobs which
meet the jct requirements specified by the user is maximised. Such an agent
would therefore be able to minimise the blocking rate whilst also accounting
for the flexible and dynamic jct specifications of the user. Following this
logic, we now describe our PAC-ML (partitioning for asynchronous computing
with machine learning) approach for learning to partition computation jobs
with rl and a gnn.

###  5.1 Markov Decision Process Formulation

Since allocating cluster resources for jobs arriving dynamically in time is a
sequential decision making process, formulating problems such as job
partitioning as an mdp is a natural approach and facilitates the application
of many traditional and state-of-the-art rl algorithms (Mao et al., 2016;
Addanki et al., 2019b; Paliwal et al., 2020).

States. A new job $j$ arriving at time step $t$ is comprised of a dag
$G(O,D,g_{j})$ with node operations $O$, edge dependencies $D$, and any other
job statistics which might be recorded $g_{j}$. Similarly, the state of the
cluster at time $t$ is made up of the number of workers available, the jobs
currently running on the cluster, and so on. To compress the state of the
cluster and the job requesting to be placed into a representation suitable as
input for a neural network at time step $t$, we encode this information into
five feature vectors:

  1. 1.

Per-operation features $o_{i}\forall i\in\\{1,...,|O|\\}$ ($5$ features): (i)
The compute cost (run time in seconds on an A100 GPU); (ii) a binary variable
indicating whether the operation has the greatest compute cost in the job;
(iii) the memory cost (byte occupancy); (iv) a binary variable indicating
whether the operation has the greatest memory cost in the job; and (v) the
node depth with respect to the source node. The compute and memory costs are
normalised by the highest compute and memory cost operations in the job, and
the node depth is normalised by the depth of the deepest node.

  2. 2.

Per-dependency features $d_{i}\forall i\in\\{1,...,|D|\\}$ ($2$ features): (i)
The size (in bytes) of the edge dependency normalised by the largest
dependency in the job; and (ii) a binary indicator of whether the dependency
is the largest in the job.

  3. 3.

Global job features $g_{j}$ ($15$ features): (i) The number of operations;
(ii) the number of dependencies; (iii) the sequential job completion time;
(iv) the maximum acceptable job completion time; the maximum acceptable job
completion time fraction $\beta$ both (v) raw and (vi) normalised; (vii) the
total memory cost of all of the operations; (viii) the total size of all of
the dependencies; (ix) the number of training steps which need to be
performed; the (x) mean and (xi) median of the operation compute costs; the
(xii) mean and (xiii) median of the operation memory costs; and (xiv) the mean
and (xv) median of the dependency sizes. Each feature is normalised by the
highest respective value of the feature across all job types.

  4. 4.

Global cluster features $g^{t}_{C}$ ($2$ features): (i) The number of occupied
workers; and (ii) the number of jobs running. Both features are normalised by
the total number of workers in the cluster $N_{W}$.

  5. 5.

Global action features $g^{t}_{U}$ ($\frac{N_{W}}{2}$ features): A binary
vector indicating the validity of each possible partitioning decision given
the state of the cluster and the RAMP rules defined by (Ottino et al., 2022).

Actions. Given the state $s^{t}$ encapsulating both the job requesting to be
placed and the current state of the cluster, the partitioning agent uses a
policy $\pi(s^{t})$ to select a number of times $u^{t}$ up to which to
partition each operation in the job’s computation graph (using a similar
minimum operation run time quantum discretisation scheme to Khani et al.
(2021)), where $u^{t}_{i}\forall i\in\\{0,1,...,\frac{N_{W}}{2}\\}$ (i.e.
there are $\big{(}\frac{N_{W}}{2}+1\big{)}$ possible discrete actions). Note
that $u^{t}=0$ enables the agent to reject a job without placing it, $u^{t}=1$
places the job onto one worker and runs it sequentially, and
$1<u^{t}\leq\frac{N_{W}}{2}$ attempts to distribute the job’s operations
across up to $u^{t}$ workers. In our setting and given the RAMP rules of
Ottino et al. (2022), an invalid partitioning action is one which is at least
one of: (i) An odd number (except $u^{t}=1$), or either (ii) greater than the
number of workers available or (iii) has no valid RAMP placement shape given
the current state of the cluster (see Section 3).

Rewards. As a consequence of the RAMP rules defined by Ottino et al. (2022),
which require that the worker and network resources allocated to a given job
are reserved exclusively for that job for the duration of its run time, we are
able to perform a deterministic lookahead to evaluate what the overall
completion time, $\text{JCT}_{j}$, of the job will be as soon as it is placed.
Subsequently, when a job $j$ arrives at time step $t$, we can immediately
determine whether or not the cluster met the $\text{JCT}^{\text{acc}}_{j}$
specified by the user. This enables the use of a simple per-step $+1$/$-1$
reward scheme,

$r^{t+1}=\begin{cases}1,&\text{if
}\text{JCT}_{j}\leq\beta\cdot\text{JCT}_{j}^{\text{seq}}\\\
-1,&\text{otherwise}\end{cases},$ (1)

which when aggregated and maximised over the course of an episode corresponds
to maximally meeting the specified per-job completion time requirements and
therefore minimising the user-defined blocking rate.

Transitions. In our hybrid time- and event-drive simulation, when the agent
makes a partitioning decision at time step $t$, the environment transitions to
the next step $t+1$ by fast-forwarding its internal simulated wall clock time,
$t_{\text{wc}}$, to when the next job arrives and requests to be placed,
updating the states of any running and completed jobs and their corresponding
compute and network resources as necessary. The episode terminates when
$t_{\text{wc}}=T_{\text{wc}}^{\text{max}}$.

###  5.2 PAC-ML Learning Setup

Reinforcement learning algorithm. To find a policy which maximises the
expected return when partitioning jobs, we used the state-of-the-art Ape-X DQN
(Horgan et al., 2018) rl algorithm; a distributed and highly scalable value-
based method (see Appendix 9.7 for algorithm details and hyperparameters).

Neural network architecture. To make the learning of value and policy
functions tractable in large state-action spaces, we approximated them with a
custom-built message passing gnn implemented using the open-source PyTorch
(Paszke et al., 2019) and DGL (Wang et al., 2019) libraries. Refer to Appendix
9.6 for further architectural details.

##  6 Experimental Setup

All code for reproducing the experiments and links to the generated data sets
are provided at https://github.com/cwfparsonson/ddls.

Figure 6: The four $\beta$ distributions used in our experiments in order to
measure the capability of each partitioner to cater to different user-defined
maximum acceptable completion time requirement settings. In each $\beta_{X}$
experiment setting, each new job generated was assigned a $\beta$ value
sampled from $\beta_{X}$ in order to get the maximum acceptable job completion
time, $\beta\cdot\text{JCT}^{\text{seq}}$ (see Section 4).

Simulation environment. We built an open-source Gym environment (Brockman et
al., 2016) to simulate the RAMP ocs system of Ottino et al. (2022) in an rl-
compatible manner. We used a hybrid time- and event- driven simulation
approach where we kept track of the internal simulation wall clock time
$t_{wc}$, enabling the measurement of time-based metrics, but only took a
partitioning decision when needed (i.e. when a new job demand arrived at the
cluster), aiding efficiency since no discrete steps were needlessly simulated.
All our experiments used similar cluster parameters to Ottino et al. (2022).
We used $N_{W}=32$ $(N_{C}=4,N_{R}=4,N_{S}=2)$ A100 GPUs with $80$ GB memory
capacity, $2$ THz memory frequency, and a peak computational power of $130$
Tflop/s. We assumed an intra-GPU propagation latency of $50$ ns, a negligible
ocs circuit reconfiguration latency of $1$ ns, a worker input-output latency
of $100$ ns, and a total worker communication capacity of $1.6$ TB/s
(resulting in a per-transceiver bandwidth of $\frac{1.6\times 10^{12}}{N_{C}}$
B/s). All experiments were run up to a simulated wall clock time of
$T_{\text{wc}}^{\text{max}}=10^{6}$ s (around $12$ days) of continuous cluster
operation with dynamic job arrivals and were repeated across $3$ random seeds,
with the subsequent min-max confidence intervals for each measurement metric
reported. More details of the simulation environment are provided in Appendix
9.4.

Compute jobs. We used the computation graph time and memory profiles of five
real deep learning job types open-accessed with Microsoft’s PipeDream research
(Narayanan et al., 2019, 2021) (see Appendix 9.5 for details). These jobs
encompassed image classification (AlexNet (Krizhevsky et al., 2012), ResNet-18
(He et al., 2016), SqueezeNet-10 (Iandola et al., 2016), and VGG-16 (Simonyan
and Zisserman, 2014)) and natural language processing (GNMT (Wu et al., 2016))
tasks, thereby testing the generality of the approaches we considered. All
jobs arrived to the cluster dynamically and stochastically throughout the
simulation period, with the inter-arrival time fixed at $1000$ s to control
the load rate. Each job was ran for $N_{iter}=50$ training iterations, where
one training iteration consists of one forward and backward pass through the
neural network.

Partitioning. When partitioning the operations in a job’s computation graph,
we allowed the partitioning agents to split each operation up to
$\frac{N_{W}}{2}$ times (the environment’s ‘maximum partitioning degree’). We
followed Khani et al. (2021) by $(1)$ assuming a linear dependency between the
total number of operation splits and each split’s compute time; and $(2)$
choosing a minimum quantum of computation time, $\tau$, and splitting
operations up to a number of times which would result in sub-operations with a
compute time no smaller than $\tau$ in order to maximise GPU utilisation. We
set $\tau=10$ ms. As such, a given partitioning action $u^{t}$ set the
maxmimum partitioning degree of the job, but individual operations within the
job could be split fewer times depending on their initial compute time and
$\tau$. Note that although this restricts each operation to be distributed
across a maximum of $u^{t}$ servers, the total number of workers used by all
operations in the job can still be greater than $u^{t}$ depending on the
operation placement heuristic’s choices.

Maximum acceptable job completion times. In our setting, a partitioner would
ideally be able to take an arbitrary job with an arbitrary maximum acceptable
job completion time, $\beta\cdot\text{JCT}^{\text{seq}}$, and partition the
job such that the completion time requirement is satisfied for as many
dynamically arriving jobs as possible (thereby minimising the user-defined
blocking rate; see Section 4). To test each partitioner’s ability to do this,
we ran experiments using four $\beta$ distributions
($\beta_{A},\beta_{B},\beta_{C},\text{ and }\beta_{D}$; see Figure 6). For
each $\beta_{X}$ experiment, when one of the five possible jobs was randomly
generated to arrive at the cluster, a $\beta$ value, discretised to two
decimal places, was randomly sampled from the experiment’s $\beta_{X}$
distribution and assigned to the job. By sampling a broad range of $\beta$
values from a selection of $\beta_{X}$ distributions, we ensured that we could
analyse the performance of each partitioning agent under different completion
time requirement settings and subsequently measure the capability of each
method to cater for different user-defined requirements.

| Heuristics | RL
---|---|---
| Random | Paramax | Paramin | PAC-ML
$\beta_{A}$ | $0.517^{+0.015}_{-0.015}$ | $0.262^{+0.002}_{-0.003}$ | $0.309^{+0.014}_{-0.015}$ | $\boldsymbol{0.203^{+0.007}_{-0.009}}$
$\beta_{B}$ | $0.601^{+0.007}_{-0.008}$ | $0.263^{+0.006}_{-0.004}$ | $0.396^{+0.006}_{-0.003}$ | $\boldsymbol{0.258^{+0.007}_{-0.003}}$
$\beta_{C}$ | $0.505^{+0.016}_{-0.012}$ | $0.267^{+0.004}_{-0.006}$ | $0.307^{+0.015}_{-0.012}$ | $\boldsymbol{0.117^{+0.003}_{-0.003}}$
$\beta_{D}$ | $0.465^{+0.004}_{-0.006}$ | $0.263^{+0.006}_{-0.004}$ | $0.142^{+0.027}_{-0.046}$ | $\boldsymbol{0.099^{+0.008}_{-0.007}}$
Table 1: Blocking rate performance of the partitioning agents on the four
$\beta$ distributions (best in bold). Results are given as the mean across $3$
seeds, and error bars denote the corresponding min-max confidence intervals.

Partitioner baselines. We considered three heuristic baseline partitioning
strategies. $(1)$ Most prior works partition a given job across as many
workers as are available up to a pre-defined environment maximum partition
degree (Khani et al., 2021; Wang et al., 2022). We refer to this strategy as
‘$\text{Para}_{max}$’. $(2)$ Given the low network overhead (see Figure 3) and
contentionless nature of RAMP, and given the operations’ linear split-compute
time dependency of our environment, a reasonable estimate for the completion
time of a job with sequential run time $\text{JCT}^{\text{seq}}$ distributed
across $u^{t}$ workers would be
$\text{JCT}\approx\frac{\text{JCT}^{\text{seq}}}{u^{t}}$. Therefore, in light
of our objective to minimise the user-defined blocking rate, we introduce a
new partitioning strategy, ‘$\text{Para}_{min}$’, which partitions the job up
to the estimated minimum amount of parallelisation needed to satisfy the job’s
completion time requirements, $u^{t}=\lceil\frac{1}{\beta}\rceil$ (i.e. the
estimated speed-up factor needed). $(3)$ For completeness, we also ran a
‘Random’ partitioning baseline, which selects a partitioning degree randomly
from amongst the number of available workers.

Metrics recorded. To measure the performance of our partitioning agents, we
recorded the following key metrics. $(1)$ User-defined blocking rate (which we
abbreviate to ‘blocking rate’): The fraction of arrived jobs which had their
completion time requirements met by the cluster. $(2)$ Offered throughput: The
total ‘information size’ of the original jobs (i.e. before partitioning was
applied) processed per unit time. Since the open-access PipeDream job profiles
used in our experiments did not contain per-operation flop/s (computational
load) information, we summed the jobs’ operation and dependency sizes
(measured in bytes (B)) to get the total ‘information size’ of each job. The
load rate could then be defined as the rate of job information arriving at the
cluster per unit time, and the corresponding offered throughput as the rate at
which this total job information was processed by the cluster. For a full list
of metric definitions, refer to Appendix 9.2.

Figure 7: Validation performances (higher is better) of each partitioning
agent evaluated across three seeds normalised with respect to the best-
performing partitioner in each $B_{X}$ environment.

##  7 PAC-ML Partitioning Results & Discussion

###  7.1 Performance of the PAC-ML Partitioner

Comparison to the baseline partitioners. To test the performance of each
partitioning agent under different completion time requirement settings, we
ran our experiments across four different $\beta$ distributions (see Section
6). We visualise the relative blocking rate and throughput performance
differences between the agents in Figure 7, where an agent’s ‘score’ is its
normalised performance relative to the best-performing agent with respect to a
given metric. We evaluate these scores as
$\text{score}_{\text{blocking}}=\big{(}\frac{\text{best\\_blocking\\_rate}}{\text{blocking\\_rate}}\big{)}$,
and
$\text{score}_{\text{throughput}}=\big{(}\frac{\text{throughput}}{\text{best\\_throughput}}\big{)}$
for each agent (refer to Appendix 9.8 for all raw metric values). As shown in
Table 1 and Figure 7, our PAC-ML agent achieved the best blocking rate across
all four $\beta$ distributions, beating its nearest rival by
$22.5\%,1.90\%,56.2\%,\text{ and }30.3\%$ for $\beta_{A,B,C,D}$ respectively.

Comparison amongst the baseline partitioners. Figure 7 visualises the
performance of the best PAC-ML agents on each of the four $\beta$ distribution
environments compared to the baseline heuristic performances. Interestingly,
the best baseline in terms of blocking rate for $\beta_{A,B,C}$ is
$\text{Para}_{max}$, but this switches to $\text{Para}_{min}$ for $\beta_{D}$.
On $\beta_{B}$, PAC-ML achieved roughly equivalent performance to Paramax by
learning that, on this $\beta$ demand distribution, maximum parallelisation
led to the lowest blocking rates. This shows that different partitioning
strategies have varying relative performances under different cluster
settings. A key advantage of PAC-ML is therefore that the question of which
partitioning strategy is best for a given environment need not be addressed by
sub-optimal hand-crafted heuristics or environment-specific hyperparameter
tuning. Instead, we have demonstrated in Table 1 and Figure 7 that PAC-ML can
automatically learn performant partitioning strategies in arbitrary
environment settings.

###  7.2 Analysis of the PAC-ML Partitioner

Offered throughput analysis. One risk of optimising only for the blocking rate
when training the PAC-ML agent is that it maximises the number of jobs
accepted by prioritising small low-information jobs at the cost of a sub-
optimal offered throughput; a key metric when measuring a cluster’s quality of
service to users. Figure 7 shows that the offered throughput improves with the
blocking rate, with the PAC-ML agent ultimately achieving the best throughput
across all four $\beta$ distributions.

Figure 8: Mean per-job blocking rates of the five job types considered for
each partitioning agent under each $\beta_{X}$ setting plotted against the
number of operations (ops.), number of dependencies (deps.), the total job
information size, and the sequential run time of the job were it ran on a
single device (JCTseq).

Bias analysis. An important question is whether there is any bias in the kinds
of jobs the PAC-ML agent learns to prioritise in order to minimise the
blocking rate. To investigate this, Figure 8 shows the blocking rate vs. the
original characteristics for each of the five jobs considered (see Appendix
9.5 for a summary of these characteristics) for each $\beta_{X}$ distribution
environment. The PAC-ML agent had little to no bias across the jobs relative
to the other partitioners, with all jobs attaining approximately the same
blocking rate. There was a slight bias towards the larger jobs with greater
sequential completion times and more information to process, which is likely
due to the fact that larger jobs occupy more resources and therefore
inherently become favoured over smaller jobs.

##  8 Conclusion & Further Work

In conclusion, we have introduced a new partitioning strategy called PAC-ML.
Leveraging rl and a gnn, PAC-ML learns to partition computation jobs
dynamically arriving at a cluster of machines such that the number of jobs
which meet arbitrary user-defined completion time requirements is maximised
without the need for hand-crafted heuristics or environment-dependent
hyperparameter tuning. We tested our partitioner on the recently proposed RAMP
optical architecture (Ottino et al., 2022) across four distributions of user-
defined completion time requirements, demonstrating up to $56.2\%$ lower
blocking rates relative to the canonical maximum parallelisation strategies
used by most prior works when partitioning five real deep learning jobs. We
hope that our work will spur a new avenue of research into developing
partitioning strategies for distributed computing. In this section, we outline
potentially interesting areas of further work.

Exceeding completion time expectations. In this work, we rewarded PAC-ML with
$+1$ for completing a job within the user-defined maximum acceptable
completion time and $-1$ for failing to do so. Although minimising the
blocking rate is crucial for users, it would also be desirable to minimise the
jct as much as possible. An interesting area of further study would therefore
be to incorporate this objective into the reward function, perhaps by
combining the jct speed-up factor or offered throughput with the blocking rate
via multi-objective rl (Hayes et al., 2022).

Real-world experiments. Our work has considered real open-access deep learning
computation graph profiles but on a simulated optical architecture. A natural
but significant next step would be to implement PAC-ML in a real distributed
cluster. An important question would be whether an agent trained in a
simulated environment would be capable of inferring in a real cluster at test
time, or if real-world training would be needed.

Generalisation to unseen environments. This study ran PAC-ML in an environment
which had the same load rate, $\beta$ distribution, cluster network size, and
job computation graphs at train and test time. An interesting research
question would be whether PAC-ML would be able to learn on one set (or a
distribution) of these parameters and then generalise to a new set at test
time, or if it would need to leverage existing or new state-of-the-art methods
in gnn (Knyazev et al., 2019; Garg et al., 2020; Fan et al., 2021) and rl
(Cobbe et al., 2019; Wang et al., 2020; Kirk et al., 2021) generalisation.

Robustness to stochastic inter-arrival times. In our experiments, we fixed the
inter-arrival rate in order to fix the load rate. However, real clusters have
variable inter-arrival times (Parsonson et al., 2022a). Handling highly
stochastic environments is a known challenge for rl (Mao et al., 2019), and
therefore presents an interesting future research avenue for PAC-ML.

Combining the virtual plane. In our work, we have considered the job
partitioning task in isolation of the job placement and scheduling tasks.
However, prior works have found the merging of the latter sub-tasks into a
single resource management problem beneficial to performance (Paliwal et al.,
2020). An interesting area of further work would be to combine PAC-ML into a a
single algorithm which handled job partitioning, placement, and scheduling via
methods such hierarchical rl (Barto and Mahadevan, 2003; Vezhnevets et al.,
2017; Mirhoseini et al., 2018a; Paliwal et al., 2020; Zhang et al., 2021) or
multi-agent rl (Foerster, 2018).

##  9 Appendix

###  9.1 Extended Background

####  9.1.1 Parallelisation

There are three main types of deep learning parallelism; data parallelism,
model parallelism, and hybrid parallelism.

Data parallelism. Data parallelism (Slotnick et al., 1962) is where an
identical copy of the dnn model is sent to each worker. The input training
data is parallelised by sampling a training batch, splitting it into non-
overlapping micro-batches, training each worker on its own micro-batch, and
updating the workers’ local model parameters using some method to synchronise
the gradients of the parameters with respect to the training loss after each
training iteration. This synchronisation step is commonly referred to as
AllReduce, and can be performed using various techniques. Data parallelism can
be applied to any dnn model regardless of its architecture, enables the use of
large data sets (which are crucial for scaling model performance (Hoffmann et
al., 2022)), and facilitates the use of large training batch sizes which can
lead to smoother and faster convergence. This is a form of weak scaling, where
the jct is decreased by reducing the total number of training iterations
needed via increasing the amount of data processed per iteration as the number
of workers is increased (Khani et al., 2021). However, it scales poorly for
large models with many parameters since all parameters must fit onto a single
worker and then be synchronised at the end of each training step, and has the
constraint that the training data must be i.i.d. in order for parameter
updates to be computed and summed across workers to attain the updated model
parameters.

Model parallelism. Model parallelism (Karakus et al., 2021) is where the dnn
model is partitioned (split) and a part of the model is sent to each worker.
In the dnn forward pass, a training batch is sampled, copied, and sent to each
worker which holds layer-$1$ of the dnn. The layer-$1$ worker(s) then compute
the layer-$1$ output(s) and forward them to the worker(s) which hold
layer-$2$, and so on. In the backward pass, the gradients of the model
parameters with respect to the training loss are computed by starting at the
worker(s) which hold the final layer and propagating these gradients back to
the layer-$1$ workers, after which the partitioned model will be globally
synchronised. Layer outputs, gradients, and activations are exchanged during
the training iteration using a synchronisation step commonly referred to as
AllGather. Model parallelism facilitates the use of very large models which
otherwise would not fit onto a single worker and caters for time-efficient
parallelisation of computational operations where possible. This is a form of
strong scaling, where the jct and per-worker memory utilisation are decreased
via increasingly partitioning different parts the job across more workers as
the number of workers is increased (Khani et al., 2021). However, passing
gradients between workers during training can create a large communication
overhead (Mirhoseini et al., 2017, 2018b), and expert domain knowledge of the
specific model architecture is needed to know how to split the model across
multiple workers.

Hybrid parallelism. Hybrid parallelism (Dean et al., 2012) is where a
combination of data and model parallelism is used to strive for the benefits
of both. This can be extended to include pipeline parallelism (Huang et al.,
2019; Narayanan et al., 2019), where intra-batch parallelism (data and model
parallelism) are combined with inter-batch parallelism (pipelining) where
multiple micro-batches are processed simultaneously where possible. Hybrid
parallelism can result in higher worker utilisation and the advantages of both
model and data parallelism, but requires complex bidirectional pipelining
across different inputs, careful model parameter versioning to ensure correct
computations of the gradients during the backward pass, and each stage
allocated across workers must be load-balanced to ensure roughly equivalent
computational times between workers in order to maximise peak pipeline
throughput.

####  9.1.2 Neural Networks as Function Approximators

Neural networks. Neural networks are a composition of linear and non-linear
(activation) functions connected in a chain to form a dag. Each function in
the chain is a layer parameterised by a set of weights and biases which, given
enough parameters, can be trained to approximate any universal function
(Hornik et al., 1989; Montufar et al., 2014). Neural networks with multiple
intermediary (hidden) layers between input and output are referred to as deep
neural networks and have powerful expressivity capabilities when approximating
complex non-linear functions (Hornik et al., 1989; Montufar et al., 2014).

Graph neural networks. Whereas standard dnn are restricted to handling only
vector- and grid-structured inputs (e.g. sentences, images, etc.), gnn are
generalised dnn architectures which can handle graph-structured data as inputs
(e.g. job dag). Most current gnn use the message passing paradigm by mapping
each node and edge onto a vector embedding space before performing additional
graph-level embeddings and readouts if desired.

Specifically, each gnn layer usually performs four stages: (i) On each edge in
the input graph, use a message function to generate a message (representation)
to pass from a source node to a set of destination nodes, where each node
stores the message(s) it receives in its mailbox; (ii) on each node in the
input graph, apply an aggregate function (a vanilla reduce operation such as
mean, sum, max, min, etc., or a trainable function) to the messages in its
mailbox to generate an intermediate aggregate representation of its
neighbourhood; (iii) pass the intermediate aggregate representation through a
trainable function to produce a final vector embedding for each node; and
(optional) (iv) if desired, at the end of the final GNN layer, pass the node
embeddings through a trainable function to produce a graph-level
representation. Crucially, the parameters of all message, aggregation, and
forward pass functions are shared across nodes, enabling gnn to be inductive
in that they can generalise to unseen nodes and graphs.

####  9.1.3 Reinforcement Learning Algorithm

Here we break down the key background components of the rl approach used for
PAC-ML.

Q-learning. Q-learning (Watkins, 1989) is the canonical value-based algorithm
which can be applied to a sequential decision making process formalised as an
mdp. It is an off-policy temporal difference algorithm whose goal is to learn
the action value function mapping state-action pairs to their expected
discounted future return when following a policy $\pi$;
$Q^{\pi}(s,u)=\mathbb{E}_{\pi}\big{[}\sum_{t^{\prime}=t+1}^{\infty}\gamma_{t^{\prime}-1}r(s_{t^{\prime}})|s_{t}{=}s,u_{t}{=}u\big{]}$.
By definition, an optimal policy $\pi_{*}$ will select an action which
maximises the true Q-value $Q_{*}(s,u)$,
$\pi_{*}(s)=\arg\max_{u^{\prime}}Q_{*}(s,u^{\prime})$.

Concretely, the classical Q-learning algorithm maintains an action value look-
up table $Q(s,u)$ mapping all possible state-action pairs to their predicted
discounted return. The return is the sum of future rewards over the remainder
of the episode. During training, Q-learning follows an exploration-
exploitation policy. The simplest such policy is $\epsilon$-greedy, where a
random action is sampled with probability $\epsilon\in[0,1]$ and the best
action, according to the current $Q$ table, is sampled with probably
$1-\epsilon$. At each time step $t$, the agent in state $s_{t}$ uses this
policy to select an action $u_{t}$ which it performs in the environment to
transition to the next state $s_{t+1}$ and receive a reward $r_{t+1}$.
$Q(s,u)$ is then updated according to:

$Q(s_{t},u_{t})\leftarrow
Q(s_{t},u_{t})+\alpha\cdot\bigg{(}r_{t}+\gamma\cdot\max_{u^{\prime}}Q(s_{t+1},u^{\prime})-Q(s_{t},u_{t})\bigg{)}.$
(2)

On the right-hand side of Eq. 2, $Q(s_{t},u_{t})$ is the agent’s estimate of
the discounted return of taking action $u_{t}$ in state $s_{t}$, $\alpha$ is
the learning rate, $\gamma$ is the factor by which to discount future rewards
to their present value, and $\max_{u^{\prime}}Q(s_{t+1},u^{\prime})$ is an
estimate of the future value of being in state $s_{t+1}$ and taking an
‘optimal’ action according to $Q$. The
$r_{t}+\gamma\cdot\max_{u^{\prime}}Q(s_{t+1},u^{\prime})$ term is called the
temporal difference target, and and the collective
$r_{t}+\gamma\cdot\max_{u^{\prime}}Q(s_{t+1},u^{\prime})-Q(s_{t},a_{t})$ term
the temporal difference error. As such, the
$\max_{u^{\prime}}Q(s_{t+1},u^{\prime})$ term treats $Q$ as an oracle from
which optimal actions can be sampled. Although $Q$ is usually randomly
initialised and changes at each update step, the general idea is that, with
stable learning and sufficient exploration, $Q$ will converge on the true
$Q_{*}$ function.

As a side note, Q-learning is a temporal difference algorithm because, rather
than using the actual returns to update $Q$ in Eq. 2 as done my Monte Carlo
methods, it uses a bootstrapped estimate of the future returns
$\max_{u^{\prime}}Q(s_{t+1},u^{\prime})$. Furthermore, it is an off-policy
algorithm because the policy used to select the action $u_{t}$ at the current
time step, such as $\epsilon$-greedy sampling of $Q$, is different to the
policy used to select the next-state action $u^{\prime}$ when evaluating the
temporal difference target, such as greedy sampling of $Q$. This is as opposed
to on-policy temporal difference algorithms, such as SARSA, which use the same
action selection policy for both the current time step and for future time
steps when bootstrapping.

Deep Q-learning. Many practical problems have an extremely large number of
possible state-action combinations. For example, the game of Go has over
$10^{700}$ possible sequences; far more than the number of atoms in the
universe (Silver et al., 2016). As such, modelling the action value function
with a tabular approach is intractable given practical memory constraints. To
enable Q-learning to be scaled to complex problems, dqn (Mnih et al., 2013)
approximates the true Q-function with a dnn parameterised by $\theta$ such
that $Q_{\theta}(s,u)\approx Q_{*}(s,u)$.

Concretely, during training at each time step $t$, $Q_{\theta}(s,u)$ is used
with an exploration strategy such as $\epsilon$-greedy to select an action and
add the observed transition $T=(s_{t},u_{t},r_{t+1},\gamma_{t+1},s_{t+1})$ to
a replay memory buffer (Lin, 1992). The network’s parameters $\theta$ are then
optimised with stochastic gradient descent to minimise the mean squared error
loss between the online network’s predictions and a bootstrapped estimate of
the Q-value,

$J_{DQN}(Q)=\big{[}r_{t+1}+\gamma_{t+1}\max_{u^{\prime}}Q_{\bar{\theta}}(s_{t+1},u^{\prime})-Q_{\theta}(s_{t},u_{t})\big{]}^{2},$
(3)

where $t$ is a time step uniform randomly sampled from the buffer and
$Q_{\bar{\theta}}$ a target network with parameters $\bar{\theta}$ which are
periodically copied from the acting online network. The target network is not
directly optimised, but is used to provide the bootstrapped Q-value estimates
for the loss function. Only periodically updating the target network rather
than at each learning step leads to lower variance in the bootstrapped targets
at each step. This helps helps to stabilise learning and leads to better
convergence (Mnih et al., 2013).

Double DQN. In the traditional Q-learning update rule of Eq. 2 and the DQN
loss of Eq. 3, the Q-function used to select and evaluate an action for the
temporal difference target is the same;
$\max_{u^{\prime}}Q(s_{t+1},u^{\prime})$ for Eq. 2, and
$\max_{u^{\prime}}Q_{\bar{\theta}}(s_{t+1},u^{\prime})$ for Eq. 3. However,
this can lead to an overestimation bias where the chosen action $u^{\prime}$
is incorrectly over-valued because the same function which perceives
$u^{\prime}$ as being best is also being asked to evaluate it. This can lead
to high variance updates, unstable learning, and convergence on local minima.
Double DQN (van Hasselt et al., 2015) reduces overestimation by decomposing
the max operation in the temporal difference target into action selection and
action evaluation and performing these two tasks with two separate networks.

Concretely, action $u^{\prime}$ is greedily selected according to the online
network $Q_{\theta}$ and evaluated with the separate target network
$Q_{\bar{\theta}}$. The loss term from Eq. 3 then becomes:

$J_{DDQN}(Q)=\big{[}r_{t+1}+\gamma_{t+1}Q_{\bar{\theta}}(s_{t+1},\max_{u^{\prime}}Q_{\theta}(s_{t+1},u^{\prime}))-Q_{\theta}(s_{t},u_{t})\big{]}^{2}.$
(4)

Prioritised experience replay. Vanilla DQN replay buffers are sampled
uniformly to obtain transitions for network updates. A preferable approach is
to more frequently sample transitions from which there is much to learn.
Prioritised experience replay (Schaul et al., 2016) deploys this intuition by
sampling transitions with probability $p_{t}$ proportional to the last
encountered absolute temporal difference error,

$p_{t}\propto|r_{t+1}+\gamma_{t+1}\max_{u^{\prime}}Q_{\bar{\theta}}(s_{t+1},u^{\prime})-Q_{\theta}(s_{t},u_{t})|^{\omega},$
(5)

where $\omega$ is a tuneable hyperparameter for shaping the probability
distribution. New transitions are added to the replay buffer with maximum
priority to ensure all experiences will be sampled at least once to have their
errors evaluated.

n-step Q-learning. Traditional Q-learning uses the target network’s greedy
action at the next step to bootstrap a Q-value estimate for the temporal
difference target. Alternatively, to improve learning speeds and help with
convergence (Sutton and Barto, 2018; Hessel et al., 2017), forward-view multi-
step targets can be used (Sutton and Barto, 2018), where the $n$-step
discounted return from state $s$ is

$r_{t}^{(n)}=\sum_{k=0}^{n-1}\gamma_{t}^{(k)}r_{t+k+1},$ (6)

resulting in an $n$-step DQN loss of

$J_{DQN_{n}}(Q)=\big{[}r^{(n)}_{t}+\gamma^{(n)}_{t}\max_{u^{\prime}}Q_{\bar{\theta}}(s_{t+n},u^{\prime})-Q_{\theta}(s_{t},u_{t})\big{]}^{2}.$
(7)

Dueling DQN. Traditional dqn approaches use a dnn architecture which is not
specific to rl. Subsequently, when learning the Q-function, the entire dnn
architecture must learn to estimate the state value and the action advantage
for each action in order to learn the state-action function $Q^{\pi}(s,u)$ of
being in state $s$, taking action $u$, and following policy $\pi$. However, in
many problems where bootstrapped Q-learning is applied, the most important
objective is to learn to estimate the value of each state rather than the
effect of each action for each state. This is especially true in environments
and individual states where future transitions are mainly influenced by
factors other than the agent’s actions.

Leveraging the insight that in many states it is unnecessary to estimate the
value of each action choice, Wang et al. (2015) developed a new dnn
architecture, termed ‘dueling DQN’, which is better suited to the Q-learning
task. Concretely, the dueling architecture uses the same core dnn as standard
DQN. However, rather than following the initial encoding with a single
sequence of fully connected layers to get a Q-value for each possible action
in the current state, dueling DQN uses two separate streams of fully connected
layers. One stream, parameterised by $\beta$, estimates the state value
function $V_{\theta,\beta}(s)$ (the estimated future discounted return of the
current state regardless of future actions taken), and the other stream,
parameterised by $\alpha$, estimates the relative action advantage function
$A_{\theta,\alpha}(s,u)$ (the relative difference in the future discounted
return of each action).

The outputs of the two streams are then combined via a special aggregation
function to recover the state-action value function $Q$. Crucially, $V(s)$ and
$A(s,u)$ must be combined into $Q(s,u)$ in such a way that they are
independently identifiable from the output $Q$ values alone in order for
backpropagation to be able to calculate the appropriate loss and weight
updates for the separate $V(s)$ and $A(s,u)$ streams. As such, a simple
$Q(s,u)=V(s)+A(s,u)$ aggregation function to get the Q-values from the two
streams does not suffice. Instead, the authors tried two different aggregation
schemes.

The first aggregation method subtracted the advantage of the maximum advantage
action from all advantages to make the argmax action’s advantage $0$ and the
rest $<0$,

$Q_{\theta,\alpha,\beta}=V_{\theta,\beta}(s)+\bigg{(}A_{\theta,\alpha}(s,u)-\max_{u^{\prime}}A_{\theta,\alpha}(s,u^{\prime})\bigg{)},$
(8)

thus enabling $V(s)$ to be recovered at the argmax action’s Q-value.

The second aggregation method subtracted the mean advantage from all action
advantages to centre the advantage values around $0$ (i.e. to have a mean of
$0$),

$Q_{\theta,\alpha,\beta}=V_{\theta,\beta}(s)+\bigg{(}A_{\theta,\alpha}(s,u)-\frac{1}{|A|}\sum_{u^{\prime}}A_{\theta,\alpha}(s,u^{\prime})\bigg{)}.$
(9)

This makes $V(s)$ recoverable from $Q(s,u)$ by estimating the $V(s)$ value
which, when subtracted from each $A(s,u)$ value, leads to a set of $A(s,u)$
values which have a mean of $0$. In practice, this second approach of using
the mean was found to lead to more stable learning since using a mean
operation resulted in lower variance targets between learning steps compared
to when a max operation was used.

As with standard Q-learning, the output of the dueling network is a set of
Q-values (one for each action), therefore no change to the underlying
algorithm other than a slight adjustment of the network architecture was
required. By decomposing the Q-function approximator in this way, dueling dqn
is able to attain superior policy evaluation in the presence of many similar-
value actions, and the authors demonstrated their architecture achieving
state-of-the-art performance on the Atari $2600$ games.

Ape-X DQN. Noting that state-of-the-art ml performance is often achieved with
more computation, more powerful models, and larger training data sets, Horgan
et al. (2018) proposed Ape-X; a parallelisation approach to off-policy
experience replay rl. Concretely, rather than using a single actor-learner
setup, Ape-X decouples acting from learning. It distributes many actors across
a set of CPU cores each with their own instance of the environment. Each actor
retains a copy of a dnn shared across actors which it uses for action
selection to accumulate experiences in parallel with other actors. These
experiences are then communicated to a central shared replay buffer, where a
single learner mounted on a GPU uses prioritised experience replay to sample
the most important experiences for learning. Learner sampling, gradient
computation, and network updates are done asynchronously with one another on
separate threads, as are the periodic updates made to the actors’ networks
with the latest shared learner network. By using multiple actors in parallel,
not only can orders of magnitude more transition data be attained for
learning, but also a broader diversity of experiences can be collected by
allocating a different exploration strategy to each actor and thereby avoid
local optima in difficult exploration and large state-action space settings.
For $N_{actors}$ distributed actors, Horgan et al. (2018) used a per-actor
$\epsilon$-greedy exploration strategy whereby each actor $i$ had a fixed
exploration probability
$\epsilon_{i}=\epsilon^{1+\frac{i}{N_{actors}-1}\cdot\alpha}$ where
$\epsilon=0.4$ and $\alpha=0.7$. The authors demonstrated their approach
achieving new state-of-the-art results on Atari in a fraction of the training
time of prior works.

###  9.2 Metric Definitions

Table 2 summarises the metric jargon used throughout our manuscript.

Metric |  Description
---|---
Job completion time |  Time between job arriving and being completed.
Sequential job completion time |  Time it would take to complete a job were its operations ran sequentially on a single device.
Maximum acceptable job completion time |  Maximum time allowed to complete a job.
Speed-up factor |  Factor difference between sequential job completion time and actual job completion time.
Network overhead |  Fraction of the job completion time spent communicating information between workers when no computation was taking place.
Blocking rate |  Fraction of the arrived jobs which were successfully serviced across a given period of time.
Job information size |  Summed sizes (in bytes) of a job’s operations and dependencies.
Cluster throughput |  Total partitioned job information processed per unit time by the cluster.
Offered throughput |  Total original job information processed per unit time by the cluster.
Load rate |  Amount of job information arriving at the cluster per unit time.
Job inter-arrival time |  Time between when two jobs arrived at the cluster.
Table 2: Descriptions of the various metrics referred to throughout this
manuscript.

###  9.3 Experimental Hardware

All environment simulations were ran on Intel Xeon ES-2660 CPUs, and all
learner network training and inference was done on either a V100 or an A100
GPU.

###  9.4 Additional Simulation Details

####  9.4.1 Code Structure

We built a core RAMP simulation environment which followed a Gym-like
interface (Brockman et al., 2016) but without inheriting from a Gym
environment object to allow additional flexibility. We then built a wrapper
‘job partitioning’ environment which did conform to the Gym interface but used
our core RAMP simulation environment to perform the internal RAMP simulation
logic. Our code base is publicly available at
https://github.com/cwfparsonson/ddls for further practical implementation
details.

####  9.4.2 Job Allocation Procedure

When a job arrives at the cluster, our environment uses the following ordered
sequence of task executions to allocate the job:

  1. 1.

Op. partitioning: Partition the job dag’s operations to attain a ‘partitioned’
job dag.

  2. 2.

Op. placement: Place the operations in the partitioned job dag onto a sub-set
of cluster workers.

  3. 3.

Op. scheduling: For each worker, schedule the priority of its placed
operations to resolve conflicts where $\geq 2$ operations are ready to be
executed at the same time.

  4. 4.

Dep. placement: Given the placed operations and the data dependencies which
must be exchanged between operations, place the dependencies onto cluster
communication links.

  5. 5.

Dep. scheduling: For each communication link, schedule the priority of its
placed dependencies to resolve conflicts where $\geq 2$ dependencies are ready
to be communicated at the same time.

####  9.4.3 Job Allocation Methods

Each of the above allocation procedure tasks can be performed by any
algorithm, heuristic, or learning agent. In our work, we use the following
methods:

  1. 1.

Op. partitioning: PAC-ML, Paramax, Paramin, or Random. See the main manuscript
for details.

  2. 2.

Op. placement: A first-fit heuristic customised for the requirements of RAMP.
See Section 9.4.4 below for details.

  3. 3.

Op. scheduling: Shortest remaining processing time (Cai et al., 2016; Alizadeh
et al., 2013; Hong et al., 2012). Given a set of operations placed on a
worker, the operation with the shortest remaining run time will have the
highest priority and therefore be executed first wherever two operations on
the same worker request to be executed at the same time.

  4. 4.

Dep. placement: Shortest path & first-fit. Given a set of operation
placements, for any dependencies which need to be transferred through the
network (i.e. for dependencies with size $>0$ and whose parent operation is
placed on a separate worker from the child operation), $(1)$ first-fit select
a path from the $k-$shortest path with available light channel(s), and $(2)$
first-fit select an available channel.

  5. 5.

Dep. scheduling: Shortest remaining processing time. Given a set of
dependencies placed on a communication link channel, the dependency with the
shortest remaining processing time (i.e. the lowest amount of information left
to be transferred) will have the highest priority and therefore be
communicated first wherever two dependencies on the same link channel request
to be transported at the same time.

####  9.4.4 First-Fit Operation Placement in RAMP

The original RAMP paper of Ottino et al. (2022) did not specify an operation
placement heuristic which conformed to the RAMP placement rules (see Section
3). Here, we propose a simple first-fit heuristic which conforms to these
rules whilst making the placement problem tractable for large cluster
networks.

The basic idea behind partitioning and placement in the scenario described in
this work is to exploit the network efficiencies of RAMP as much as possible.
In particular, this means maximising the use of RAMP’s highly efficient
collective operations. For a generic partitioned dag, in the backward pass,
collectives happen for each operation when weights/gradients are shared
between sub-operations. If both a parent and child operation are placed on the
same set of (RAMP symmetry adherent) workers, then when the parent
communicates its output to the child’s input in the forward pass this will
also constitute a collective operation. As such the placement heuristic
implemented here seeks to primarily maximise the amount that these two
conditions are encountered. Given some operation, $o$, that has been
partitioned into $N$ equal sub-operations, $o_{i}$ and needs to be placed, the
placement is handled as:

  1. 1.

If a parent of $o$ has been partitioned and placed across $N$ servers which
adhere to the RAMP symmetry conditions, and if these servers each have enough
memory to store $o_{i}$, then place $o$ across this set of $N$ servers. This
ensures collective operations can happen in both the forward and backward
pass.

  2. 2.

Otherwise, check if a set of $N$ workers can be found in the network that
adheres to the RAMP symmetry requirements. This is achieved by sliding the
various possible symmetric shapes over the topology until a suitable one (or
none) is found. This ensures collective operations in the backward pass only.

Allocating in this way ensures that every partitioned operation can exploit
RAMP’s efficient collective operation process on the backward pass, and where
possible can also exploit it on the forward pass when receiving information
from (one of) its parents.

####  9.4.5 Evaluating the job completion time

The time to complete each operation was taken from the real computation job
profiles of the dnn jobs considered (see Section 9.5). To calculate the
communication time of point-to-point information transfers and of the MPI
collectives, we used the equations and code of Ottino et al. (2022).

####  9.4.6 Possible Causes of a Job Being Blocked

A job is blocked when either $\text{JCT}>\beta\cdot\text{JCT}^{\text{seq}}$
(i.e. failing to meet user’s chosen JCT requirement) or when the cluster does
not have enough available resources to service the job. The possible causes of
this latter form of blocking are:

  * 1.

Prior jobs using up too many cluster resources when later jobs arrive;

  * 2.

the minimum operation run time quantum not being low enough to partition the
operations enough times to lead to the desired JCT;

  * 3.

mounted worker operation scheduling conflicts for partitioned operations
mounted on the same worker leading to longer run times, since one worker can
only execute one operation at a time; and

  * 4.

excessive communication overheads incurring from over-partitioning of the job.

###  9.5 Job Computation Graph Data Sets

All computation graphs used in our experiments were taken from the open-access
PipeDream computation graph data set (Narayanan et al., 2019). Figure 9 shows
a visualisation of the key computation graph characteristics for each neural
network model considered, where the numbers reported are for one training
iteration (i.e. one forward and backward pass through the model). Table 3
reports the same characteristics but in tabular form. Finally, for
completeness, Figure 10 shows the actual job dag of the models used.

Figure 9: Visualisation of the characteristics of the deep learning
computation graphs used for our experiments before partitioning. The bottom
left sub-figure contains the model colour code scheme for all other sub-
figures. The statistics shown are for the operations and dependencies which
need to be executed and satisfied to conduct one training iteration.
Therefore, to carry out $N_{iter}$ training steps, the computation graph would
need to be executed $N_{iter}$ times. Computation time units are reported in
seconds, and memory units in bytes.

Model | # ops. | $\text{JCT}^{\text{seq}}$ | Max. op. comp. time |  $\Sigma$ op. mem. | Max. op. mem. | Depth | # deps. |  $\Sigma$ dep. size | Max. dep. size
---|---|---|---|---|---|---|---|---|---
ResNet-18 | $142$ | $36\,668.35$ | 473.625 | $17.258\,66\text{\times}{10}^{9}$ | $0.822\,121\,2\text{\times}{10}^{9}$ | $60$ | $159$ | $18.733\,29\text{\times}{10}^{9}$ | $0.822\,083\,6\text{\times}{10}^{9}$
VGG-16 | $82$ | $34\,525.35$ | $113.330$ | $30.625\,30\text{\times}{10}^{9}$ | $1.644\,315\text{\times}{10}^{9}$ | $80$ | $83$ | $29.467\,06\text{\times}{10}^{9}$ | $1.644\,167\text{\times}{10}^{9}$
GNMT | $96$ | $4470.80$ | $15.88$ | $2.368\,447\text{\times}{10}^{9}$ | $3.269\,491\text{\times}{10}^{8}$ | $30$ | $117$ | $1.027\,801\text{\times}{10}^{9}$ | $0.194\,437\,1\text{\times}{10}^{9}$
SqueezeNet-10 | $136$ | $38\,000.15$ | $474.637$ | $24.962\,62\text{\times}{10}^{9}$ | $1.168\,007\text{\times}{10}^{9}$ | $102$ | $153$ | $27.910\,09\text{\times}{10}^{9}$ | $1.167\,950\text{\times}{10}^{9}$
AlexNet | $46$ | $36\,061.15$ | $635.902$ | $3.046\,234\text{\times}{10}^{9}$ | $0.198\,339\,6\text{\times}{10}^{9}$ | $44$ | $47$ | $2.422\,161\text{\times}{10}^{9}$ | $0.198\,246\,4\text{\times}{10}^{9}$
Table 3: Summary of the characteristics of the deep learning computation
graphs used for our experiments before partitioning. The statistics shown are
for the operations (‘ops.’) and dependencies (‘deps.’) which need to be
executed and satisfied to conduct one training iteration. Therefore, to carry
out $N_{iter}$ training steps, the computation graph would need to be executed
$N_{iter}$ times. Computation (‘comp.’) time units are reported in seconds,
and memory (‘mem.’) units in bytes.

Figure 10: Deep learning computation graphs used for our experiments before
partitioning. Each computation graph represents the operations and
dependencies which need to be executed and satisfied to conduct one forward
and one backward pass through the neural network. Therefore, to carry out
$N_{iter}$ training steps, the computation graph would need to be executed
$N_{iter}$ times.

###  9.6 Neural Network Architecture

Figure 11: Schematic of the dnn architecture with $|L|$ gnn layers used to
parameterise the policy of PAC-ML. The gnn is similar to that of GraphSAGE
with mean pooling (Hamilton et al., 2017). Each GNN layer $l\in L$ contains a
node, edge, and reduce dnn module and ultimately learns to create an embedded
representation for each node in a given job DAG. These per-node embeddings are
then passed, along with any global job, cluster, and action features, to a
readout module. The readout module ultimately generates scores for each
possible action, which enables an action to be selected following a given
exploration-exploitation policy being followed. For clarity, this figure only
shows the GNN embedding-generation process for node $1$. See accompanying text
for a detailed explanation of this architecture and the accompanying figure.

As shown in Fig. 11, we used a message passing gnn similar to GraphSAGE with
mean pooling (Hamilton et al., 2017) to parameterise the PAC-ML policy. Table
4 summarises the hyperparameters used for the components of this dnn. We note
that we did not perform extensive hyperparamter tuning on the GNN
architecture. Below is a detailed explanation of this architecture.

GNN. First, the GNN layer takes in the DAG’s node and edge features and
generates an embedding for each node and edge in the graph. Then, each local
node’s nearest neighbour ($1$-hop away) sends the local node a message
(‘message passing’) which is the neighbouring nodes’ embeddings concatenated
with their connected edges’ embeddings. These messages are stored in the local
node’s ‘mailbox’, which now contains information about the node’s
neighbourhood. To ensure consistent dimensioning with the received messages, a
dummy zero-padded edge embedding is concatenated with the local node’s
embedding. Next, the reduce module takes the local and message embeddings and
generates a reduced representation for each. Finally, to generate a layer-$l$
output embedding for the local node, the element-wise mean of the reduced
embeddings is taken (‘mean pooling’). Note that this embedding process is done
for each node in the DAG, but for clairty Fig. 11 only follows node $1$.

If $l<L$ (i.e. if this is not the last GNN layer), these final node embeddings
are used as new features for the original DAG’s nodes and are passed to the
next GNN layer. If $l\equiv L$, then the node embeddings are passed to the
readout module. Note that $(1)$ the node, edge, and reduce modules are shared
across the aforementioned operations within a given GNN layer when generating
node embeddings, but not across different GNN layers, and $(2)$ the
$l^{th}$-layer’s output node embeddings will contain information about the
node’s neighbourhood from up to $l$ hops away.

Readout. The readout module takes the GNN’s node embeddings and the job’s and
cluster’s global features as input. To convert the node-level embeddings of
the GNN into a representation of the overall job DAG, their element-wise mean
is taken. To generate an embedding capturing the global job, cluster, and
action information, a global dnn module is used. The DAG and global embeddings
are then concatenated and passed to a logit module, which in turn generates a
vector of (optionally masked) scores for each possible action in the
environment. Finally, based on these scores and the exploration-exploitation
policy being followed, an action is selected.

Parameter | Value
---|---
Message passing # hidden dimensions | 64
Message passing # output dimensions | 32
Reduce module # hidden dimensions | 64
Reduce module # output dimensions |  64 if $l<L$, else $16$
Global module # hidden dimensions | 8
Global module # output dimensions | 8
Logit module RLlib FC net # layers | 1
Logit module RLlib FC net # hidden dimensions | 256
All modules’ activation | ReLU
GNN # layers $L$ | 2
Apply action mask | False
Table 4: Hyperparamters used for the PAC-ML ApeX-DQN DNN policy architecture
shown in Fig. 11. Note that the ‘message passing’ dimensions refer to the
dimensions of the concatenated node and edge modules’ embeddings, so the
dimensions of these modules’ hidden and output embeddings will be half the
corresponding ‘message passing’ dimension. Due to the RLlib implementation of
Ape-X DQN, we did not apply an action mask, but instead included the action
mask in the global features given to the model and used the reward signal to
train the agent to avoid selecting invalid actions.

###  9.7 Reinforcement Learning Algorithm

Approach. Given the stochastic nature of our dynamic cluster environment
setting, we hypothesised that a value-based rl method would be best suited to
our setting (Mao et al., 2019). We did try the PPO (Schulman et al., 2017)
actor-critic method but found performance to be worse, although we leave a
full analysis of alternative rl algorithms to future work.

As stated in the main manuscript, we used the state-of-the-art value-based
Ape-X DQN rl algorithm (Horgan et al., 2018) to attain the PAC-ML policy.
Concretely, we used the Ape-X parallelisation approach with double Q-learning
action selection-evaluation (van Hasselt et al., 2015) and multi-step
bootstrapped learning targets (Sutton and Barto, 2018; Hessel et al., 2017),
prioritised experience replay (Schaul et al., 2016), a dueling dqn network
architecture (Wang et al., 2015), and a per-actor $\epsilon$-greedy
exploration algorithm. For a breakdown of each of these components, refer to
Appendix 9.1.3.

Hyperparameters. To select the algorithm hyperparameters, we conducted a
Bayesian search across the search space summarised in Table 5, with
simulations conducted in a light 32-worker RAMP environment with a maximum
simulation run time of $2\text{\times}{10}^{5}$ seconds to speed up the
search. We adopted similar search ranges to those used by Kurach et al.
(2019); Hoffman et al. (2020); Parsonson et al. (2022b). For each set of
hyperparameters, we ran the algorithm for $100$ learner steps (a.k.a. training
epochs), and performed a validation across $3$ seeds at each learner step (see
Figure 12). We selected the parameter set with the highest episode return
across the $3$ seeds (see Table 5). We also report the importance of each
parameter with respect to the total episode return. The importance is
calculated by training a random forest with all algorithm hyperparameters as
inputs and the episode return as the target output, with the per-feature
(hyperparameter) importance values predicted by random forest reported
accordingly (fab, 2018; how, 2018). All our experiments used the same per-
actor $\epsilon$-greedy exploration as Horgan et al. (2018).

Figure 12: Validation performance of the Ape-X DQN hyperparameter sweep. Each
agent was trained for $100$ learner steps, and at each learner step a
validation was performed across $3$ seeds - the mean metrics with their min-
max interval bands are plotted for each hyperparameter set.

Parameter | Search Range | Best Value | Importance
---|---|---|---
Discount factor $\gamma$ | $\\{0.99,0.993,0.997,0.999,0.9999\\}$ | $0.999$ | $0.004$
Learning rate |  Log-uniform values ( $1\text{\times}{10}^{-7}$, $1\text{\times}{10}^{-3}$ ) | $4.121\text{\times}{10}^{-7}$ | $0.045$
$v_{min}$ | $\\{-1,-10,-100,-200,-1000\\}$ | $-1000$ | $0.01$
$v_{max}$ | $\\{1,10,100,200,1000\\}$ | $1000$ | $0.004$
Target network update frequency |  $\\{$ $1\text{\times}{10}^{3}$, $1\text{\times}{10}^{4}$, $1\text{\times}{10}^{5}$ $\\}$ | $1\text{\times}{10}^{5}$ | $0.001$
Prioritised replay $\alpha$ | $\\{0.1,0.4,0.5,0.6,0.7,0.8,0.9\\}$ | $0.9$ | $0.04$
Prioritised replay $\beta$ | $\\{0.1,0.4,0.5,0.6,0.7,0.8,0.9\\}$ | $0.1$ | $0.047$
$n$-step | $\\{1,3,5,10\\}$ | $3$ | $0.227$
$\\#$ CPU workers | $32$ | $32$ | $-$
$\\#$ GPU workers | $1$ | $1$ | $-$
Batch mode | Truncated episodes | Truncated episodes | $-$
Rollout length | $50$ | $50$ | $-$
Train batch size | $512$ | $512$ | $-$
Optimiser | Adam | Adam | $-$
Dueling | True | True | $-$
$\\#$ atoms | $1$ | $1$ | $-$
Noisy | False | False | $-$
Double $Q$ | True | True | $-$
Replay buffer capacity | $100\,000$ | $100\,000$ | $-$
Learning starts | $10\,000$ | $10\,000$ | $-$
Prioritised replay TD-error $\epsilon$ | $1\text{\times}{10}^{-6}$ | $1\text{\times}{10}^{-6}$ | $-$
Table 5: Ape-X DQN training parameter sweep search range, best value found,
and corresponding parameter importance.

We note that our rl algorithms were implemented using the open-source RLlib
library (Liang et al., 2018) and hyperparameter tuning was done using Weights
& Biases (Biewald, 2020).

####  9.7.1 Final Learning Curves

For completeness, Figure 13 shows the learning curves of the tuned PAC-ML
agents in each $\beta_{X}$ environment superimposed on the baseline agents’
performances. At each learner step, the PAC-ML agent was evaluated across
three seeds in the validation environment.

Figure 13: Validation curves of the PAC-ML agent trained in four different
$\beta$ distribution environments. At each learner step (update to the GNN),
the agent was evaluated across 3 seeds, with the mean blocking rate, offered
throughput, jct, and jct speed-up (relative to the jobs’ sequential run time
$\text{JCT}^{\text{seq}}$) performance metrics reported as well as their min-
max confidence intervals. For reference, the performances of the baseline
heuristic partitioners are also plotted.

###  9.8 Additional Experimental Results

Figure 14: Validation performances of each partitioning agent evaluated across
three seeds, with the mean blocking rate, offered throughput, jct, and jct
speed-up (relative to the jobs’ sequential run time $\text{JCT}^{\text{seq}}$)
performance metrics reported.

Figure 14 shows the performance of the agents in terms of raw blocking rate,
throughput, jct, and jct speed-up.

##  10 Funding and Acknowledgments

## Funding

EPSRC Distributed Quantum Computing and Applications EP/W032643/1; the
Innovate UK Project on Quantum Data Centres and the Future 10004793; OptoCloud
EP/T026081/1; TRANSNET EP/R035342/1; the Engineering and Physical Sciences
Research Council EP/R041792/1 and EP/L015455/1; the Alan Turing Institute; and
Horizon Europe Dynamos.

## References

  * how (2018) , 2018.  Intro to machine learning: Lesson 4.  URL: https://www.youtube.com/watch?v=0v93qHDqq_g.
  * fab (2018) , 2018.  Introduction to hyperparameters.  URL: https://forums.fast.ai/t/wiki-lesson-thread-lesson-4/7540.
  * Addanki et al. (2019a) Addanki, R., Venkatakrishnan, S.B., Gupta, S., Mao, H., Alizadeh, M., 2019a.  Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learning. Curran Associates Inc., Red Hook, NY, USA.
  * Addanki et al. (2019b) Addanki, R., Venkatakrishnan, S.B., Gupta, S., Mao, H., Alizadeh, M., 2019b.  Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learning. Curran Associates Inc., Red Hook, NY, USA.
  * Alizadeh et al. (2013) Alizadeh, M., Yang, S., Sharif, M., Katti, S., McKeown, N., Prabhakar, B., Shenker, S., 2013.  Pfabric: Minimal near-optimal datacenter transport.  SIGCOMM Comput. Commun. Rev. 43, 435–446.  URL: https://doi.org/10.1145/2534169.2486031, doi:10.1145/2534169.2486031.
  * Ballani et al. (2020) Ballani, H., Costa, P., Behrendt, R., Cletheroe, D., Haller, I., Jozwik, K., Karinou, F., Lange, S., Shi, K., Thomsen, B., Williams, H., 2020.  Sirius: A flat datacenter network with nanosecond optical switching, in: Proceedings of the Annual Conference of the ACM Special Interest Group on Data Communication on the Applications, Technologies, Architectures, and Protocols for Computer Communication, Association for Computing Machinery, New York, NY, USA. p. 782–797.  URL: https://doi.org/10.1145/3387514.3406221, doi:10.1145/3387514.3406221.
  * Bao et al. (2018) Bao, Y., Peng, Y., Wu, C., Li, Z., 2018.  Online job scheduling in distributed machine learning clusters, in: IEEE INFOCOM 2018 - IEEE Conference on Computer Communications, IEEE Press. p. 495–503.  URL: https://doi.org/10.1109/INFOCOM.2018.8486422, doi:10.1109/INFOCOM.2018.8486422.
  * Barrett et al. (2019) Barrett, T.D., Clements, W.R., Foerster, J.N., Lvovsky, A., 2019\.  Exploratory combinatorial optimization with reinforcement learning.  Association for the Advancement of Artificial Intelligence .
  * Barrett et al. (2022) Barrett, T.D., Parsonson, C.W.F., Laterre, A., 2022.  Learning to solve combinatorial graph partitioning problems via efficient exploration.  arXiv preprint arXiv:2205.14105 URL: https://arxiv.org/abs/2205.14105, doi:10.48550/ARXIV.2205.14105.
  * Barto and Mahadevan (2003) Barto, A.G., Mahadevan, S., 2003\.  Recent advances in hierarchical reinforcement learning.  Discrete Event Dynamic Systems 13, 341–379.
  * Bello* et al. (2017) Bello*, I., Pham*, H., Le, Q.V., Norouzi, M., Bengio, S., 2017\.  Neural combinatorial optimization with reinforcement learning.  URL: https://openreview.net/forum?id=rJY3vK9eg.
  * Ben-Nun and Hoefler (2019) Ben-Nun, T., Hoefler, T., 2019\.  Demystifying parallel and distributed deep learning: An in-depth concurrency analysis.  ACM Comput. Surv. 52.  URL: https://doi.org/10.1145/3320060, doi:10.1145/3320060.
  * Bengio et al. (2021) Bengio, Y., Lodi, A., Prouvost, A., 2021.  Machine learning for combinatorial optimization: A methodological tour d’horizon.  European Journal of Operational Research 290, 405–421.  URL: https://www.sciencedirect.com/science/article/pii/S0377221720306895, doi:https://doi.org/10.1016/j.ejor.2020.07.063.
  * Benjamin et al. (2020) Benjamin, J.L., Gerard, T., Lavery, D., Bayvel, P., Zervas, G., 2020.  Pulse: Optical circuit switched data center architecture operating at nanosecond timescales.  Journal of Lightwave Technology 38, 4906–4921.  doi:10.1109/JLT.2020.2997664.
  * Benjamin et al. (2022) Benjamin, J.L., Ottino, A., Parsonson, C.W.F., Zervas, G., 2022\.  Traffic tolerance of nanosecond scheduling on optical circuit switched data center network, in: 2022 Optical Fiber Communications Conference and Exhibition (OFC), pp. 1–3.
  * Benjamin et al. (2021) Benjamin, J.L., Parsonson, C.W.F., Zervas, G., 2021.  Benchmarking packet-granular ocs network scheduling for data center traffic traces, in: OSA Advanced Photonics Congress 2021, Optica Publishing Group. p. NeW3B.3.  URL: http://opg.optica.org/abstract.cfm?URI=Networks-2021-NeW3B.3, doi:10.1364/NETWORKS.2021.NeW3B.3.
  * Bergman (2018) Bergman, K., 2018.  Empowering Flexible and Scalable High Performance Architectures with Embedded Photonics.  IPDPS .
  * Biewald (2020) Biewald, L., 2020.  Experiment tracking with weights and biases.  URL: https://www.wandb.com/. software available from wandb.com.
  * Brockman et al. (2016) Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang, J., Zaremba, W., 2016.  Openai gym.  arXiv preprint arXiv:1606.01540 .
  * Cai et al. (2016) Cai, C.X., Saeed, S., Gupta, I., Campbell, R.H., Le, F., 2016.  Phurti: Application and network-aware flow scheduling for multi-tenant mapreduce clusters, in: 2016 IEEE International Conference on Cloud Engineering (IC2E), pp. 161–170.  doi:10.1109/IC2E.2016.21.
  * Cobbe et al. (2019) Cobbe, K., Klimov, O., Hesse, C., Kim, T., Schulman, J., 2019.  Quantifying generalization in reinforcement learning, in: Chaudhuri, K., Salakhutdinov, R. (Eds.), Proceedings of the 36th International Conference on Machine Learning, PMLR. pp. 1282–1289.  URL: https://proceedings.mlr.press/v97/cobbe19a.html.
  * Dai et al. (2017) Dai, H., Khalil, E.B., Zhang, Y., Dilkina, B., Song, L., 2017.  Learning combinatorial optimization algorithms over graphs, in: Proceedings of the 31st International Conference on Neural Information Processing Systems, Curran Associates Inc., Red Hook, NY, USA. p. 6351–6361.
  * Dean et al. (2012) Dean, J., Corrado, G.S., Monga, R., Chen, K., Devin, M., Le, Q.V., Mao, M.Z., Ranzato, M., Senior, A., Tucker, P., Yang, K., Ng, A.Y., 2012.  Large scale distributed deep networks, in: Proceedings of the 25th International Conference on Neural Information Processing Systems - Volume 1, Curran Associates Inc., Red Hook, NY, USA. p. 1223–1231.
  * Dong et al. (2020) Dong, K., Luo, Y., Yu, T., Finn, C., Ma, T., 2020\.  On the expressivity of neural networks for deep reinforcement learning, in: Proceedings of the 37th International Conference on Machine Learning, JMLR.org.
  * Fan et al. (2021) Fan, S., Wang, X., Shi, C., Cui, P., Wang, B., 2021\.  Generalizing graph neural networks on out-of-distribution graphs.  URL: https://arxiv.org/abs/2111.10657, doi:10.48550/ARXIV.2111.10657.
  * Foerster (2018) Foerster, J., 2018.  Deep multi-agent reinforcement learning.  Ph.D. thesis. University of Oxford.
  * Furber (2016) Furber, S., 2016.  Large-scale neuromorphic computing systems.  Journal of Neural Engineering 13, 051001.  URL: https://doi.org/10.1088/1741-2560/13/5/051001, doi:10.1088/1741-2560/13/5/051001.
  * Gao et al. (2018) Gao, Y., Chen, L., Li, B., 2018.  Spotlight: Optimizing device placement for training deep neural networks, in: Dy, J., Krause, A. (Eds.), Proceedings of the 35th International Conference on Machine Learning, PMLR. pp. 1676–1684.  URL: https://proceedings.mlr.press/v80/gao18a.html.
  * Garg et al. (2020) Garg, V., Jegelka, S., Jaakkola, T., 2020.  Generalization and representational limits of graph neural networks, in: III, H.D., Singh, A. (Eds.), Proceedings of the 37th International Conference on Machine Learning, PMLR. pp. 3419–3430.  URL: https://proceedings.mlr.press/v119/garg20c.html.
  * Gasse et al. (2019) Gasse, M., Chételat, D., Ferroni, N., Charlin, L., Lodi, A., 2019.  Exact combinatorial optimization with graph convolutional neural networks, in: Advances in Neural Information Processing Systems 32.
  * Gerard et al. (2020) Gerard, T., Parsonson, C., Shabka, Z., Bayvel, P., Lavery, D., Zervas, G., 2020\.  Swift: Scalable ultra-wideband sub-nanosecond wavelength switching for data centre networks.  arXiv URL: https://arxiv.org/abs/2003.05489, doi:10.48550/ARXIV.2003.05489, arXiv:arXiv:2003.05489.
  * Gerard et al. (2021) Gerard, T., Parsonson, C., Shabka, Z., Thomsen, B., Bayvel, P., Lavery, D., Zervas, G., 2021.  Ai-optimised tuneable sources for bandwidth-scalable, sub-nanosecond wavelength switching.  Opt. Express 29, 11221–11242.  URL: http://opg.optica.org/oe/abstract.cfm?URI=oe-29-7-11221, doi:10.1364/OE.417272.
  * Halim and Ismail (2019) Halim, A.H., Ismail, I., 2019\.  Combinatorial optimization: Comparison of heuristic algorithms in travelling salesman problem.  Archives of Computational Methods in Engineering 26, 367–380.
  * Hamilton et al. (2017) Hamilton, W.L., Ying, R., Leskovec, J., 2017.  Inductive representation learning on large graphs.  URL: https://arxiv.org/abs/1706.02216, doi:10.48550/ARXIV.1706.02216.
  * van Hasselt et al. (2015) van Hasselt, H., Guez, A., Silver, D., 2015.  Deep reinforcement learning with double q-learning arXiv:1509.06461.
  * Hayes et al. (2022) Hayes, C.F., Radulescu, R., Bargiacchi, E., Källström, J., Macfarlane, M., Reymond, M., Verstraeten, T., Zintgraf, L.M., Dazeley, R., Heintz, F., Howley, E., Irissappane, A.A., Mannion, P., Nowé, A., de Oliveira Ramos, G., Restelli, M., Vamplew, P., Roijers, D.M., 2022\.  A practical guide to multi-objective reinforcement learning and planning.  Auton. Agents Multi Agent Syst. 36, 26.  URL: https://doi.org/10.1007/s10458-022-09552-y, doi:10.1007/s10458-022-09552-y.
  * He et al. (2016) He, K., Zhang, X., Ren, S., Sun, J., 2016.  Deep residual learning for image recognition, in: 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770–778.  doi:10.1109/CVPR.2016.90.
  * Hessel et al. (2017) Hessel, M., Modayil, J., van Hasselt, H., Schaul, T., Ostrovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M., Silver, D., 2017\.  Rainbow: Combining improvements in deep reinforcement learning.  arXiv arXiv:1710.02298.
  * Hoffman et al. (2020) Hoffman, M.W., Shahriari, B., Aslanides, J., Barth-Maron, G., Momchev, N., Sinopalnikov, D., Stanczyk, P., Ramos, S., Raichuk, A., Vincent, D., Hussenot, L., Dadashi, R., Dulac-Arnold, G., Orsini, M., Jacq, A., Ferret, J., Vieillard, N., Ghasemipour, S.K.S., Girgin, S., Pietquin, O., Behbahani, F., Norman, T., Abdolmaleki, A., Cassirer, A., Yang, F., Baumli, K., Henderson, S., Friesen, A., Haroun, R., Novikov, A., Colmenarejo, S.G., Cabi, S., Gulcehre, C., Paine, T.L., Srinivasan, S., Cowie, A., Wang, Z., Piot, B., de Freitas, N., 2020.  Acme: A research framework for distributed reinforcement learning URL: https://arxiv.org/abs/2006.00979, doi:10.48550/ARXIV.2006.00979, arXiv:arXiv:2006.00979.
  * Hoffmann et al. (2022) Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D.d.L., Hendricks, L.A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., Driessche, G.v.d., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Rae, J.W., Vinyals, O., Sifre, L., 2022\.  Training compute-optimal large language models URL: https://arxiv.org/abs/2203.15556, doi:10.48550/ARXIV.2203.15556, arXiv:arXiv:2203.15556.
  * Hong et al. (2012) Hong, C.Y., Caesar, M., Godfrey, P.B., 2012.  Finishing flows quickly with preemptive scheduling, in: Proceedings of the ACM SIGCOMM 2012 Conference on Applications, Technologies, Architectures, and Protocols for Computer Communication, Association for Computing Machinery, New York, NY, USA. p. 127–138.  URL: https://doi.org/10.1145/2342356.2342389, doi:10.1145/2342356.2342389.
  * Hopfield and Tank (1985) Hopfield, J., Tank, D., 1985\.  "neural" computation of decisions in optimization problems.  Biological Cybernetics 52, 141–152.  doi:10.1007/BF00339943.
  * Horgan et al. (2018) Horgan, D., Quan, J., Budden, D., Barth-Maron, G., Hessel, M., van Hasselt, H., Silver, D., 2018.  Distributed prioritized experience replay, in: International Conference on Learning Representations.  URL: https://openreview.net/forum?id=H1Dy---0Z.
  * Hornik et al. (1989) Hornik, K., Stinchcombe, M., White, H., 1989.  Multilayer feedforward networks are universal approximators.  Neural Networks 2, 359–366.
  * Huang et al. (2019) Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, M.X., Chen, D., Lee, H., Ngiam, J., Le, Q.V., Wu, Y., Chen, Z., 2019\.  GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism. Curran Associates Inc., Red Hook, NY, USA.
  * Iandola et al. (2016) Iandola, F.N., Han, S., Moskewicz, M.W., Ashraf, K., Dally, W.J., Keutzer, K., 2016\.  Squeezenet: Alexnet-level accuracy with 50x fewer parameters and <0.5mb model size.  arXiv preprint arXiv:1602.07360 URL: http://arxiv.org/abs/1602.07360. cite arxiv:1602.07360Comment: In ICLR Format.
  * Kaplan et al. (2020) Kaplan, J., McCandlish, S., Henighan, T., Brown, T.B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., Amodei, D., 2020.  Scaling laws for neural language models.  CoRR abs/2001.08361.  URL: https://arxiv.org/abs/2001.08361, arXiv:arXiv:2001.08361.
  * Karakus et al. (2021) Karakus, C., Huilgol, R., Wu, F., Subramanian, A., Daniel, C., Cavdar, D., Xu, T., Chen, H., Rahnama, A., Quintela, L., 2021\.  Amazon sagemaker model parallelism: A general and flexible framework for large model training URL: https://arxiv.org/abs/2111.05972, doi:10.48550/ARXIV.2111.05972, arXiv:arXiv:2111.05972.
  * Khadka et al. (2021) Khadka, S., Aflalo, E., Mardar, M., Ben-David, A., Miret, S., Mannor, S., Hazan, T., Tang, H., Majumdar, S., 2021.  Optimizing memory placement using evolutionary graph reinforcement learning, in: International Conference on Learning Representations.  URL: https://openreview.net/forum?id=-6vS_4Kfz0.
  * Khani et al. (2021) Khani, M., Ghobadi, M., Alizadeh, M., Zhu, Z., Glick, M., Bergman, K., Vahdat, A., Klenk, B., Ebrahimi, E., 2021.  Sip-ml: High-bandwidth optical network interconnects for machine learning training, in: Proceedings of the 2021 ACM SIGCOMM 2021 Conference, Association for Computing Machinery, New York, NY, USA. p. 657–675.  URL: https://doi.org/10.1145/3452296.3472900, doi:10.1145/3452296.3472900.
  * Kirk et al. (2021) Kirk, R., Zhang, A., Grefenstette, E., Rocktäschel, T., 2021\.  A survey of generalisation in deep reinforcement learning.  URL: https://arxiv.org/abs/2111.09794, doi:10.48550/ARXIV.2111.09794.
  * Knyazev et al. (2019) Knyazev, B., Taylor, G.W., Amer, M.R., 2019.  Understanding Attention and Generalization in Graph Neural Networks. Curran Associates Inc., Red Hook, NY, USA.
  * Krizhevsky et al. (2012) Krizhevsky, A., Sutskever, I., Hinton, G.E., 2012.  Imagenet classification with deep convolutional neural networks, in: Proceedings of the 25th International Conference on Neural Information Processing Systems - Volume 1, Curran Associates Inc., Red Hook, NY, USA. p. 1097–1105.
  * Kurach et al. (2019) Kurach, K., Raichuk, A., Stanczyk, P., Zając, M., Bachem, O., Espeholt, L., Riquelme, C., Vincent, D., Michalski, M., Bousquet, O., Gelly, S., 2019.  Google research football: A novel reinforcement learning environment URL: https://arxiv.org/abs/1907.11180, doi:10.48550/ARXIV.1907.11180, arXiv:arXiv:1907.11180.
  * Li et al. (2021) Li, Q., Xu, J., Cao, C., 2021\.  Scheduling distributed deep learning jobs in heterogeneous cluster with placement awareness, in: Proceedings of the 12th Asia-Pacific Symposium on Internetware, Association for Computing Machinery, New York, NY, USA. p. 217–228.  URL: https://doi.org/10.1145/3457913.3457936, doi:10.1145/3457913.3457936.
  * Liang et al. (2018) Liang, E., Liaw, R., Nishihara, R., Moritz, P., Fox, R., Goldberg, K., Gonzalez, J., Jordan, M., Stoica, I., 2018.  RLlib: Abstractions for distributed reinforcement learning, in: Dy, J., Krause, A. (Eds.), Proceedings of the 35th International Conference on Machine Learning, PMLR. pp. 3053–3062.  URL: https://proceedings.mlr.press/v80/liang18b.html.
  * Lin (1992) Lin, L.J., 1992.  Self-improving reactive agents based on reinforcement learning, planning and teaching.  Mach. Learn. 8, 293–321.  URL: https://doi.org/10.1007/BF00992699, doi:10.1007/BF00992699.
  * Mao et al. (2016) Mao, H., Alizadeh, M., Menache, I., Kandula, S., 2016\.  Resource management with deep reinforcement learning, in: Proceedings of the 15th ACM Workshop on Hot Topics in Networks, Association for Computing Machinery, New York, NY, USA. p. 50–56.  URL: https://doi.org/10.1145/3005745.3005750, doi:10.1145/3005745.3005750.
  * Mao et al. (2019) Mao, H., Venkatakrishnan, S.B., Schwarzkopf, M., Alizadeh, M., 2019\.  Variance reduction for reinforcement learning in input-driven environments, in: 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, OpenReview.net.  URL: https://openreview.net/forum?id=Hyg1G2AqtQ.
  * Mayer and Jacobsen (2020) Mayer, R., Jacobsen, H.A., 2020\.  Scalable deep learning on distributed infrastructures: Challenges, techniques, and tools.  ACM Comput. Surv. 53.  URL: https://doi.org/10.1145/3363554, doi:10.1145/3363554.
  * Mirhoseini et al. (2018a) Mirhoseini, A., Goldie, A., Pham, H., Steiner, B., Le, Q.V., Dean, J., 2018a.  A hierarchical model for device placement, in: International Conference on Learning Representations.  URL: https://openreview.net/forum?id=Hkc-TeZ0W.
  * Mirhoseini et al. (2018b) Mirhoseini, A., Goldie, A., Pham, H., Steiner, B., Le, Q.V., Dean, J., 2018b.  A hierarchical model for device placement, in: International Conference on Learning Representations.  URL: https://openreview.net/forum?id=Hkc-TeZ0W.
  * Mirhoseini et al. (2017) Mirhoseini, A., Pham, H., Le, Q.V., Steiner, B., Larsen, R., Zhou, Y., Kumar, N., Norouzi, M., Bengio, S., Dean, J., 2017\.  Device placement optimization with reinforcement learning, in: Proceedings of the 34th International Conference on Machine Learning - Volume 70, JMLR.org. p. 2430–2439.
  * Mishra et al. (2021) Mishra, V., Benjamin, J.L., Zervas, G., 2021.  Monet: heterogeneous memory over optical network for large-scale data center resource disaggregation.  Journal of Optical Communications and Networking 13, 126–139.  doi:10.1364/JOCN.419145.
  * Mnih et al. (2013) Mnih, V., Kavukcuoglu, K., Silver, D., Graves, A., Antonoglou, I., Wierstra, D., Riedmiller, M., 2013.  Playing atari with deep reinforcement learning arXiv:arXiv:1312.5602.
  * Montufar et al. (2014) Montufar, G., Pascanu, R., Cho, K., Bengio, Y., 2014\.  On the number of linear regions of deep neural networks, in: Proceedings of the 27th International Conference on Neural Information Processing Systems - Volume 2, MIT Press, Cambridge, MA, USA. p. 2924–2932.
  * Narayanan et al. (2019) Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N., Granger, G., Gibbons, P., Zaharia, M., 2019\.  Pipedream: Generalized pipeline parallelism for dnn training, in: ACM Symposium on Operating Systems Principles (SOSP 2019).  URL: https://www.microsoft.com/en-us/research/publication/pipedream-generalized-pipeline-parallelism-for-dnn-training/.
  * Narayanan et al. (2021) Narayanan, D., Phanishayee, A., Shi, K., Chen, X., Zaharia, M., 2021.  Memory-efficient pipeline-parallel dnn training, in: International Conference on Machine Learning (ICML 2021).  URL: https://www.microsoft.com/en-us/research/publication/memory-efficient-pipeline-parallel-dnn-training/.
  * NVIDIA (2020) NVIDIA, 2020.  Nvidia selene: Leadership-class supercomputing infrastructure.  https://www.nvidia.com/en-us/on-demand/session/supercomputing2020-sc2019/.
  * NVIDIA (2022) NVIDIA, 2022.  Nvidia ai platform delivers big gains for large language models.  https://developer.nvidia.com/blog/nvidia-ai-platform-delivers-big-gains-for-large-language-models/.
  * OpenAI (2018) OpenAI, 2018.  Ai and compute.  https://openai.com/blog/ai-and-compute/.
  * Ottino et al. (2022) Ottino, A., Benjamin, J., Zervas, G., 2022.  Ramp: A flat nanosecond optical network and mpi operations for distributed deep learning systems.  arXiv URL: https://arxiv.org/abs/2211.15226, doi:10.48550/ARXIV.2211.15226.
  * Paliwal et al. (2020) Paliwal, A., Gimeno, F., Nair, V., Li, Y., Lubin, M., Kohli, P., Vinyals, O., 2020\.  Reinforced genetic algorithm learning for optimizing computation graphs, in: 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020\.  URL: https://openreview.net/forum?id=rkxDoJBYPB.
  * Parsonson et al. (2022a) Parsonson, C.W., Benjamin, J.L., Zervas, G., 2022a.  Traffic generation for benchmarking data centre networks.  Optical Switching and Networking 46, 100695.
  * Parsonson et al. (2022b) Parsonson, C.W.F., Laterre, A., Barrett, T.D., 2022b.  Reinforcement learning for branch-and-bound optimisation using retrospective trajectories URL: https://arxiv.org/abs/2205.14345, doi:10.48550/ARXIV.2205.14345, arXiv:arXiv:2205.14345.
  * Parsonson et al. (2020) Parsonson, C.W.F., Shabka, Z., Chlupka, W.K., Goh, B., Zervas, G., 2020.  Optimal control of soas with artificial intelligence for sub-nanosecond optical switching.  Journal of Lightwave Technology 38, 5563–5573.  doi:10.1109/JLT.2020.3004645.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., Chintala, S., 2019.  Pytorch: An imperative style, high-performance deep learning library, in: Advances in Neural Information Processing Systems 32. Curran Associates, Inc., pp. 8024–8035.  URL: http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf.
  * Raja et al. (2021) Raja, A.S., Lange, S., Karpov, M., Shi, K., Fu, X., Behrendt, R., Cletheroe, D., Lukashchuk, A., Haller, I., Karinou, F., Thomsen, B., Jozwik, K., Liu, J., Costa, P., Kippenberg, T.J., Ballani, H., 2021.  Ultrafast optical circuit switching for data centers using integrated soliton microcombs.  Nature Communications 12, 5867\.  URL: https://doi.org/10.1038/s41467-021-25841-8, doi:10.1038/s41467-021-25841-8.
  * Schaul et al. (2016) Schaul, T., Quan, J., Antonoglou, I., Silver, D., 2016\.  Prioritized experience replay arXiv:1511.05952.
  * Schulman et al. (2017) Schulman, J., Wolski, F., Dhariwal, P., Radford, A., Klimov, O., 2017.  Proximal policy optimization algorithms.  CoRR abs/1707.06347.  URL: http://dblp.uni-trier.de/db/journals/corr/corr1707.html#SchulmanWDRK17.
  * Silver et al. (2016) Silver, D., Huang, A., Maddison, C.J., Guez, A., Sifre, L., van den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., Dieleman, S., Grewe, D., Nham, J., Kalchbrenner, N., Sutskever, I., Lillicrap, T., Leach, M., Kavukcuoglu, K., Graepel, T., Hassabis, D., 2016\.  Mastering the game of Go with deep neural networks and tree search.  Nature 529, 484–489.  doi:10.1038/nature16961.
  * Simonyan and Zisserman (2014) Simonyan, K., Zisserman, A., 2014\.  Very deep convolutional networks for large-scale image recognition.  CoRR abs/1409.1556.  URL: http://arxiv.org/abs/1409.1556.
  * Slotnick et al. (1962) Slotnick, D.L., Borck, W.C., McReynolds, R.C., 1962.  The solomon computer, in: Proceedings of the December 4-6, 1962, Fall Joint Computer Conference, Association for Computing Machinery, New York, NY, USA. p. 97–107.  URL: https://doi.org/10.1145/1461518.1461528, doi:10.1145/1461518.1461528.
  * Smith et al. (2022) Smith, S., Patwary, M., Norick, B., LeGresley, P., Rajbhandari, S., Casper, J., Liu, Z., Prabhumoye, S., Zerveas, G., Korthikanti, V., Zhang, E., Child, R., Aminabadi, R.Y., Bernauer, J., Song, X., Shoeybi, M., He, Y., Houston, M., Tiwary, S., Catanzaro, B., 2022\.  Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model.  arXiv URL: https://arxiv.org/abs/2201.11990, doi:10.48550/ARXIV.2201.11990, arXiv:arXiv:2201.11990.
  * Sutton and Barto (2018) Sutton, R.S., Barto, A.G., 2018\.  Reinforcement Learning: An Introduction.  Second ed., The MIT Press.  URL: http://incompleteideas.net/book/the-book-2nd.html.
  * Vezhnevets et al. (2017) Vezhnevets, A.S., Osindero, S., Schaul, T., Heess, N., Jaderberg, M., Silver, D., Kavukcuoglu, K., 2017.  FeUdal networks for hierarchical reinforcement learning, in: Precup, D., Teh, Y.W. (Eds.), Proceedings of the 34th International Conference on Machine Learning, PMLR. pp. 3540–3549.  URL: https://proceedings.mlr.press/v70/vezhnevets17a.html.
  * Wang et al. (2020) Wang, K., Kang, B., Shao, J., Feng, J., 2020.  Improving generalization in reinforcement learning with mixture regularization.  URL: https://arxiv.org/abs/2010.10814, doi:10.48550/ARXIV.2010.10814.
  * Wang et al. (2019) Wang, M., Zheng, D., Ye, Z., Gan, Q., Li, M., Song, X., Zhou, J., Ma, C., Yu, L., Gai, Y., Xiao, T., He, T., Karypis, G., Li, J., Zhang, Z., 2019.  Deep graph library: A graph-centric, highly-performant package for graph neural networks URL: https://arxiv.org/abs/1909.01315, doi:10.48550/ARXIV.1909.01315.
  * Wang et al. (2022) Wang, W., Khazraee, M., Zhong, Z., Jia, Z., Mudigere, D., Zhang, Y., Kewitsch, A., Ghobadi, M., 2022\.  Topoopt: Optimizing the network topology for distributed dnn training.  arXiv preprint arXiv:2202.00433 arXiv:arXiv:2202.00433.
  * Wang et al. (2015) Wang, Z., Schaul, T., Hessel, M., van Hasselt, H., Lanctot, M., de Freitas, N., 2015\.  Dueling network architectures for deep reinforcement learning.  URL: http://arxiv.org/abs/1511.06581. cite arxiv:1511.06581Comment: 15 pages, 5 figures, and 5 tables.
  * Watkins (1989) Watkins, C.J.C.H., 1989.  Learning from Delayed Rewards.  Ph.D. thesis. King’s College, Oxford.
  * Williams (1992) Williams, R.J., 1992.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  Mach. Learn. 8, 229–256.  URL: https://doi.org/10.1007/BF00992696, doi:10.1007/BF00992696.
  * Williamson and Shmoys (2011) Williamson, D.P., Shmoys, D.B., 2011\.  The Design of Approximation Algorithms.  Cambridge University Press.
  * Wu et al. (2016) Wu, Y., Schuster, M., Chen, Z., Le, Q.V., Norouzi, M., Macherey, W., Krikun, M., Cao, Y., Gao, Q., Macherey, K., Klingner, J., Shah, A., Johnson, M., Liu, X., Kaiser, L., Gouws, S., Kato, Y., Kudo, T., Kazawa, H., Stevens, K., Kurian, G., Patil, N., Wang, W., Young, C., Smith, J., Riesa, J., Rudnick, A., Vinyals, O., Corrado, G., Hughes, M., Dean, J., 2016.  Google’s neural machine translation system: Bridging the gap between human and machine translation.  arXiv preprint arXiv:1609.08144 URL: http://arxiv.org/abs/1609.08144. cite arxiv:1609.08144.
  * Zervas et al. (2018) Zervas, G., Yuan, H., Saljoghei, A., Chen, Q., Mishra, V., 2018.  Optically disaggregated data centers with minimal remote memory latency: Technologies, architectures, and resource allocation [invited].  Journal of Optical Communications and Networking 10, A270–A285.  doi:10.1364/JOCN.10.00A270.
  * Zhang et al. (2021) Zhang, J., Yu, H., Xu, W., 2021.  Hierarchical reinforcement learning by discovering intrinsic options, in: International Conference on Learning Representations.  URL: https://openreview.net/forum?id=r-gPPHEjpmw.