Patent Description:
Some neural networks are recurrent neural networks. A recurrent neural network is a neural network that receives an input sequence and generates an output sequence from the input sequence. In particular, a recurrent neural network can use some or all of the internal state of the network from a previous time step in computing an output at a current time step. An example of a recurrent neural network is a long short term (LSTM) neural network that includes one or more LSTM memory blocks. Each LSTM memory block can include one or more cells that each include an input gate, a forget gate, and an output gate that allow the cell to store previous states for the cell, e.g., for use in generating a current activation or to be provided to other components of the LSTM neural network.

"<NPL>) discloses TensorFlow, an interface for expressing machine learning algorithms, and an implementation for executing such algorithms. A computation expressed using TensorFlow can be executed with little or no change on a wide variety of heterogeneous systems, ranging from mobile devices such as phones and tablets up to large-scale distributed systems of hundreds of machines and thousands of computational devices such as GPU cards.

"<NPL>) discloses a static method that uses Hopfield neural network to cluster the tasks of a parallel program for a given system. This method takes into account both load balancing and communication minimization.

"<NPL>) discloses a systematic feature selection method that reduces the feature set size based on the extent to which features affect performance. These features are used as input to a reinforcement learning technique, called NeuroEvolution of Augmenting Topologies (NEAT), that uses a genetic algorithm to evolve neural networks to automatically discover placement heuristics.

The invention is defined in the appended independent claims. Embodiments of the invention are defined in the appended dependent claims.

<FIG> illustrates a device placement system <NUM> that determines a placement for operations of a machine learning model across multiple hardware devices. The device placement system <NUM> can be implemented as computer programs on one or more computers in one or more locations.

The machine learning model being placed can be configured to receive any kind of digital data input and to generate any kind of score, classification, or regression output based on the input.

For example, if the inputs to the machine learning model are images or features that have been extracted from images, the output generated by the machine learning model for a given image may be scores for each of a set of object categories, with each score representing an estimated likelihood that the image contains an image of an object belonging to the category.

As another example, if the inputs to the machine learning model are Internet resources (e.g., web pages), documents, or portions of documents or features extracted from Internet resources, documents, or portions of documents, the output generated by the machine learning model for a given Internet resource, document, or portion of a document may be a score for each of a set of topics, with each score representing an estimated likelihood that the Internet resource, document, or document portion is about the topic.

As another example, if the inputs to the machine learning model are features of a personalized recommendation for a user, e.g., features characterizing the context for the recommendation, e.g., features characterizing previous actions taken by the user, the output generated by the machine learning model may be a score for each of a set of content items, with each score representing an estimated likelihood that the user will respond favorably to being recommended the content item.

As another example, if the input to the machine learning model is a sequence of text in one language, the output generated by the machine learning model may be a score for each of a set of pieces of text in another language, with each score representing an estimated likelihood that the piece of text in the other language is a proper translation of the input text into the other language.

As another example, if the input to the machine learning model is a sequence representing a spoken utterance, the output generated by the machine learning model may be a score for each of a set of pieces of text, each score representing an estimated likelihood that the piece of text is the correct transcript for the utterance.

In particular, the device placement system <NUM> receives input data <NUM> that specifies a machine learning model to be placed for distributed processing on a plurality of hardware devices. The hardware devices are generally heterogeneous, and can include any appropriate hardware device, e.g., a combination of any of, CPUs, GPUs, ASICs or other special-purpose hardware, FPGAs, and so on.

The input data <NUM> specifying the machine learning model may include data that represents a computational graph. The computational graph has vertices that represent operations and edges that represent data communicated between the operations.

For example, the input data <NUM> includes data that represents a computational graph G having vertices that represent M operations {o<NUM>, o<NUM>,. The M operations can be operations to train the machine learning model or operations to generate outputs from received inputs using the machine learning model once the machine learning model has already been trained. Given M operations, the device placement system <NUM> aims to determine a placement P = {p<NUM>, p<NUM>,. The placement P is an assignment of each operation oi ∈ G to a device pi that belongs to a set of D available hardware devices, i.e., pi ∈ {<NUM>,.

An example computational graph and an example placement of computational graph operations of the graph on multiple hardware devices are described in detail with reference to <FIG>.

To determine a placement, the system <NUM> trains a placement recurrent neural network <NUM> that generates outputs that define placements of the operations across the devices. Once the placement recurrent neural network <NUM> has been trained, the system <NUM> can generate a final placement. The system <NUM> runs the trained placement recurrent neural network <NUM> and uses the output of the trained placement recurrent neural network <NUM> to determine the final placement.

The system <NUM> can then schedule the machine learning model for processing by the plurality of hardware devices, i.e., causing the operations of the machine learning model to be executed according to the final placement. In some other cases, the system <NUM> can provide data identifying the final placement to another system that manages the execution of the machine learning model so that the other system can place the operations across the devices according to the final placement.

As part of training the placement recurrent neural network, the system <NUM> generates, from the input data <NUM>, a sequence of operation embeddings <NUM>. Each operation embedding in the sequence <NUM> characterizes one or more respective operations necessary to perform the processing of the machine learning model. An embedding is an ordered collection of numeric values, e.g., a vector or a matrix of floating point values or of quantized floating point values.

In some cases, the system combines multiple different embeddings to generate a single operation embedding for each operation in the sequence.

More specifically, to generate an operation embedding characterizing a particular operation, the system <NUM> generates a type embedding of an operation type of the particular operation. For example, an operation type may describe an underlying computation (e.g., matrix multiplication or two-dimensional convolution or one-dimensional convolution or non-linear activation function) of the operation, and the type embedding may be a tunable embedding vector of the operation type, i.e., so that each operation of the same type shares the same type embedding.

The system <NUM> generates an output size embedding that characterizes a size of outputs generated by the particular operation. For instance, the system <NUM> may record the size of each of the outputs (e.g., output tensors) generated by the particular operation and concatenate the recorded sizes to generate into an output shape (e.g., a fixed-size zero-padded list). The output shape is the output size embedding of the particular operation.

The system <NUM> generates an adjacency embedding (e.g., a one-hot encoding vector) that identifies operations that provide input to and receive output generated by the particular operation.

The system <NUM> generates the operations embedding characterizing the particular operations from a combination of embeddings. The combination of embeddings used by the system <NUM> can vary. For example, in some cases, the system <NUM> may combine three kinds of embeddings described above, e.g., the type embedding, the output size embedding, and the adjacency embedding, to generate the operation embedding characterizing the particular operation. For example, the system <NUM> may concatenate the type embedding, the output size embedding, and the adjacency embedding to generate the operation embedding. In some other cases, the system <NUM> may combine two of the three kinds of embeddings described above to generate the operation embedding. In some other cases, the system <NUM> may combine one or more of the three kinds of embeddings described above with a new kind of embedding to generate the operation embedding.

In some implementations, as part of generating the sequence of operation embeddings, the system <NUM> determines that two or more of the operations represented by vertices in the computational graph are to be co-located on the same device, and in response, the system <NUM> generates a single operation embedding that characterizes the two or more operations.

During each iteration of the training of the placement recurrent neural network <NUM>, the system <NUM> processes the sequence of operation embeddings <NUM> using the placement recurrent neural network <NUM> in accordance with current values of network parameters of the placement recurrent neural network <NUM>. The placement recurrent neural network <NUM> is configured to process the sequence of operation embeddings <NUM> in accordance with the current values to generate a network output <NUM>. The network output <NUM> defines a placement of the operations characterized by the operation embeddings in the sequence across the plurality of devices.

In particular, the placement recurrent neural network <NUM> is configured to generate, for each of the operation embeddings in the sequence <NUM>, a set of scores that includes a respective score for each hardware device in the set of available hardware devices. A respective score for each hardware device is a likelihood that represents how likely it is that the hardware device is the best device to assign the operation characterized by the operation embedding. The neural network <NUM> is then configured to select a device for each of the operations using the set of scores for the operation embedding. In some cases, the neural network <NUM> may select the device that has the highest score according to the set of scores for the operation embedding. In some other cases, the neural network <NUM> may sample a device from the plurality of devices according to probabilities defined by the set of scores for the operation embedding characterizing the operation.

Once a device is selected for each of the operations, the neural network <NUM> outputs the network output <NUM> that defines a placement of the operations across the plurality of hardware devices.

The system <NUM> may schedule the machine learning model for processing by the plurality of hardware devices by placing the operations on the plurality of devices according to the placement defined by the network output <NUM>.

Generally, the placement recurrent neural network <NUM> can be a recurrent neural network that includes a sequence-to-sequence model with Long Short-Term Memory (LSTM) neural network layers and a content-based attention mechanism. An example sequence-to-sequence model is described in <NPL>. An example content-based attention mechanism is described in <NPL>.

The architecture of the placement recurrent neural network <NUM> is divided into two parts: an encoder recurrent neural network <NUM> and a decoder neural network <NUM>.

The encoder recurrent neural network <NUM> is configured to receive as input a sequence of operation embeddings. The encoder recurrent neural network <NUM> processes the sequence of operation embeddings to generate a respective encoder hidden state for each of the operation embeddings.

For each of the operation embeddings, the decoder neural network <NUM> is configured to receive a decoder input and to process the decoder input and the encoder hidden states to generate a set of scores for the operation embedding. The decoder input for each of the operation embeddings after a first operation embedding in the sequence identifies a device selected for the one or more operations represented by the preceding operation embedding in the sequence. For the first operation embedding, the decoder neural network <NUM> may process only the encoder hidden states to generate the set of scores for the first operation embedding.

An example architecture of the placement recurrent neural network <NUM> is described in more detail below with reference to <FIG>.

The system <NUM> can update values of the network parameters of the placement recurrent neural network <NUM> based on running time using the process described in detail below with reference to <FIG>.

<FIG> shows an example architecture of a placement recurrent neural network. The placement recurrent neural network includes an encoder <NUM> and a decoder <NUM>.

The encoder recurrent neural network <NUM> is configured to receive as input a sequence of operation embeddings (e.g., embeddings <NUM>, <NUM>. The sequence of operation embeddings characterize operations that are necessary to perform the processing of a machine learning model on a plurality of hardware devices. The encoder recurrent neural network <NUM> processes the sequence of operation embeddings to generate a respective encoder hidden state for each of the operation embeddings. For example, as shown in <FIG>, the encoder recurrent neural network <NUM> generates a list of encoder hidden states e<NUM>, e<NUM>,. eM for operation embeddings x<NUM>, x<NUM>,. , xM, where M is the number of operations that are necessary to perform the processing of the machine learning model.

The decoder neural network <NUM> maintains LSTM hidden states d<NUM>, d<NUM>,. , dM and is configured to output a device for a respective operation embedding at each decoding time step. Each decoding time step corresponds to one operation embedding.

In particular, for each of the operation embeddings, the decoder neural network <NUM> is configured to receive a decoder input and to process the decoder input and a set of appropriate encoder hidden states to generate a set of scores for the operation embedding. The decoder neural network <NUM> may generate the set of scores for the operation embedding using a softmax neural network layer. The set of scores for the operation embedding includes a respective score for each hardware device in the plurality of hardware devices. A respective score for each hardware device is a likelihood that represents how likely it is that the hardware device is the best device to assign the operation characterized by the operation embedding. The decoder neural network <NUM> is then configured to select a device for each of the operations using the set of scores for the operation embedding. In some cases, the decoder <NUM> may select the device that has the highest score according to the set of scores for the operation embedding. In some other cases, the decoder <NUM> may sample a device from the plurality of devices according to probabilities defined by the set of scores for the operation embedding characterizing the operation.

The decoder neural network <NUM> can use an attention mechanism to determine a set of appropriate encoder hidden states to be used at each decoding time step. The decoder input for each of the operation embeddings after the first operation embedding in the sequence identifies a device selected for the one or more operations represented by the preceding operation embedding in the sequence. For example, the decoder input <NUM> for the second operating embedding in the sequence is a device embedding that identifies the device <NUM> that is selected for the first operating embedding. For the first operation embedding, the decoder neural network <NUM> may process only the appropriate encoder hidden states to generate the set of scores for the first operation embedding.

<FIG> is a flow diagram of an example process <NUM> for training a placement recurrent neural network (e.g., the placement recurrent neural network <NUM> of <FIG>) to update values of network parameters of the placement recurrent 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 device placement system, e.g., the device placement system <NUM> of <FIG>, appropriately programmed in accordance with this specification, can perform the process <NUM>.

Generally, given a computational graph G having vertices that represent M operations {o<NUM>, o<NUM>,. oM} that are necessary to perform the processing (or training) of a machine learning model on a plurality of hardware devices, it is desirable for the trained placement recurrent neural network to determine a placement that requires a minimal time to perform the processing of the machine learning model under the placement. A placement P = {p<NUM>, p<NUM>,. , pM} is an assignment of an operation oi ∈ G to a device pi that belongs to a set of D hardware devices, i.e., pi ∈ {<NUM>,. Let r(P) denote the time that it takes to perform a complete execution of M operations in the computational graph G under the placement P (hereafter referred to as running time). The system trains the placement recurrent neural network to find P such that the execution time r(P) is minimized.

To update values of the network parameters of the placement recurrent neural network (e.g., from initial values or current values of the network parameters), the system repeatedly performs steps <NUM>-<NUM> as follows.

The system processes a current sequence of operation embeddings using the placement recurrent neural network in accordance with current values of network parameters of the placement recurrent neural network to select one or more placements (e.g., K placements) of the operations across the plurality of devices (step <NUM>).

For example, to select K placements, the system can run the placement recurrent neural network K times to draw K placements from a probability distribution of placements defined by the placement recurrent neural network. That is, the system provides a batch of K identical input examples to the placement recurrent neural network. Each input example in the batch is the same current sequence of operation embeddings. For each input examples in the batch, the placement recurrent neural network is configured to process the current sequence of operation embeddings through an encoder recurrent neural network and a decoder neural network to generate a placement in accordance with a probability distribution of placement defined by the placement recurrent neural network (i.e., defined by a softmax neural network layer of the placement recurrent neural network) in the manner as described in detail above with reference to <FIG>.

The system performs step <NUM> for each selected placement. In particular, the system performs the processing of the machine learning model with the operations across the plurality of devices according to the placement, and then determines a time required for the processing to complete. That is, for each selected placement, the system can monitor the processing of the machine learning model with the operations placed according to the selected placement and identify the time required for the processing to complete.

The system adjusts the current values of the parameters using a reinforcement learning technique that uses a reward derived from the times required for the processing to complete for each of the selected placements (step <NUM>). Generally, the reward is higher when the running time is shorter to encourage the placement neural network to generate placements that have shorter running times.

In particular, the system trains the placement recurrent neural network to maximize the reward by minimizing the following objective function that defines a negative of the reward:
<MAT>
where θ denotes the network parameters of the placement recurrent neural network, π(<IMG>; θ) denotes a probability distribution of placements defined by the placement recurrent neural network, <IMG> is the square root of the running time r(P). The objective function computes an expected value of the square roots of K running times for the K placements selected by the system (e.g., by running the placement recurrent neural network K times to draw K placements from the probability distribution of placements π(<IMG>; θ)) given the computational graph <IMG>.

The system determines a gradient of the objective function using a reinforcement learning technique, e.g., a REINFORCE technique. Using a REINFOCE technique, the gradient of the objective function can be expressed as follows:
<MAT>
where p is the probability assigned to the placement by the placement recurrent neural network, i.e., the product of the scores generated by the placement recurrent neural network for the individual placements.

To reduce the variance of the gradient, the system can use a baseline term B, which is a moving average of the K running times, to approximate the gradient as follows:
<MAT>.

The system can backpropagate the gradient of the objective function to determine a respective updated value for each of the network parameters of the placement recurrent neural network.

While adjusting the current values of the network parameters, the system can optionally adjust the operation embeddings in the current sequence as part of the reinforcement learning technique, i.e., by backpropagating gradients into the operation embeddings.

<FIG> illustrates an example computational graph <NUM> and an example placement of computational graph operations of the graph <NUM> on multiple devices of a system <NUM>. The computational graph <NUM> includes vertices that represent operations <NUM>-<NUM>. The computational graph <NUM> further includes edges <NUM>-<NUM> that represent data communicated between the operations <NUM>-<NUM>. For example, the computational graph <NUM> can be a TensorFlow computational graph whose vertices represent operations and edges represent multidimensional data arrays (tensors) communicated between the operations.

Each of the operations <NUM>-<NUM> can be assigned to one of the devices <NUM>-<NUM> of the system <NUM> according to a particular placement, e.g., a placement determined by a device placement system (e.g., the device placement system <NUM> of <FIG>) using a placement recurrent neural network (e.g., the placement recurrent neural network <NUM> of <FIG>). <FIG> shows an example placement in which operation <NUM> and <NUM> are assigned to device <NUM>, operations <NUM>, <NUM>, and <NUM> are assigned to device <NUM>, operation <NUM> is assigned to device <NUM>, and operations <NUM> and <NUM> are assigned to device <NUM>.

The devices <NUM>-<NUM> are generally heterogeneous, and can include any appropriate hardware devices, e.g., a combination of any of, CPUs, GPUs, ASICs or other special-purpose hardware, FPGAs, and so on. For example, in some implementations, the system <NUM> can be a single machine with one or more CPUs and one or more GPUs and each of the operations <NUM>-<NUM> can be assigned to one of the CPUs and GPUs according to the particular placement. In some other implementations, the system <NUM> can include multiple machines with a mix of GPUs and ASICs or FPGAs and each of the operations <NUM>-<NUM> can be assigned to a device of one of the machines according to the particular placement.

Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non transitory program carrier for execution by, or to control the operation of, data processing apparatus. The computer storage medium is not, however, a propagated signal.

Claim 1:
A computer-implemented method comprising:
receiving data (<NUM>) specifying a machine learning model to be placed for distributed processing on a plurality of hardware devices;
generating, from the data specifying the machine learning model, a sequence of operation embeddings (<NUM>), wherein each operation embedding in the sequence characterizes one or more respective operations necessary to perform the processing of the machine learning model;
processing the sequence of operation embeddings using a placement recurrent neural (<NUM>) network in accordance with first values of a plurality of network parameters of the placement recurrent neural network, wherein the first values have been determined through training of the placement recurrent neural network using a reinforcement learning technique that uses a reward derived from execution times for different placements of operations across the plurality of hardware devices,
wherein the placement recurrent neural network is configured to process the sequence of operation embeddings in accordance with the first values to generate a network output (<NUM>) defining a placement of the operations characterized by the operation embeddings in the sequence across the plurality of hardware devices,
wherein the placement recurrent neural network is configured to generate, for each of the operation embeddings in the sequence, a set of scores that includes a respective score for each of the plurality of hardware devices, and
wherein processing the sequence of operation embeddings comprises selecting a hardware device for each of the operations using the set of scores for the operation embedding characterizing the operation; and
scheduling the machine learning model for processing by the plurality of hardware devices by placing the operations on the plurality of hardware devices according to the placement defined by the network output,
wherein the placement recurrent neural network comprises:
an encoder recurrent neural network (<NUM>) configured to process the sequence of operation embeddings to generate a respective encoder hidden state for each of the operation embeddings; and
a decoder neural network (<NUM>) configured to, for each of the operation embeddings:
receive a decoder input; and
process the decoder input and the encoder hidden states to generate the set of scores for the operation embedding.