Scheduling operations on a computation graph

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for scheduling operations represented on a computation graph. One of the methods includes receiving, by a computation graph system, a request to generate a schedule for processing a computation graph, obtaining data representing the computation graph, generating separators for the computation graph; and generating a schedule to perform operations represented in the computation graph, wherein generating the schedule includes initializing the schedule with a zero node, for each node in the separators: determining whether the node has any predecessor nodes in the computation graph, when the node has any predecessor nodes, adding the predecessor nodes to the schedule, and adding the node in the schedule, and adding each node in each subgraph that is neither in the separators nor a predecessor of any node in the separators on the computation graph to the schedule.

Description

[0001] Cross Reference to Related Applications

[0002] This application claims priority to U.S. Patent Application No. 62 / 875,433, filed July 17, 2019, the entirety of which is incorporated by reference herein. TECHNICAL FIELD

[0003] This specification relates to optimizing peak memory usage when executing a computation graph. BACKGROUND

[0004] Neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current values of a respective set of parameters.

[0005] To allow neural networks to perform better on a variety of tasks, e.g., speech, language, and vision processing tasks, neural networks are designed with larger sized layers, a larger number of layers, and / or increasingly complex and intricate data dependencies. However, training large and complex neural networks is difficult due to high computational cost and required memory capacity. To facilitate computation speed, specialized hardware such as GPUs and AI accelerators are often employed when training complex neural network models. In increasingly close neural networks, the peak memory capacity of the hardware does not meet the peak demand for memory when training these large models, thus turning storage capacity into a bottleneck for training. SUMMARY

[0006] This specification generally describes techniques for scheduling operations of a computation graph for execution by one or more devices to optimize peak memory usage.

[0007] According to one aspect, a method of scheduling operations of a computation graph includes receiving, by a computation graph system, a request to generate a schedule for processing a computation graph, obtaining data representing the computation graph, generating a separator for the computation graph, and generating a schedule to perform operations represented in the computation graph. The schedule defines a sequence of operations represented by the computation graph to be performed according to an ordering. The computation graph includes a plurality of nodes and directed edges, where each node represents a respective operation. Each directed edge from a respective first node to a respective second node represents that the second operation represented by the respective second node requires at least one output generated by performing the first operation represented by the respective first node as input. The separator satisfies a property that removing the nodes in the separator from the computation graph and connecting edges of other nodes to the separator causes the remaining nodes and edges of the computation graph to form a plurality of connected components. Generating the schedule includes: (1) initializing the schedule with zero nodes; (2) for each node in the separator, determining whether the node has any predecessor nodes in the computation graph, when the node has any predecessor nodes, adding the predecessor nodes to the schedule, and adding the node in the schedule, and (3) adding each node in each subgraph on the computation graph that is neither in the separator nor a predecessor of any node in the separator to the schedule.

[0008] The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following advantages.

[0009] The described systems can be used to optimize memory usage when computing a directed or undirected acyclic graph, i.e., when performing the operations represented by the graph on one or more computing devices. For example, the described techniques can be used to optimize memory usage when computing graphs for training large and complex neural networks. To optimize memory usage when computing a graph, a system generates a schedule that specifies a sequence of the operations represented in the graph in an order of execution. The schedule generated by the system is valid and reduces the peak memory usage of the computation graph by utilizing re-materialization such that only necessary operation inputs and outputs are kept in memory at any given time during execution. Some implementations can be used to meet peak memory requirements when computing a directed or undirected acyclic graph to a given memory size, e.g., the memory available on a particular hardware.

[0010] When generating the schedule, the system only rearranges the order in which each operation is executed, deciding which operations to keep in memory, if so, when and for how long to keep the operations. Because none of the operations represented in the graph are modified, this technique of scheduling operations avoids the risk of reducing the accuracy of the final output of the computation graph that exists when using other techniques to reduce memory usage, such as reusing memory regions and communication between CPU and GPU memory.

[0011] Further, the schedule can be generated in a faster clock time compared to traditional schedule generation techniques. The system can generate the schedule based on graph properties of the computation graph, including a tree width output of a tree decomposition of the computation graph. The system can further reduce memory cost using other graph properties while executing the computation graph according to the schedule.

[0012] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims. BRIEF DESCRIPTION OF DRAWINGS

[0013] Figure 1 A schedule of an example of an example directed acyclic computation graph is shown.

[0014] Figure 2 is a schematic diagram of an example operation scheduling system.

[0015] Figure 3A is an example directed acyclic computation graph.

[0016] Figure 3B is an example tree decomposition of the non-directed counterpart in Figure 3A

[0017] is an example tree decomposition of the non-directed counterpart in Figure 3C Figure 3B

[0018] Figure 4 is a flow diagram of an example process for scheduling operations of a computation graph based on a tree decomposition.

[0019] Figure 5 is a flow diagram of an example process of an operation scheduling system scheduling operations of a computation graph.

[0020] Figure 6 illustrates an example process of generating a schedule of executing operations in a computation graph based on a path decomposition.

[0021] Figure 7 illustrates an example process of modifying a computation graph so as to define a long spine.

[0022] Figure 8 illustrates an example process for scheduling operations of a computation graph based on a path decomposition. DETAILED DESCRIPTION

[0023] ​​The specification describes a system implemented as a computer program on one or more computers in one or more locations that schedules operations for processing an input computation graph to reduce peak memory requirements for executing the operations. To reduce peak memory usage, the system can receive an input computation graph representing operations and input dependencies of the operations on other operations in the input computation graph. The system can then generate a schedule representing a sequence of operations to execute. The system can generate the schedule by identifying where intermediate inputs between operations can be re-materialized to reduce or completely eliminate the need to store intermediate inputs in memory until they are needed.

[0024] A directed acyclic computation graph can represent operations for processing an input by a neural network. A neural network can be configured to receive any kind of numerical data input as a network input and generate any kind of network output based on the network input, i.e., any kind of score, classification, or regression output. For example, in image classification, if the input to the neural network is an image or features that have been extracted from an image, the output generated by the neural network for a given image can be scores for each of a set of object classes, where each score represents an estimated likelihood that the image contains an object belonging to that class. The neural network processes the input to generate the network output by performing a plurality of operations. Each operation can receive zero or more inputs and can generate an output. Examples of operations that can be represented in a graph representing a neural network include activation functions, e.g., non-linear mappings that take an intermediate variable based on input from a previous layer and generate an output for a next layer, forward gradient propagation operations, e.g., operations related to computing and storing intermediate variables for each layer and outputs for the neural network in order from a first input layer to an output layer, backward gradient propagation operations, e.g., operations related to computing and storing gradients of a loss function with respect to parameters of each layer in the neural network from the output layer to the first input layer, and optimization operations, e.g., operations used by an optimizer such as (stochastic) gradient descent, momentum, RMSProp, and Adam to assist in updating parameters of the neural network based on computed gradients of the loss function. The inputs and outputs of any given operation can be, e.g., scalar values, categorical values, vectors, matrices, or higher-order tensors.

[0025] Figure 1 A schedule for the example directed acyclic computation graph 100 is shown. As Figure 1As shown, the plurality of operations can be represented as a directed acyclic computation graph 100 having a plurality of nodes (103, 105, 107, 109, and 111) and edges (135, 157, 179, 191, 131, and 159). Each node of the computation graph represents a respective operation of the plurality of operations. Each node consumes a set of inputs from its incoming edge(s), performs its respective operation on the inputs, and outputs the results of the operation to any nodes connected to the node by outgoing edges. For example, node 105 takes as input the output of the node 103 operation via incoming edge 135, performs its respective operation, and sends the output via two outgoing edges 157 and 159 as input to nodes 107 and 109, respectively. In general, throughout the computation, both the inputs and outputs of the computation graph should be stored in memory, and the order in which the computation nodes are executed, i.e., the schedule, will determine the peak memory usage. As Figure 1 As shown, when storing the output of each node takes one memory unit, computing the nodes in the order of {103, 105, 107, 109, 111} will require four memory units, while computing them in the order of {103, 105, 107, 109, 103, 111} will require only three memory units. The latter order involves here re-materializing the output of node 103, i.e., generating the output of node 103 again before it is needed as input to another node, instead of always keeping it in memory. The combination of ordering the execution of the operations and possibly re-materializing some intermediate nodes is referred to in this specification as a schedule.

[0026] Figure 2 An example operation scheduling system 200 is shown that includes a tree decomposition engine 215 and a schedule engine 220.

[0027] The operation scheduling system 200 is implemented as a computer program on one or more computers in one or more locations that implement the systems, components, and techniques described below.

[0028] The operation scheduling system 200 can take as input a computation graph 205, e.g., a directed acyclic graph. The computation graph 205 can be directed acyclic.

[0029] The tree decomposition engine 215 takes as input the computation graph 205 and performs a tree decomposition process to generate a plurality of sets of tree decomposition outputs 230, (1) a plurality of sets of nodes of the computation graph 205, and (2) a directed path connecting each set of the plurality of sets. In short, the set of nodes in (1) can be referred to as a “bag,” and the directed path in (2) can be referred to as a “tree.” Note here that by applying one or more decomposition processes to each computation graph, the computation graph can have a plurality of different possible tree decompositions.

[0030] The scheduling engine 220 takes the original computation graph 205 and the tree decomposition output 230 as input and outputs a schedule 210 for the computation graph 205.

[0031] More specifically, a directed acyclic computation graph can be represented by G = (V, E), where V is a set of nodes that includes all nodes of the computation graph G, and E is a set of edges that includes all directed edges of the computation graph G. For any two nodes u and v in the set of nodes V, a directed edge represents (u, v) a data dependency from u to v. A data dependency from u to v means that the operation represented by node u generates an output that is input to the operation represented by node v. Thus, the node u operation must be executed before the node v operation is executed.

[0032] Additionally, one or some of the nodes are designated as a set of final nodes f. At the end of the computation of the computation graph, the nodes in f need to be kept in memory, e.g., because the outputs of the nodes in f are the outputs of the machine learning model represented by the computation graph. In this specification, we assume that these dummy nodes have out-degree zero, i.e., the output of each node in the set of nodes f will not be used by any other node in the same graph.

[0033] A schedule for a computation graph is a sequence of nodes belonging to V that has the following properties: (1) the set of final nodes f is represented in the schedule, and (2) each node in the schedule only occurs after all of its predecessors have been added to the schedule at least once. For a directed edge (u, v), a predecessor of node v is a node u that is connected by an outgoing path from u to v in the computation graph. The second requirement prevents the schedule from including a node in the schedule before all other nodes in the computation graph that the node has a data dependency with. If a node’s predecessor nodes have their own one or more predecessors, those predecessor nodes appear in the schedule before the predecessor node, and so on until a node without predecessors is reached, e.g., an input node that receives input to the computation graph.

[0034] Additionally, a node can appear in the schedule more than once. As described below, the system can generate a schedule that can add the same node to the schedule more than once, trading additional computation time for reduced memory requirements to execute the computation graph.

[0035] A schedule naturally implies a time and memory bound for the computation graph G. The length of a node includes the time required to execute the corresponding operation of that node, and the length of a schedule is the sum of the lengths of each node in the schedule of the computation graph G. Thus, the lower bound on the length of a schedule is the time required to execute each operation of the computation graph G once.

[0036] Peak memory usage for scheduling is defined by the maximum memory required at any given time step according to the scheduling execution graph. The maximum memory required at a time step is defined by summing the magnitudes of the tensor outputs of all direct predecessors of a node at a time step. Therefore, the set of tensors that need to be held in memory at a time step depends on the current node, each predecessor of the current node, and all predecessors of each predecessor of the current node.

[0037] A computation graph G can be decomposed into (1) multiple sets, each set comprising one or more nodes of the computation graph (multiple bags) and (2) paths (trees) connecting each of the multiple sets. Each bag is a subset of the nodes in the computation graph, and a tree is a tree structure connecting all the bags. A tree is a non-directed acyclic computation graph in which any two nodes are connected by exactly one edge, or equivalently, each node has at most one nearest predecessor. This decomposition is called tree decomposition, which satisfies the following three properties: (1) every node in the original computation graph G is in at least one bag; (2) for each directed edge (u,v) in the original computation graph, the two nodes of that edge appear together in at least one bag; and (3) for each node in the original computation graph, the bags containing the same node form a connected subgraph of the tree. In other words, a node can be assigned to more than one bag, and all bags containing a given node must form a connected subgraph of the tree. Each connected subgraph can be viewed as a subtree.

[0038] like Figure 3A As shown, the directed acyclic computation graph 301 has nodes A to J, and edge 313 is directed because node I is a predecessor of node J (311). In some implementations, the system transforms the directed computation graph into an undirected computation graph by removing the directions of edges in computation graph 301 before performing tree decomposition. Figure 3B As shown, in the undirected computation graph 303 based on computation graph 301, each edge has no direction. Figure 3C An example tree decomposition output 305 is shown for an undirected computation graph 303, where the tree decomposition output is (1) multiple packets and (2) a tree connecting these packets. The tree decomposition output 305 satisfies the three properties of tree decomposition:

[0039] (1) Each node in the original undirected computation graph 303 is in at least one bag. For example, node D is in bags 331, 333, and 335. For another example, node B is in bags 333, 335, and 337.

[0040] (2) For each edge in the computation graph 303, the two end nodes of the edge are assigned to the same bag at least once in the tree decomposition. For example, edge 323 in the computation graph 303 has two end nodes A and B, and now in the tree decomposition output 305, nodes A and B are together in bag 335. As another example, edge 325 in the computation graph 303 connects two end nodes C and D, so in the tree decomposition output 305, nodes C and D both appear in bag 333 and bag 341.

[0041] (3) For any node that is assigned to more than one bag, all bags containing the node can form a sub-tree. For example, bags 333, 335, and 337 all contain node B, and bags 333, 335, and 337 and the edges connecting these bags can together form a connected component subgraph of the tree decomposition output 305, as emphasized by the dashed lasso in Figure 3C .

[0042] The tree width (treewidth) tw(G) of a computation graph G is the minimum width of all tree decomposition outputs, and the width of a tree decomposition output is defined as the maximum bag size minus 1, where the bag size of a given bag is based on the number of nodes in the bag. For example, the width of the tree decomposition output 305 is 1. Suppose there are only two possible tree decomposition outputs from the undirected computation graph 303, one with width 2 and the other with width 5. Then, the minimum width among the three tree decomposition outputs is 1, i.e., the tree width tw(G) of the computation graph 303 is 1.

[0043] In some implementations, the system maintains a width threshold, i.e., a predetermined threshold, such that the width of any tree decomposition output of the computation graph G is less than the width threshold. In other words, the system does not select any tree decomposition output unless the tree has a width less than the tree width.

[0044] A separator is a set of nodes that, when removed from a computation graph in a way that the nodes in the separator and the edges linking to each node in the separator are removed, separates the remaining nodes and edges of the computation graph, i.e., the nodes and edges that are not in the separator, into multiple connected component subgraphs such that the remaining nodes and edges form the multiple connected component subgraphs of the computation graph. Each connected component subgraph includes multiple nodes connected to each other with edges, but no connected component subgraph is connected to any other connected component subgraph.

[0045] A separator for tree decomposition is a bag that, when removed from the tree structure in a way that the nodes in the separator and the edges connecting nodes from other bags to the nodes in the separator are removed, is selected from all bags in the tree decomposition that form multiple connected sub-trees.

[0046] A balanced separator for tree decomposition is a separator that, when removed from the tree structure as described above, leaves the remaining bags and edges forming multiple connected sub-trees, where the size of the sub-trees is at most half the size of the original tree of the tree decomposition.

[0047] In some implementations, the system can generate partitions of connected sub-graphs, where each sub-graph is much smaller in size than the original graph.

[0048] For example, the system can do this by generating a balanced separator for a tree decomposition of the computation graph and then removing the balanced separator from the tree decomposition.

[0049] The system can then efficiently generate a schedule for the computation graph by recursively generating a schedule for each connected sub-graph resulting from the partitioning. The recursion can be repeated until a threshold condition is met. The threshold condition can be, for example, that the peak memory requirement for executing the directed or undirected acyclic graph falls below a threshold corresponding to the available memory for executing the directed or undirected acyclic graph on a particular hardware.

[0050] In some implementations, this process can be performed in parallel for each connected sub-graph. Since the partitions are characterized by the fact that each connected sub-graph is independent of each other, i.e., not connected together by any edges, the system can independently generate a schedule for each sub-graph without data dependency issues.

[0051] The system can then add each schedule obtained for each sub-graph to an effective final schedule.

[0052] A system or other device configured to execute the computation graph can then execute the computation graph according to the final schedule.

[0053] Figure 4 A flowchart of an example process 400 illustrating the operations for scheduling a computation graph based on tree decomposition. For convenience, the process 400 will be described as performed by a system of one or more computers located in one or more locations. For example, an operation scheduling system, e.g., the operation scheduling system 200 of Figure 2 , appropriately programmed, can perform the process 400.

[0054] To schedule operations using the process 400, the system takes as input (a) the acyclic computation graph G and (b) a tree decomposition having (1) a plurality of bags and (2) a tree obtained from the computation graph G by the tree decomposition process. In some embodiments, the system takes as input (a) a subset of nodes of the computation graph G as a set of constraints, (b) a tree decomposition based on the set of constraints, (c) a subset of nodes of the set of constraints as a set of computations, and (d) the computation graph G. That is, a device that executes the operations of a computation graph according to a schedule can use the annotations of any given node in the schedule to determine whether to store the output of the operation represented by the node in memory or discard it once it has been consumed as input by another operation. In some embodiments, adding annotations to a node can be performed by defining a scalar variable associated with the node that describes whether to save the output of the node in memory.

[0055] As previously described, the system finds a balanced separator of the computation graph G (410). The balanced separator can cause the remaining bags and edges in the tree decomposition to form a set of connected component subgraphs (subtrees) after removal, where the size of the connected component subgraphs is at most half the size of the original tree of the tree decomposition.

[0056] The system then determines the set of connected component subgraphs after removing the balanced separator and all nodes associated with the separator from all bags (420). A bag containing any identical nodes from the computation graph G can form a connected component subgraph using the corresponding tree edges of the tree decomposition. Because of the way the separator is identified, if a node is in two or more connected component subgraphs, the node must also be in the balanced separator. Thus, it is guaranteed that the remaining nodes and edges can form the set of connected component subgraphs after removal of the balanced separator.

[0057] Next, the system initializes a schedule list (430). Typically, the schedule list is initialized to an empty set.

[0058] Thereafter, the system finds all predecessors of each node in the separator and adds the node, optionally with a corresponding annotation, to the schedule (440). For example, the system can iterate over the nodes in the separator according to a random order or according to a fixed order when generating the schedule. Specifically, for a given node in the separator, the system finds all direct predecessors of the node and adds the direct predecessors with annotations to the schedule, then the system finds all predecessors of the direct predecessors and adds all predecessors with annotations to the schedule, and finally the system adds the node with annotations to the schedule until all predecessors of the node are found and added to the schedule.

[0059] Further, the system finds all final nodes in each connected component subgraph and adds them to the schedule (450). All final nodes are nodes that are not in the separator and are not predecessors of any other node. Each final node will also be added to the schedule optionally with an annotation.

[0060] Finally, the system optimizes the schedule length by releasing unnecessary balanced separator nodes (460). A unnecessary balanced separator node is a node that was originally included in the balanced separator by step 410 but is not added to the schedule.

[0061] In some embodiments, steps 440 and 450 can be called recursively based on a divide-and-conquer strategy. For example, when step 440 is first executed, some internal parameters, such as the limit set and the computation set, are updated. The system then calls again from step 410 with the updated internal parameters as input and performs the same process. As the system recursively calls the process, the updated internal parameters can affect the range of searching predecessors and / or tracking the schedule. The system recursively stops calling the process when certain criteria are met. For example, the criteria can be that the input parameters of the last recursive call are undefined. As another example, the criteria can include that the peak memory requirement for executing the schedule falls below a threshold, such as the available memory for executing the schedule on a particular hardware. Similarly, for step 450, the system calls again recursively from step 410 and performs the same process. Recursive calling guarantees that the size of the tree decomposition decreases at each recursive call, thereby guaranteeing an efficient schedule, and each operation is executed at least once.

[0062] Figure 5 A flowchart of an example process 500 illustrating the operation of the operation scheduling system to schedule a computation graph is shown. Process 500 will be described as being performed by a system of one or more computers located in one or more locations. For example, an operation scheduling system, such as operation scheduling system 200 of Figure 2 , appropriately programmed, can perform process 500.

[0063] The operation scheduling system receives a request to generate a schedule for processing a computation graph (510), obtains data representing the computation graph (520), generates a separator for the computation graph (530), and then generates a schedule to execute the operations represented in the computation graph (540). To generate the schedule, the system first initializes a schedule of the operations to be executed in the computation graph (540a), then for each node in the separator, determines whether the node has any predecessor nodes, adds the predecessor nodes to the schedule if there are any, and adds the node to the schedule (540b), and finally adds to the schedule each node in each connected component subgraph that is neither in the separator nor a predecessor to any node in the separator on the computation graph (540c). Generating the separator in step 530 is described aboveFigure 4 described in more detail in step 410. The predecessors of the nodes from the separator are added to the schedule in step 540b, which is described in more detail above in step 440. The nodes that are neither in the separator nor predecessors are added to the schedule in step 540c, which is described in more detail above in step 450. Figure 4 Figure 4

[0064] In some embodiments, the system can alternatively use a path decomposition, i.e., instead of a tree decomposition, to generate a schedule for a computation graph. In these embodiments, the computation graph defines a directed path that traverses each node in the computation graph.

[0065] A path in a computation graph G is a sequence of distinct nodes in the computation graph such that each pair of consecutive nodes in the sequence has an edge from the computation graph connecting them. A path is called a directed path when each pair of nodes in the sequence is connected by a directed edge, where the edge goes from the earlier node in the sequence to the node immediately after it in the sequence. In this specification, a directed path that traverses each node in a computation graph is also called a long spine of the computation graph. In some embodiments, if a computation graph does not include a long spine, the system can modify the computation graph, e.g., by adding nodes and edges, so that the computation defines a long spine.

[0066] For a directed acyclic computation graph that defines a long spine, the system can generate a schedule of execution operations based on the long spine of the computation graph. For a directed acyclic computation graph that does not define a long spine, the system can generate a schedule of execution operations in the graph by first generating a path decomposition for the computation graph. A path decomposition is a particular tree decomposition of a computation graph that specifies (1) a plurality of sets, each set comprising one or more nodes of the computation graph (a bag) and (2) a directed path (a tree) through each set in the plurality of sets, but with the additional requirement that the tree defines a path that sequentially traverses each bag in the plurality of bags. The plurality of sets of nodes of the path decomposition are ordered according to the path defined by the path decomposition. For example, the plurality of bags is defined as a set of bags X in the path decomposition, e.g., X = {X1, X2, X3}, where X1, X2, X3 represent respective bags, or sets of nodes in the plurality of sets in the path decomposition, in an order such that bag X1 precedes bag X2 (which precedes bag X3) in the path.

[0067] In some embodiments in which a computation graph does not define a long spine, the system can modify the computation graph to define a long spine so that a path decomposition can still be formed from the computation graph. In one embodiment, a long spine can be added to a computation graph while controlling the path width of its path decomposition. Note here that the definition of the path width pw(G) of a computation graph G is the minimum width of any path decomposition of the computation graph G, similar to the definition of the tree width tw(G).​​

[0068] Figure 6 An example process 600 is illustrated that generates a schedule of executing operations in a computation graph based on a path decomposition. The process 600 will be described as being performed by a system of one or more computers located in one or more locations. For example, an operation scheduling system, such as the operation scheduling system 200 of Figure 2 FIG. 1, suitably programmed, is capable of performing the process 600.

[0069] First, the scheduling operation system receives a request to generate a schedule for processing a computation graph (610), the schedule defining a sequence of nodes in the computation graph, and each node representing an operation.

[0070] The system then obtains data representing a computation graph having a plurality of nodes and directed edges (620). Each directed edge connects a preceding node to a consecutive node that takes an output of the preceding node as an input.

[0071] Next, the system determines whether the computation graph already defines a long spine that traverses each node in the computation graph (630).

[0072] If the computation graph already defines a long spine, the system generates a path decomposition for the computation graph (640), where the path decomposition includes (1) a plurality of node sets (packets) and (2) a directed path through each node set of the plurality of node sets.

[0073] If the computation graph does not already define a long spine, the system modifies the computation graph to add a new long spine that traverses each node in the computation graph (680). Adding a long spine to a computation graph while controlling a path width of a path decomposition of the computation graph will be described in more detail below. Once the long spine is added to the computation graph, the system generates a path decomposition for the computation graph as described above (640).

[0074] Next, the system initializes a schedule of executing operations in the computation graph (650). For example, the initial schedule can be an empty set with zero nodes.

[0075] Thereafter, the system identifies a particular node set, i.e., packet, from all node sets in the path decomposition, i.e., from all packets (660).

[0076] For each node in the particular node set and in an order according to the long spine in the computation graph, the system determines whether the node has any predecessors, adds each predecessor of the node to the schedule, and finally adds the node to the schedule (670). Adding these nodes to the schedule will be described in more detail below.

[0077] As described above, long ridges can be added to a computation graph while controlling the width of its path decomposition. In some implementations, long ridges are added to the computation graph such that the path width of the computation graph is only increased by a factor of 2.

[0078] Figure 7 Figure illustrates an example process 700 of modifying a computation graph to maintain a low path width of the computation graph while defining long ridges to the computation graph. The process 700 will be described as being performed by a system of one or more computers located in one or more locations. For example, an operation scheduling system, such as the operation scheduling system 200 of Figure 2, that is suitably programmed, can perform the process 700. Figure 2

[0079] The system first adds a node u to the computation graph, connects it to all other nodes, and adds it to all node sets (all bags) in the path decomposition (710), then the system sets node u as the current final node of the directed paths in the path decomposition (720). In some implementations, the system can be called recursively to obtain the current final node that has no in-neighbors. Note here that nodes in the computation graph are processed in topological order, so any node is not processed until its in-neighbors have been processed. An in-neighbor of a node is a node that is connected to the node in topological order. For example, there must be at least one in-neighbor node r of a node v in the path decomposition such that when node r has been processed, node v has no unexplored in-neighbor nodes, so v should then be processed.

[0080] Next, the system adds a path starting from node u to connect with the immediate next node (730). The path should cover all nodes that are processed in a recursive manner. Then, the system updates the path decomposition of the computation graph to handle the new path (740), and finally, the system removes node u from the computation graph and from all bags of the updated path decomposition of the graph (750).

[0081] When generating a schedule for a graph, the system can rely on structural properties of the path decomposition of a computation graph with long ridges. For example, any directed path (e.g., a long ridge) that passes between two node sets (two bags) of the path decomposition must also pass between any intermediate node sets (intermediate bags). As another example, each last node of a bag in the path decomposition of a computation graph can be ordered based on (1) the final node of a directed path (e.g., a long ridge) and (2) the bag that contains the final node.

[0082] As another example, given a computation graph with a path decomposition and one of its subgraphs, removing all nodes from each bag of the computation graph that are not in the subgraph results in the remaining part of the path decomposition being a valid path decomposition for the subgraph. By removing nodes, edges connecting these nodes are also removed. Thus, the remaining nodes and edges of the original path decomposition form a valid path decomposition.​

[0083] As another example, the scheduling of subgraphs of a computation graph based on path decomposition can be interleaved, and the interleaved scheduling of the computation graph can be merged, flattened, or compressed together to form a schedule for the computation graph.

[0084] The aforementioned structural features allow the divide-and-conquer strategy to recursively schedule operations on the right and left sides of the path decomposition. This recursion can be repeated until one or more criteria are met. These criteria may include peak memory requirements for performing the scheduling falling below a threshold, such as available memory for performing the scheduling on specific hardware.

[0085] Figure 8 The illustration depicts an example process 800 for scheduling operations on a computation graph based on path decomposition. This process 800 will be described as being executed by a system of one or more computers located in one or more locations. For example, an operation scheduling system that is appropriately programmed, such as... Figure 2 The operation scheduling system 200 is capable of executing this process 800.

[0086] Figure 8 The system first divides the path decomposition of the computation graph into left path decomposition and right path decomposition (810). The computation graph has been defined with long ridges that traverse each node in the computation graph in sequence. The partitioning of the path decomposition is based on a specific node in the computation graph. Nodes ordered after a specific node in the ridge are removed from each bag in the path decomposition. For example, the set of multiple sets of nodes obtained from the path decomposition (e.g., the set of all bags) is defined as X, where X = (X1, X2, ..., X...). b Define u as a specific node. Then, the system removes any nodes in X that follow node u in order according to the long ridge, resulting in a new path decomposition X' = (X'1, X'2, ..., X'). b The set of nodes (packets) containing a specific node u will divide the new path decomposition X' into a left path decomposition and a right path decomposition. For example, if a specific node u is in pack X... i In the case of ', the new path decomposition X' is divided into left path decomposition X'. l =(X' i-1 , X' i-2 , ..., X'1) and right decomposition X r =(X' i+1 , X' i+2 , ..., X' b ).

[0087] Then, the system takes the left path decomposition as input, determines whether any predecessor of a specific node u exists in the multiple sets of nodes in the left path decomposition according to the order of the long ridges, and returns the interleaved schedule (820) in which the annotations are annotated to all predecessors of the specific node u. Note here that steps 810 and 820 can be called recursively to further divide the left path decomposition into smaller sizes, which reflects the divide-and-conquer technique.

[0088] Similarly, the system takes the right path decomposition as input, determines whether any predecessor of a specific node u exists in the multiple sets of nodes in the right path decomposition according to the order of the long ridges, and returns an interleaved schedule annotating all predecessors to the specific node u (830). Again, steps 810 and 830 can be called recursively. In some implementations, steps 820 and 830 can share the same system architecture without distinguishing between left path decomposition and right path decomposition. For example, steps 820 and 830 can be the same system that recursively calls back to step 810.

[0089] Finally, the system outputs a schedule for computation graph by combining interleaved schedules from both left and right path decompositions (840). To this end, the system can perform merging, flattening, and / or compressing operations to integrate interleaved schedules into the schedule based on the structural characteristics of the path decomposition.

[0090] As an explanation of integrated interleaved scheduling, we first define σ as the scheduler and u as the node. Then, we define the cascading operation. Make The schedule σ represents the node u followed by the schedule σ′. Define a set of nodes X, where the interleaved schedules on X are a set of tuples {L = {<σ1, x1>, <σ2, x2>, ..., <σ′>}. k x k >} makes each x i ∈X and each σ i It's scheduling. Then the merge operation is defined as...

[0091]

[0092] For each i,

[0093] Next, the flattening operation will be defined as

[0094]

[0095] make And i1 < i2 < ... < i k′ The index is such that X′={x ij}, where k′ is the length of the node set X′. For convenience, define where i1= 0. Finally, the compression operation is defined as;

[0096]

[0097] Furthermore, Condense(X', L) = Condense(X', Merge(L)) and Flatten(L) = Flatten(Merge(L)) = Flatten(Condense(X', L)).

[0098] The description herein, together with the systems and computer program components, utilizes the term "configured" in connection with systems that are to be configured to perform a particular operation or action. This means that the system has installed on it software, firmware, hardware, or a combination of them that in operation causes the system to perform the operation or action. In connection with one or more computer programs that are to be configured to perform a particular operation or action, this means that the one or more programs include instructions that when executed by a data processing apparatus cause the apparatus to perform the operation or action.

[0099] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or additionally, the program instructions can be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

[0100] The term "data processing apparatus" refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0101] A computer program, which can also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it can be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, sub programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and

[0102] In this specification, the term "database" is used broadly to refer to any set of data: the data need not be structured in any particular way, or structured at all, and it can be stored on storage devices in one or more locations. Thus, for example, an index database can include multiple sets of data, each of which can be organized and accessed differently.

[0103] Similarly, in this specification the term "engine" is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers.

[0104] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0105] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read-only memory or a random access memory or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto-optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

[0106] Computer readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media and memory devices, including by way of example semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto-optical disks; and CD-ROM and DVD-ROM disks.

[0107] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device of the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device of the user, e.g., a smart phone that is running a messaging application, and receiving responsive messages from the user in return.

[0108] The data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing the public and compute-intensive portions of machine learning training or production, i.e., inference, workloads.

[0109] The machine learning models can be implemented and deployed using a machine learning framework, such as a TensorFlow framework, a Microsoft Cognitive Toolkit framework, an Apache Singa framework, or an Apache MXNet framework.

[0110] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back-end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front-end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.

[0111] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at a server from the device.

[0112] While this specification contains many specifics, these should not be construed as limitations on the scope of any invention or of any claims hereafter presented, but rather as descriptors of features that can be part of some embodiments. Certain features that are described in this specification in the context of separate embodiments can also combine in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments. Moreover, although features can be described above as acting in certain combinations and even initially claimed as such, one or more features from a claimed combination can in some cases be excised from the combination and the claimed combination can be directed to a sub-combination or a variation of a sub-combination.

[0113] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order described or in sequential order, or that all illustrated operations be performed, to implement desirable results. In certain circumstances, multitasking and parallel processing can be beneficial. Additionally, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems are often integrally combined on a single software product or packaged into multiple software products.

[0114] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the acts recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order or sequential order illustrated in order to achieve desirable results. In some instances, multitasking and parallel processing can be beneficial.

Claims

1. A method for scheduling operations on a computation graph, the method comprising: The computation graph system receives a request to generate a schedule for processing the computation graph, wherein the schedule defines a sequence of operations to be executed according to an order represented by the computation graph; Obtain data representing the computation graph, which includes multiple nodes and directed edges, wherein each node represents a corresponding operation. Each directed edge from the corresponding first node to the corresponding second node represents a second operation represented by the corresponding second node that requires at least one output generated by executing the first operation represented by the corresponding first node as input; A separator for generating the computation graph, wherein the separator satisfies the property that removing nodes from the separator and connecting other nodes to edges of the separator from the computation graph results in the remaining nodes and edges of the computation graph forming multiple connected components; and Generating the schedule to execute the operations represented in the computation graph, wherein generating the schedule includes: The schedule is initialized using a zero node; For each node in the separator: Determine whether the node has any predecessor nodes in the computation graph. When the node has any predecessor node, add the predecessor node to the schedule, and Add the node to the schedule, and Each node in each subgraph of the predecessor of a node that is neither in the separator nor any node in the separator is added to the schedule.

2. The method of claim 1, further comprising processing the computation graph according to the scheduling.

3. The method according to claim 1, wherein, Each operation represented in the computation graph is an operation associated with processing the input of the neural network.

4. The method according to claim 1, wherein, Generating the separator includes generating a tree decomposition for the computation graph.

5. The method according to claim 4, wherein, The separator is a balanced separator for the tree decomposition, and wherein generating the tree decomposition for the computation graph includes generating the tree decomposition to be smaller than a predetermined width.

6. The method according to any one of claims 1 to 5, wherein, The schedule includes annotations, each annotation corresponding to a corresponding node in the schedule and indicating whether the output generated by performing the corresponding operation corresponding to the corresponding node is stored in memory.

7. A method for scheduling operations on a computation graph, the method comprising: The computation graph system receives a request to generate a schedule for processing the computation graph, wherein the schedule defines a sequence of operations to be executed according to an order represented by the computation graph; Obtain data representing the computation graph, which includes multiple nodes and directed edges. Each node represents a corresponding operation, and Each directed edge from the corresponding first node to the corresponding second node represents a second operation represented by the corresponding second node that requires at least one output generated by executing the first operation represented by the corresponding first node as input; Determine whether the computation graph includes a first directed path through the computation graph that traverses each node in the computation graph, including traversing a single node; In response to determining that the computation graph includes the first directed path: Generate a path decomposition, wherein the path decomposition defines (i) a plurality of sets, each of the plurality of sets comprising one or more nodes of the computation graph, and (ii) a second directed path through each of the plurality of sets. The schedule is initialized using a zero node. A set that identifies one of the plurality of sets of the path decomposition, including the aforementioned node. For each node in the set and in the order according to the first directed path: Determine whether the node has any predecessor nodes in the computation graph. When the node has any predecessor node in the computation graph, add the predecessor node to the schedule, and The node is added to the schedule after the predecessor node.

8. The method according to claim 7, further comprising: It is determined that the computation graph does not include the first directed path; In response to determining that the computation graph does not include the first directed path, the computation graph is modified, wherein, as a result of modifying the computation graph, the computation graph defines a third directed path through each node of the computation graph, including through the first node; and In response to the request, the schedule is generated, including: The schedule is initialized using a zero node. Generate the path decomposition that defines the multiple sets. The set that identifies the plurality of sets includes the node mentioned above. For each node in the set and in the order according to the third directed path, add the node to the schedule: Determine whether the node has any predecessor nodes in the computation graph. When the node has any predecessor node in the computation graph, add the predecessor node to the schedule, and The node is added to the schedule after the predecessor node.

9. The method according to claim 7, further comprising: After generating the path decomposition for the computation graph: The path decomposition of the computation graph is divided into a left path decomposition and a right path decomposition, wherein the left path decomposition includes each of the plurality of sets ordered before the first set according to the second directed path, and wherein the right path decomposition includes each of the plurality of sets ordered after the first set according to the second directed path; and Determining that one or more nodes in the computation graph are predecessors of the node includes performing an operation, the operation including: One or more nodes in the computation graph are determined to be predecessors of the node and are included in the left path decomposition, and in response, each of the one or more nodes in the left path decomposition is added to the schedule, and One or more nodes of the computation graph are determined to be predecessors of the node and are included in the right path decomposition, and in response, each of the one or more nodes in the right path decomposition is added to the schedule.

10. The method according to claim 7 further comprises: After generating the path decomposition, all nodes in the path decomposition that were ordered after the first node in the first directed path are removed.

11. The method according to any one of claims 7 to 10, wherein, Generating the schedule includes optimizing the separator size.

12. A system for scheduling operations on a computation graph, the system comprising one or more computers and one or more storage devices storing instructions, the instructions, when executed by the one or more computers, causing the one or more computers to perform operations according to any one of claims 1 to 11.

13. A computer storage medium storing instructions, said instructions, when executed by said one or more computers, causing said one or more computers to perform the operations of the method according to any one of claims 1 to 11.

Images