Hierarchical device placement with reinforcement learning

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.