Device placement optimization with reinforcement learning

1. A computer-implemented method for training a placement neural network, the placement neural network configured to process a sequence of operation embeddings with each operation embedding in the sequence characterizing one or more respective operations that are part of performing the processing of a machine learning model to generate a network output defining a placement of the operations characterized by the operation embeddings in the sequence across a plurality of devices, the method comprising:

processing a training sequence of operation embeddings using the placement neural network in accordance with current values of network parameters of the placement neural network to select K placements of the operations across the plurality of devices;

for each of the selected K placements:

performing the processing of the machine learning model with the operations across the plurality of devices according to the placement, and

determining a running time required for the processing to complete; and

adjusting the current values of the network parameters using a reinforcement learning technique that uses a reward derived from the running times required for the processing to complete for each of the selected K placements.

2. The method of claim 1 , wherein processing the training sequence of operation embeddings to select the K placements comprises:

running the placement neural network K times to draw the K placements from a probability distribution of placements defined by the placement neural network.

3. The method of claim 2 , wherein running the placement neural network K times to draw the K placements comprises:

providing a batch of K identical input examples to the placement neural network, wherein each input example in the batch is the same training sequence of operation embedding, and

for each input example in the batch, processing the training sequence of operation embeddings through an encoder recurrent neural network and a decoder neural network of the placement neural network to generate a respective placement in accordance with a probability distribution of placements defined by the placement neural network.

4. The method of claim 1 , wherein the reinforcement learning technique is a REINFORCE technique.

5. The method of claim 1 , wherein adjusting the current values of the network 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 K placements comprises:

adjusting the current values of the network parameters to optimize an objective function.

6. The method of claim 5 , wherein the objective function defines a negative of the reward and computes an expected value of square roots of K running times for the selected K placements.

7. The method of claim 5 , wherein the reinforcement learning technique comprises: using a baseline that is a moving average of the K running times to approximate a gradient of the objective function, and backpropagating the gradient of the objective function to determine adjusted values for the network parameters that minimize the objective function.

8. The method of claim 1 , wherein adjusting the current values of the network parameters further comprises: adjusting the operation embeddings in the training sequence as part of the reinforcement learning technique.

9. One or more non-transitory computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform operations for training a placement neural network, the placement neural network configured to process a sequence of operation embeddings with each operation embedding in the sequence characterizing one or more respective operations that are part of performing the processing of a machine learning model to generate a network output defining a placement of the operations characterized by the operation embeddings in the sequence across a plurality of devices, the operations comprising:

processing a training sequence of operation embeddings using the placement neural network in accordance with current values of network parameters of the placement neural network to select K placements of the operations across the plurality of devices;

for each of the selected K placements:

performing the processing of the machine learning model with the operations across the plurality of devices according to the placement, and

determining a running time required for the processing to complete; and

adjusting the current values of the network parameters using a reinforcement learning technique that uses a reward derived from the running times required for the processing to complete for each of the selected K placements.

10. The one or more non-transitory computer storage media of claim 9 , wherein the operations for processing the training sequence of operation embeddings to select the K placements comprise:

running the placement neural network K times to draw the K placements from a probability distribution of placements defined by the placement neural network.

11. The one or more non-transitory computer storage media of claim 10 , wherein the operations for running the placement neural network K times to draw the one or more K placements comprises:

providing a batch of K identical input examples to the placement recurrent neural network, wherein each input example in the batch is the same training sequence of operation embedding, and

for each input example in the batch, processing the training sequence of operation embeddings through an encoder recurrent neural network and a decoder neural network of the placement neural network to generate a respective placement in accordance with a probability distribution of placements defined by the placement neural network.

12. The one or more non-transitory computer storage media of claim 9 , wherein the reinforcement learning technique is a REINFORCE technique.

13. The one or more non-transitory computer storage media of claim 9 , wherein adjusting the current values of the network 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 K placements comprises:

adjusting the current values of the network parameters to optimize an objective function.

14. The one or more non-transitory computer storage media of claim 13 , wherein the objective function computes an expected value of square roots of K running times for the selected K placements.

15. The one or more non-transitory computer storage media of claim 13 , wherein the reinforcement learning technique comprises: using a baseline that is a moving average of the running times to approximate a gradient of the objective function, and backpropagating the gradient of the objective function to determine adjusted values for the network parameters that minimize the objective function.

16. The one or more non-transitory computer storage media of claim 9 , wherein adjusting the current values of the network parameters further comprises:

adjusting the operation embeddings in the training sequence as part of the reinforcement learning technique.

17. A system comprising:

one or more computers; and

one or more non-transitory computer storage media storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations for training a placement neural network, the placement neural network configured to process a sequence of operation embeddings with each operation embedding in the sequence characterizing one or more respective operations that are part of performing the processing of a machine learning model to generate a network output defining a placement of the operations characterized by the operation embeddings in the sequence across a plurality of devices, the operations comprising:

processing a training sequence of operation embeddings using the placement neural network in accordance with current values of network parameters of the placement neural network to select K placements of the operations across the plurality of devices;

for each of the selected K placements:

performing the processing of the machine learning model with the operations across the plurality of devices according to the placement, and

determining a running time required for the processing to complete; and

adjusting the current values of the network parameters using a reinforcement learning technique that uses a reward derived from the running times required for the processing to complete for each of the selected K placements.

18. The system of claim 17 , wherein the operations for processing the training sequence of operation embeddings to select the K placements comprise:

running the placement neural network K times to draw the K placements from a probability distribution of placements defined by the placement neural network.

19. The system of claim 18 , wherein the operations for running the placement neural network K times to draw the one or more K placements comprises:

providing a batch of K identical input examples to the placement recurrent neural network, wherein each input example in the batch is the same training sequence of operation embedding, and

for each input example in the batch, processing the training sequence of operation embeddings through an encoder recurrent neural network and a decoder neural network of the placement neural network to generate a respective placement in accordance with a probability distribution of placements defined by the placement neural network.

20. The system of claim 17 , wherein adjusting the current values of the network 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 K placements comprises:

adjusting the current values of the network parameters to optimize an objective function.