Hierarchical device placement with reinforcement learning

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 for jointly training a grouper neural network having grouper parameters and a placer neural network having placer parameters, the operations comprising:

• receiving data specifying machine learning operations that are to be distributed across a plurality of hardware devices for processing;
• generating, from the data specifying the operations, a respective operation embedding for each of the machine learning operations;
• processing the operation embeddings using the grouper neural network to generate a group output that assigns the machine learning operations to a plurality of operation groups in accordance with current values of grouper parameters of the grouper neural network;
• processing the group output using the placer neural network to generate a placer output that assigns each of the operation groups to a respective hardware device from the plurality of hardware devices in accordance with current values of placer parameters of the placer neural network;
• performing the machine learning operations across the plurality of hardware devices according to the placer output;
• determining a time required for the processing to complete;
• and
• 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.

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 for jointly training a grouper neural network having grouper parameters and a placer neural network having placer parameters, the operations comprising:

• receiving data specifying machine learning operations that are to be distributed across a plurality of hardware devices for processing;
• generating, from the data specifying the operations, a respective operation embedding for each of the machine learning operations;
• processing the operation embeddings using the grouper neural network to generate a group output that assigns the machine learning operations to a plurality of operation groups in accordance with current values of grouper parameters of the grouper neural network;
• processing the group output using the placer neural network to generate a placer output that assigns each of the operation groups to a respective hardware device from the plurality of hardware devices in accordance with current values of placer parameters of the placer neural network;
• performing the machine learning operations across the plurality of hardware devices according to the placer output;
• determining a time required for the processing to complete;
• and
• 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.