Constrained device placement using neural networks
AssigneeGoogle Inc.
InventorsXinfeng Xie, Azalia Mirhoseini, James Laudon, Phitchaya Mangpo Phothilimthana, Sudip Roy, Prakash Janardhana Prabhu, Ulysse Beaugnon, Yanqi Zhou
Life of this patent
Estimating the term…
Abstract
Systems and methods for determining a placement for computational graph across multiple hardware devices. One of the methods includes generating a policy output using a policy neural network and using the policy output to generate a final placement that satisfies one or more constraints.
Claims21
A method performed by one or more computers, the method comprising:
- obtaining graph data specifying a computational graph to be executed on a plurality of hardware devices, the computational graph comprising a plurality of nodes representing operations and a plurality of edges that represent data dependencies between the operations represented by the plurality of nodes;
- obtaining constraint data specifying one or more constraints on the execution of the computational graph;
- and
- generating a final placement that assigns each node of the computational graph to a respective hardware device of the plurality of hardware devices and that satisfies the one or more constraints, the generating comprising:
- processing the graph data using a placement neural network to generate a policy output that comprises, for each node, a respective score distribution that includes a respective score for each of the plurality of hardware devices;
- and
- generating the final placement by assigning the nodes one after the other according to a node order, comprising, for each particular node:
- identifying a subset of the hardware devices that would satisfy the one or more constraints if the particular node were assigned the hardware device given the assignment of any nodes that precede the particular node in the node order;
- and
- assigning, using the policy output, the particular node to a hardware device in the subset of devices.
The method of claim 10, wherein the reward measures a throughput of the execution of the computational graph, a latency of the execution of the computational graph, or both.
(canceled)
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 comprising:
- obtaining graph data specifying a computational graph to be executed on a plurality of hardware devices, the computational graph comprising a plurality of nodes representing operations and a plurality of edges that represent data dependencies between the operations represented by the plurality of nodes;
- obtaining constraint data specifying one or more constraints on the execution of the computational graph;
- and
- generating a final placement that assigns each node of the computational graph to a respective hardware device of the plurality of hardware devices and that satisfies the one or more constraints, the generating comprising:
- processing the graph data using a placement neural network to generate a policy output that comprises, for each node, a respective score distribution that includes a respective score for each of the plurality of hardware devices;
- and
- generating the final placement by assigning the nodes one after the other according to a node order, comprising, for each particular node:
- identifying a subset of the hardware devices that would satisfy the one or more constraints if the particular node were assigned the hardware device given the assignment of any nodes that precede the particular node in the node order;
- and
- assigning, using the policy output, the particular node to a hardware device in the subset of devices.
A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one or more computers to perform operations comprising:
- obtaining graph data specifying a computational graph to be executed on a plurality of hardware devices, the computational graph comprising a plurality of nodes representing operations and a plurality of edges that represent data dependencies between the operations represented by the plurality of nodes;
- obtaining constraint data specifying one or more constraints on the execution of the computational graph;
- and
- generating a final placement that assigns each node of the computational graph to a respective hardware device of the plurality of hardware devices and that satisfies the one or more constraints, the generating comprising:
- processing the graph data using a placement neural network to generate a policy output that comprises, for each node, a respective score distribution that includes a respective score for each of the plurality of hardware devices;
- and
- generating the final placement by assigning the nodes one after the other according to a node order, comprising, for each particular node:
- identifying a subset of the hardware devices that would satisfy the one or more constraints if the particular node were assigned the hardware device given the assignment of any nodes that precede the particular node in the node order;
- and
- assigning, using the policy output, the particular node to a hardware device in the subset of devices.
Dependencies are inferred from the claim text.