Hierarchical device placement with reinforcement learning

Patent metadata

Field	Value
Publication number	US-10438113-B2
Application number	US-201816040186-A
Country	US
Kind code	B2
Filing date	Jul 19, 2018
Priority date	Jul 21, 2017
Publication date	Oct 8, 2019
Grant date	Oct 8, 2019

Abstract

A method for determining a placement for machine learning model operations across multiple hardware devices includes receiving data specifying machine learning operations, and determining a placement that assigns each of the operations specified by the data to a respective device from the multiple hardware devices. Determining the placement includes: generating, from the data, a respective operation embedding for each of the operations; grouping the operations into multiple operation groups, comprising processing each of the respective operation embeddings using a grouper neural network having multiple grouper parameters, in which the grouper neural network is configured to, for each of the operations, process the operation embedding for the operation in accordance with first values of the grouper parameters to generate a grouper output that assigns the operation to an operation group from the multiple operation groups; and assigning each of the operation groups to a respective device from the multiple hardware devices.

First claim

What is claimed is:

1. A method comprising:
receiving data specifying machine learning operations that are to be distributed across a plurality of hardware devices for processing; and
determining a placement that assigns each of the operations specified by the received data to a respective device from the plurality of hardware devices, comprising:
generating, from the data specifying the operations, a respective operation embedding for each of the plurality of operations;
grouping the operations into a plurality of operation groups, comprising processing each of the respective operation embeddings using a grouper neural network having a plurality of grouper parameters,
wherein the grouper neural network is configured to, for each of the operations, process the operation embedding for the operation in accordance with first values of the grouper parameters to generate a grouper output that assigns the operation to an operation group from the plurality of operation groups; and

assigning each of the operation groups to a respective device from the plurality of hardware devices, comprising:
generating, for each of the operation groups, a respective group embedding by combining the operation embeddings for the operations in the operation group, and
processing the respective group embeddings for the operation groups using a placer neural network having a plurality of placer parameters,
wherein the placer neural network is configured to process the respective group embeddings in accordance with first values of the placer parameters to generate a placer output that assigns each operation group to a respective hardware device from the plurality of hardware devices.

2. The method of claim 1 , further comprising:
scheduling the operations for processing by the plurality of hardware devices by, for each operation group, placing the operations in the operation group on the hardware device to which the operation group was assigned.

3. The method of claim 1 , wherein the operations are operations necessary to train a machine learning model.

4. The method of claim 1 , wherein the operations are operations necessary to perform an inference using a machine learning model.

5. The method of claim 1 , wherein the data specifying the operations is data representing a computational graph having vertices that represent operations and edges that represent data communicated between the operations.

6. The method of claim 1 , wherein the grouper neural network is a feedforward neural network.

7. The method of claim 1 , wherein the placer neural network is a recurrent neural network.

8. The method of claim 7 , wherein the placer neural network is a sequence-to-sequence neural network.

9. The method of claim 1 , wherein the placer neural network employs learnable embeddings for the plurality of hardware devices.

10. The method of claim 1 , wherein generating, from the data specifying the operations, a respective operation embedding for each of the plurality of operations comprises:
generating a respective property embedding for each of a plurality of properties of the operation; and
combining the property embeddings to generate the operation embedding.

11. The method of claim 10 , wherein the properties comprise one or more of: an operation type of the operation, a size of outputs generated by the operation, or operations that provide input to and receive output generated by the operation.

12. The method of claim 1 , further comprising:
determining the first values of the grouper parameters and the placer parameters from initial values of the grouper parameters and placer parameters by repeatedly performing the following:
processing a current sequence of operation embeddings using the grouper neural network and the placer neural network in accordance with current values of the placer parameters and the grouper parameters to select one or more placements of the operations across the plurality of devices;
for each selected placement:
performing the operations across the plurality of devices according to the placement, and
determining a time required for the processing to complete;

adjusting the current values of the grouper parameters and the placer 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.

13. The method of claim 12 , wherein the reinforcement learning technique is a REINFORCE technique.

14. The method of claim 13 , wherein the reinforcement learning technique includes a baseline that is derived from a moving average of the required times.

15. The method of claim 12 , wherein the initial values have been determined by pre-training the grouper neural network and the placer neural network.

16. The method of claim 15 , wherein the pre-training used a different set of machine learning operations, a different plurality of hardware devices, or both.

17. The method of claim 12 , further comprising initializing the initial values of the grouper parameters and the placer parameters using a parameter initialization technique.

18. A system comprising one or more computers and one or more non-transitory computer-readable storage media storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations comprising:
receiving data specifying machine learning operations that are to be distributed across a plurality of hardware devices for processing; and
determining a placement that assigns each of the operations specified by the received data to a respective device from the plurality of hardware devices, comprising:
generating, from the data specifying the operations, a respective operation embedding for each of the plurality of operations;
grouping the operations into a plurality of operation groups, comprising processing each of the respective operation embeddings using a grouper neural network having a plurality of grouper parameters,
wherein the grouper neural network is configured to, for each of the operations, process the operation embedding for the operation in accordance with first values of the grouper parameters to generate a grouper output that assigns the operation to an operation group from the plurality of operation groups; and

assigning each of the operation groups to a respective device from the plurality of hardware devices, comprising:
generating, for each of the operation groups, a respective group embedding by combining the operation embeddings for the operations in the operation group, and
processing the respective group embeddings for the operation groups using a placer neural network having a plurality of placer parameters,
wherein the placer neural network is configured to process the respective group embeddings in accordance with first values of the placer parameters to generate a placer output that assigns each operation group to a respective hardware device from the plurality of hardware devices.

19. One or more non-transitory computer-readable storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform operations comprising:
receiving data specifying machine learning operations that are to be distributed across a plurality of hardware devices for processing; and
determining a placement that assigns each of the operations specified by the received data to a respective device from the plurality of hardware devices, comprising:
generating, from the data specifying the ope

Assignees

• Google Llc

Inventors

• Steiner Benoit
• Goldie Anna Darling
• Dean Jeffrey Adgate
• Pham Hieu Hy
• Mirhoseini Azalia
• Le Quoc V

Classifications

• G06N3/044

  Recurrent networks, e.g. Hopfield networks · CPC title

• G06N3/045

  Combinations of networks · CPC title

• G06F9/5066 Primary

  Algorithms for mapping a plurality of inter-dependent sub-tasks onto a plurality of physical CPUs (mappping at compile time, see G06F8/451) · CPC title

• G06F16/9024

  Graphs; Linked lists (G06F16/9027 takes precedence) · CPC title

• G06N3/063

  using electronic means · CPC title