Patent Description:
One or more portions of a neural network may be implemented by corresponding individual computing device(s) (e.g. one device may implement one layer), so that the multiple devices collectively implement the neural network. Due to the large number and size of operations generally required to generate the outputs in the neural network, one device can consume significant computer resources and take a significant amount of time to perform its task. The time and computational resources collectively required by the multiple devices depend on the task each device is required to perform and the scheduling of those tasks.

Workloads to be executed by one or more devices may be represented as computation graphs and the one or more devices may execute the computation graph in order to execute the workload. <NPL>) proposes a search space for parallelization strategies to parallelize a DNN in the Sample, Operation, Attribute and Parameter dimensions (SOAP), and a deep learning framework that uses guided randomized search of the SOAP space. <NPL>) proposes a neural architecture search in which, for a given candidate architecture, a graph hypernetwork generates the weights for a neural network of that architecture.

<NPL>) generally relates to superoptimization, which requires the estimation of the best program for a given computational task. In order to deal with large programs, superoptimization techniques perform a stochastic search. This involves proposing a modification of the current program, which is accepted or rejected based on the improvement achieved. This paper proposes learning a proposal distribution over possible modifications using Reinforcement Learning.

<NPL>) proposes a general, black-box Population Based Training (PBT) framework that distributes many asynchronous "trials" (a small number of training steps with warm-starting) across a cluster, coordinated by the PBT controller.

<NPL>) introduces a cost-based optimization framework for fusion plans and describes its end-to-end integration into Apache SystemML. The paper presents techniques for candidate exploration and selection of fusion plans, as well as code generation of local and distributed operations over dense, sparse, and compressed data.

<NPL>) proposes a systematic approach to reduce the memory consumption of deep neural network training. The paper focus on reducing the memory cost to store the intermediate feature maps and gradients during training. Computation graph analysis is used for automatic in-place operation and memory sharing optimizations.

This specification describes a system implemented as computer programs on one or more computers in one or more locations that performs a method as defined by claim <NUM>. In general terms, it determines a schedule for a computation graph by combining a neural network policy with an optimization algorithm, e.g., a genetic algorithm. In other words, the system uses the neural network policy to generate one or more instance-specific proposal distributions that are used by an optimization algorithm that schedules the input computation graph for execution across a plurality of (computing) devices. Each device may be a hardware resource that performs operations independent of other devices in the multiple devices.

The generated schedule can define any of a variety of aspects of the execution of input computation graph for the plurality of devices. As a particular example, the schedule can specify which device each operation represented by a node in the graph should be assigned to and, for each device, the order in which the device should execute the operations assigned to the device.

By combining a neural network policy with an optimization algorithm, e.g., a genetic algorithm, as described in this specification, a schedule can be generated for a computation graph that effectively distributes the execution of the workload represented by the graph across a set of devices. In particular, the performance of the optimization algorithm can be significantly improved by incorporating the neural network as described without a significant increase in the time required or computational resources consumed in generating the schedule. This makes it possible to derive a schedule for the computational task such that, when the schedule is implemented, the computational task is performed collectively by the multiple devices with reduced computing resources and/or more rapidly.

The neural network used by the system described in this specification is "generalizable", that is, it can be used to generate high-quality proposal distributions for computation graphs that were not seen during training. Therefore, the system described in this specification may reduce consumption of computational resources (e.g., memory and computing power) by obviating the need to re-train the network each time a new computation graph needs to be scheduled. In particular, many existing techniques that attempt to use neural networks or other machine learning algorithms to determine a placement for a computation graph across devices require that the model that generates the placements be trained for each new graph that needs to be placed. This additional training consumes a large amount of computational resources, particularly because many of these techniques require that the candidate placements generated during training be evaluated by actually executing the graph using the candidate placement. The described technique, on the other hand, can achieve high quality performance on previously unseen graphs without any additional training.

<FIG> shows an example graph scheduling system <NUM>. The graph scheduling system <NUM> is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented.

The system <NUM> receives data representing an input computation graph <NUM> and generates a schedule <NUM> that includes graph execution decisions for executing the computation graph <NUM> across multiple devices. Optionally, the input computation graph <NUM> may be a portion of a larger computation graph. The system may include a unit operative to receive the larger computation graph and to select the input computation graph from it.

The computation graph <NUM> represents a workload to be distributed across the devices and includes nodes that represent operations and edges that represent dependencies between operations. For example, an edge from one node to another can represent that an output of an operation represented by the first node, e.g., a tensor or other data generated by the operation, is provided as an input to an operation represented by the other node.

As a particular example, the workload can be all of or a portion of a neural network inference workload or a neural network training workload. However, more generally, the computation graph represents any workload that is executed by performing multiple operations that have some kind of dependencies, e.g., data dependencies, between them. The workload may for example be a workload for a computational task which is processing real-world data collected by one or more sensors (e.g. camera(s) and/or microphone(s)), and/or a workload for a computational task which generates control signals to control an electromechanical agent operating on the real world, e.g. moving (e.g. translating and/or changing its configuration) in the real-world.

The devices can include any appropriate types of computer hardware devices, i.e., any devices that are able to perform at least some of the operations represented in the computation graph. In some implementations, the devices are heterogeneous. For example, the devices can include a combination of any of, central processing units (CPUs), graphics processing units (GPUs), application-specific integrated circuits (ASICs) or other special-purpose hardware, field-programmable gate arrays (FPGAs), and so on. In some other implementations, the devices are homogenous, i.e., only include devices of the same device type, i.e., only devices of one of the types above or only devices that are made up of the same combination of devices of the types above.

The schedule <NUM> generated by the system <NUM> can specify any of various aspects of the execution of the computation graph <NUM> on the plurality of devices, i.e., can include any of a variety of graph execution decisions.

In some cases, the schedule <NUM> assigns each operation represented in the graph <NUM> to a respective device. In some of these cases, for each device, the schedule <NUM> also specifies the order in which the device should execute the operations that are assigned to the device.

In some cases, the schedule <NUM> specifies which tensors that are generated while executing the graph should be prioritized for transfer between devices when multiple tensors need to be transferred from one device to another in order to execute the graph according to the schedule.

In some cases, the schedule <NUM> specifies multiple operations that should be fused into a single operation during the execution of the computation graph <NUM>. That is, a single device should perform the specified operations, e.g. as if they were a single operation.

In some cases, the schedule <NUM> specifies which tensors, i.e., which data that is generated by operations represented by nodes, should be stored for later use by other nodes and which tensors should be re-computed. For example, the tensor may not be transferred from a first device which generates it to at least one other device which employs it; instead, the other device (or a third device) may generate the tensor afresh for use by the other device. This may have the advantage that memory space is not consumed by storing the tensor until it is used by the other device.

In particular, the system <NUM> processes data representing the computation graph using a graph neural network <NUM> to generate one or more instance-specific proposal distributions <NUM> ("node-level distribution choices") for an optimization algorithm <NUM> that schedules the input computation graph for execution across the devices. The proposal distributions are referred to as "instance-specific" because different input computation graphs will result in different proposal distributions being generated by the neural network <NUM>.

More specifically, the graph neural network <NUM> has been trained to generate proposal distributions <NUM> that are predicted to result in the optimization algorithm <NUM> generating a schedule <NUM> that optimizes a performance metric that measures the execution of the computation graph <NUM>.

For example, the performance metric can measure the peak memory usage during the execution of the graph, the time required to execute the computation graph, or other properties of the execution that reflect the quality of the schedule. In some cases, the execution graph is subject to some constraint, e.g., on the peak memory use at any given time on any of the devices. In these cases, whenever the constraint is violated, e.g., a generated schedule causes some device to use more than a threshold amount of memory at some point during execution of the graph, the system can set performance metric to a predetermined value that indicates that the constraint was violate (and the generated schedule is not valid).

Once the graph neural network <NUM> has been used to generate the instance-specific proposal distributions <NUM>, the system <NUM> generates a schedule <NUM> for the execution of the computation graph by performing the optimization algorithm <NUM> in accordance with the generated instance-specific proposal distributions <NUM>.

What kind of proposal distributions <NUM> the system <NUM> generates is dependent on the inputs that are required by the optimization algorithm <NUM>. In other words, the system can generate proposal distributions that are appropriate for any of a variety of optimization algorithms <NUM> that can generate an optimized schedule for executing a computation graph on multiple devices.

In some examples, the optimization algorithm is a genetic algorithm. A genetic algorithm begins with an initial population of candidates and, at each of multiple iterations, modifies the population by sampling mutations, crossovers, or both. In this example, each candidate in the initial population is a different possible schedule for the graph. In the example of <FIG>, the optimization algorithm <NUM> is a genetic algorithm that is referred to as Biased Random Key Genetic Algorithm (BRKGA).

Conventionally, these algorithms use fixed distributions, i.e., distributions that are always the same for all input graphs, when determining how to modify the population at a given iteration. Instead of these fixed distributions, the system <NUM> uses the instance-specific distributions generated using the graph neural network. In the BRKGA algorithm, for example, the system generates the parameters for one or more distributions for each node in the computation graph <NUM> ("node-level distributions") and, optionally, a set of elite biases for each node in the computation graph <NUM>.

In some other examples, the optimization algorithm <NUM> is a stochastic local search algorithm. A stochastic local search algorithm samples an initial candidate and, at each of multiple iterations, adjusts the current candidate to generate a final candidate. Conventionally, these algorithms use fixed distributions for sampling the initial candidate, for adjusting the current candidate, or both. Instead, the system <NUM> uses the instance-specific proposal distributions generated using the graph neural network <NUM> to select the initial candidate, to make the local adjustments at each iteration, or both.

While the described techniques can be used to generate instance-specific proposal distributions for any optimization algorithm that optimizes any aspect of a schedule, an example follows that uses BRKGA to optimize a schedule that specifies (i) operation to device assignments (ii) operation scheduling and (iii) tensor transfer priorities.

In particular, BRKGA maintains a population of chromosomes each representing a candidate schedule.

If the computation graph <NUM> is to be scheduled across d devices and includes nodes representing o operations and produces t tensors that may potentially need to be transferred between the devices, each chromosome is an n dimensional vector that has three distinct parts: (<NUM>) o × d entries specifying op-to-device affinities for each of the o operations; (<NUM>) o entries specifying scheduling priorities for each of the o operations and (<NUM>) t × d entries specifying tensor-to-device priorities for transfers that may be needed.

Once a final candidate has been generated by BRKGA, i.e., a chromosome has been selected, the system <NUM> then obtains a schedule <NUM> from the final chromosome by performing a topological sort over the operations given their tensor dependencies, breaking ties by using the corresponding scheduling priorities for the operations.

To generate a final candidate, BRKGA performs multiple evolution steps, i.e., a fixed number of steps or runs for a fixed amount of time or runs for a fixed number of evaluation calls, and is specified by the following: <NUM>) scalar integer parameters π, πe, and πc, representing the population size, number of elites, and number of children, respectively, <NUM>) respective elite biases for each of the n entries ρi ∈ [<NUM>, <NUM>), and <NUM>) a mutant generation distribution D over [<NUM>, <NUM>]n. The procedure aims to find a chromosome that maximizes f, a function that maps a chromosome to a performance metric.

The initial population is created by sampling from D, using known good solutions, or a mixture of both. One evolution step is completed as follows.

Given elite and nonelite chromosome a and b both in [<NUM>, <NUM>]n, the crossover procedure produces a child chromosome c by independently combining entries from the parents. Specifically, for each index i ∈ <NUM>,. , n independently, let ci = ai with probability ρi and ci = bi with probability <NUM> - ρi.

Any new chromosome in the population is then evaluated to determine the fitness f of the chromosome, i.e., to determine the performance metric of the schedule defined by the chromosome. Evaluating a schedule is described in more detail below with reference to <FIG>.

Thus, BRKGA requires the probability distribution D and the elite biases in order to operate. Conventionally, each of these would be agnostic to the input graph, i.e., would be pre-configured and held constant for all input graphs. Instead, the graph neural network <NUM> is used to predict node-specific probability distributions, node-specific elite biases, or both, that are used in place of the probability distribution D and the elite biases.

For example, the system can, instead of using a single distribution D, use n independent beta distributions in place of the single distribution D, one for each entry in a given chromosome. The graph neural network <NUM> may be used to predict the parameters of the beta distributions for the entries of the chromosome that correspond to different operations represented by nodes in the graph, i.e., for a given node in the graph, the graph neural network <NUM> can be used to predict the parameters of d + <NUM> independent beta distributions, corresponding to device affinities for the operation represented by the node and the scheduling priority for the operation represented by the node. The beta distributions corresponding to the remaining d x t entries specifying tensor-to-device priorities can be set to graph-agnostic pre-determined distributions.

BRKGA can then use these beta distributions in place of D when constructing the initial population and generating new chromosomes for each new generation of the population, i.e., when generating the π - πe - πc mutated chromosomes for each new generation.

Optionally, instead of or in addition to predicting the parameters of the beta distributions for a given node, the graph neural network <NUM> can be used to predict the elite bias for the node or parameters of a probability distribution over elite bias values for the node. Thus, in these cases, the output of the graph neural network <NUM> is also used to perform the crossover procedure when generating a new generation of the population.

The operation of BRKGA is described in more detail in <NPL>.

To generate the proposal distributions using the graph neural network <NUM>, the system <NUM> processes the data representing the computational graph <NUM>, i.e., attribute vectors for the nodes and edges in the graph, using the graph neural network <NUM> to generate a respective representation vector for each node in the graph. The attribute vectors for the nodes and edges represent features of the node or edge, e.g., sizes of the tensors received as input or output of an operation or transmitted over an edge, or types of operations performed by nodes.

Any of a variety of graph neural network architectures that are configured to process graph data to generate representation vectors for nodes in the graph can be used. One example graph neural network architecture is described in <NPL>. Another example graph neural network architecture is described in <NPL>). Yet another example graph neural network architecture is described in <NPL>.

The system <NUM> then generates the proposal distributions <NUM> from the representation vectors for the nodes of the graph that are generated by the graph neural network <NUM>.

In some implementations, the system <NUM> generates the parameters of the proposal distribution(s) for a given node in the graph only from the representation vector for the given node, e.g., by processing the representation vector for the given node through a multi-layer perceptron or other neural network.

In some other implementations, the system <NUM> generates the parameters of the proposal distribution(s) for the nodes in the graph auto-regressively. That is, the system <NUM> orders the nodes and then generates the parameters of the proposal distribution(s) for a given node conditioned on the representation vector for the given node and the generated parameters for any nodes that are before the given node in the order, e.g., by processing the feature representations of the nodes in the order using a recurrent neural network.

In some implementations, the system directly generates the parameters of the distributions, e.g., by directly regressing the values of the parameters.

In other implementations, the system <NUM> can define a discrete action space in which each discrete action in the space maps to a unique set of parameters for the proposal distribution(s) needed by the optimization algorithm for a given node. That is, each action in the discrete action space is a different set of parameters for the proposal distributions needed by the optimization algorithm for a given node. For each node, the system <NUM> then generates a probability distribution over the discrete action space from the vector representation of the node (using one of the techniques above) and then samples from the probability distribution or selects the action with the highest probability to generate the proposal distributions for the node.

In some implementations, the system <NUM> then executes the computation graph on the devices in accordance with the generated schedule.

In some other implementations, the system <NUM> provides data specifying the generated schedule to another system and the other system uses the provided data to execute the computation graph on the devices.

<FIG> is a flow diagram of an example process <NUM> for scheduling a computation graph. For convenience, the process <NUM> will be described as being performed by a system of one or more computers located in one or more locations. For example, a graph scheduling system, e.g., the graph scheduling system <NUM> of <FIG>, appropriately programmed in accordance with this specification, can perform the process <NUM>.

The system obtains data representing an input computation graph (step <NUM>).

The input computation graph includes a plurality of nodes that are connected by edges, with the nodes representing operations and the edges representing dependencies between the operations. For example, the input computation graph can represent all or some of the operations required to perform an inference using a neural network or to train a neural network.

Because of the way that the graph neural network has been trained (as will be described below with reference to <FIG>), the input computation graph need not be a graph that was included in the training data used to train the graph neural network.

The system processes the data representing the input computation graph using the graph neural network (step <NUM>). As described above, the graph neural network is a neural network having a plurality of network parameters and configured to process the data representing the input computation graph in accordance with the network parameters to generate one or more instance-specific proposal distributions for the optimization algorithm. In other words, the graph neural network generates instance-specific proposal distributions that are used by the optimization algorithm instead of the default or conventional proposal distributions that would conventionally be used by the optimization algorithm.

The system generates a schedule for the input computation graph by performing the optimization algorithm in accordance with the one or more instance-specific proposal distributions generated by the graph neural network for the input computation graph (step <NUM>). In other words, the system runs the optimization algorithm using the instance-specific proposal distributions in place of the conventional proposal distributions that would be used by the optimization algorithm. For example, the system can run the optimization for a fixed number of iterations or for a fixed amount of time and then use the solution found by the algorithm after the fixed number of iterations or after the fixed amount of time as the generated schedule.

Thus, the only computational overhead introduced to the scheduling process by generating the instance-specific probability distributions is the overhead that is required to perform a forward pass through the graph neural network for the computation graph, which will typically be minimal relative to the amount of computational resources consumed by the optimization algorithm.

In some implementations, the system then executes the input computation graph on the plurality of devices by causing the plurality of devices to perform the operations represented by the nodes in the input computation graph in accordance with the generated schedule.

In some other implementations, the system provides data specifying the executed schedule to another system, which then uses the data to cause the devices to perform the operations represented by the nodes in the input computation graph in accordance with the generated schedule.

<FIG> is a flow diagram of an example process <NUM> for training the graph neural network. For convenience, the process <NUM> will be described as being performed by a system of one or more computers located in one or more locations. For example, a graph scheduling system, e.g., the graph scheduling system <NUM> of <FIG>, appropriately programmed in accordance with this specification, can perform the process <NUM>.

The system can repeatedly perform the process <NUM> for different training examples on a set of training data in order to repeatedly adjust the values of the network parameters to determine trained values of the network parameters. Each training example in the training data is data representing a different training computation graph. Thus, by performing the process <NUM> the system trains the neural network on different computation graphs and such that the trained neural network will generalize the computation graphs that are not represented in the set of training data.

The system processes a training example, i.e., data that represents a computation graph, in accordance with current values of the network parameters to generate one or more training instance-specific proposal distributions for the optimization algorithm (step <NUM>), i.e., to generate all of the proposal distributions that are required for the optimization algorithm to run.

The system generates a training schedule for the training computation graph represented by the training example by performing the optimization algorithm in accordance with the one or more training instance-specific proposal distributions (step <NUM>).

The system determines a performance metric for the execution of the training computation graph (step <NUM>).

Generally, the performance metric measures one or more properties of the execution of the training computation graph that are attempting to be optimized by the generated schedule. For example, the performance metric can measure the peak memory usage during the execution of the graph, the time required to execute the computation graph, or other properties of the execution that reflect the quality of the schedule. The performance metric can also be derived from a combination of multiple properties of the execution, e.g., as a weighted sum of the peak memory usage and the time required to execute the graph.

In some implementations, the system determines the performance metric by executing the training computation graph on the plurality of devices by causing the plurality of devices to perform the operations represented by the nodes in the input computation graph in accordance with the training schedule and measuring the properties of the execution that are used to generate the performance metric.

In some other implementations, the system maintains a cost model that models the values of the one or more properties for a given input schedule and uses the maintained cost model to determine the values of the one or more properties, i.e., without needing to execute the graph on the devices. Maintaining and using such a computationally cheap cost model enables fast optimization and may be better suited for distributed training of the graph neural network since a cost model is cheap to replicate in parallel actors, while hardware environments are not. Example techniques for constructing such a cost model for a given set of devices are described in <NPL> and <NPL>.

The system generates a reward from the performance metric (step <NUM>). Generally, the system maps the performance metric to a reward value that can be maximized to improve the performance metrics. In some cases, to improve the generalization of the neural network to different graphs, the system can generate the reward based on the performance metric and a baseline performance metric generated from the execution of the training computation graph in accordance with a baseline schedule. The baseline schedule can be, e.g., a schedule that is generated by the optimization algorithm in accordance with default, graph-agnostic proposal distributions. As a particular example, the reward can be the negative of the ratio between the performance metric and the baseline performance metric.

The system determines an update to the current values of the network parameters based on the reward (step <NUM>). This may be done using a reinforcement learning technique. For example, the system can maximize the expected reward through a policy gradient reinforcement technique, e.g., REINFORCE with or without a variance reducing baseline.

Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification, or in combinations of one or more of them.

Claim 1:
A computer-implemented method for scheduling and executing an input computation graph representing a computational task to be performed collectively by a plurality of devices, the method comprising:
obtaining (<NUM>) data representing the computation graph, the input computation graph comprising a plurality of nodes that are connected by edges, the nodes representing operations and the edges representing dependencies between the operations;
processing (<NUM>) the data representing the input computation graph using a graph neural network having a plurality of network parameters, wherein the graph neural network is configured to process the data representing the input computation graph in accordance with first values of the network parameters to generate one or more instance-specific proposal distributions for an optimization algorithm that schedules the input computation graph for execution across the plurality of devices;
generating (<NUM>) a schedule for the input computation graph by performing the optimization algorithm in accordance with the one or more instance-specific proposal distributions generated by the graph neural network for the input computation graph, wherein the optimization algorithm generates an assignment that assigns each operation to a respective device from the plurality of devices; and
executing the input computation graph on the plurality of devices by causing the plurality of devices to perform the operations represented by the nodes in the input computation graph in accordance with the generated schedule.