Scheduling artificial intelligence model partitions based on reversed computation graph
Techniques are disclosed for scheduling artificial intelligence model partitions for execution in an information processing system. For example, a method comprises the following steps. An intermediate representation of an artificial intelligence model is obtained. A reversed computation graph corresponding to a computation graph generated based on the intermediate representation is obtained. Nodes in the reversed computation graph represent functions related to the artificial intelligence model, and one or more directed edges in the reversed computation graph represent one or more dependencies between the functions. The reversed computation graph is partitioned into sequential partitions, such that the partitions are executed sequentially and functions corresponding to nodes in each partition are executed in parallel.
1. A method, comprising steps of:
obtaining an intermediate representation of an artificial intelligence model written in a source language, the intermediate representation comprising a structured text and being independent of the source language and a target language;
obtaining a reversed computation graph corresponding to a computation graph generated based on the intermediate representation, wherein nodes in the reversed computation graph represent functions related to the artificial intelligence model, and one or more directed edges in the reversed computation graph represent one or more dependencies between the functions, wherein a direction of each directed edge connecting two dependent nodes is reversed in the reversed computation graph; and
partitioning the reversed computation graph into sequential partition sets, such that the partition sets are executed sequentially and functions corresponding to nodes in each partition set are executed in parallel;
wherein partitioning the reversed computation graph into sequential partition sets further comprises:
calculating respective in-degrees for each node within the partition sets, wherein the in-degree of a given node represents a number of directed edges directed to the given node in the reversed computational graph;
identifying a first partition set comprising one or more nodes in the reversed computation graph having an in-degree of zero;
identifying a second partition set comprising one or more nodes in the reversed computation graph to which directed edges of the one or more nodes of the first partition set directly connect;
identifying at least a third partition set comprising one or more nodes in the reversed computation graph to which directed edges of the one or more nodes of the second partition set directly connect;
reordering the first, second and at least the third partition sets into a reversed order such that when a first function depends on a second function in accordance with the computation graph generated based on the intermediate representation, the second function is executed in the third partition set preceding the first function executed in the second partition set; and
sequentially executing respective functions of the at least the third partition set, the second partition set and the first partition set; and
wherein the steps are performed by one or more processing devices.
2. The method of claim 1 , further comprising determining one or more resources needed to execute functions within a given partition set prior to executing the functions in the given partition set.
3. The method of claim 2 , further comprising scheduling one or more functions in the given partition set to another partition set when at least a portion of the one or more resources needed to execute the functions within the given partition set are unavailable.
4. The method of claim 1 , wherein the first function and the second function are executed by a given resource of one or more resources needed to execute functions.
5. The method of claim 1 , wherein identifying the at least a third partition set comprises:
identifying at least a fourth partition set comprising one or more nodes to which directed edges of the one or more nodes of the third partition set directly connect.
6. The method of claim 1 , wherein at least one of the first partition set, the second partition set and the at least a third partition set comprises multiple nodes.
7. An apparatus, comprising:
a processor; and
a memory storing computer program instructions, the processor executing the computer program instructions in the memory to control the apparatus to:
obtain an intermediate representation of an artificial intelligence model written in a source language, the intermediate representation comprising a structured text and being independent of the source language and a target language;
obtain a reversed computation graph corresponding to a computation graph generated based on the intermediate representation, wherein nodes in the reversed computation graph represent functions related to the artificial intelligence model, and one or more directed edges in the reversed computation graph represent one or more dependencies between the functions, wherein a direction of each directed edge connecting two dependent nodes is reversed in the reversed computation graph; and
partition the reversed computation graph into sequential partition sets, such that the partition sets are executed sequentially and functions corresponding to nodes in each partition set are executed in parallel;
wherein partitioning the reversed computation graph into sequential partition sets further comprises:
calculating respective in-degrees for each node within the partition sets, wherein the in-degree of a given node represents a number of directed edges directed to the given node in the reversed computational graph;
identifying a first partition set comprising one or more nodes in the reversed computation graph having an in-degree of zero;
identifying a second partition set comprising one or more nodes in the reversed computation graph to which directed edges of the one or more nodes of the first partition set directly connect;
identifying at least a third partition set comprising one or more nodes in the reversed computation graph to which directed edges of the one or more nodes of the second partition set directly connect; and
reordering the first, second and at least the third partition sets into a reversed order such that when a first function depends on a second function in accordance with the computation graph generated based on the intermediate representation, the second function is executed in the third partition set preceding the first function executed in the second partition set; and
sequentially executing respective functions of the at least the third partition set, the second partition set and the first partition set.
8. The apparatus of claim 7 , wherein the processor executes the computer program instructions in the memory to further control the apparatus to determine one or more resources needed to execute functions within a given partition set prior to executing the functions in the given partition set.
9. The apparatus of claim 8 , wherein the processor executes the computer program instructions in the memory to further control the apparatus to schedule one or more functions in the given partition set to another partition set when at least a portion of the one or more resources needed to execute the functions within the given partition are unavailable.
10. The apparatus of claim 7 , wherein the first function and the second function are executed by a given resource of one or more resources needed to execute functions.
11. The apparatus of claim 7 , wherein identifying the at least a third partition set comprises:
identifying at least a fourth partition set comprising one or more nodes to which directed edges of the one or more nodes of the third partition set directly connect.
12. The apparatus of claim 7 , wherein at least one of the first partition set, the second partition set and the at least a third partition set comprises multiple nodes.
13. The apparatus of claim 7 , wherein the source language is written in different source languages.
14. A computer program product stored on a non-transitory computer-readable storage medium and comprising machine executable instructions, the machine executable instructions, when executed, causing a processing device to:
obtain an intermediate representation of an artificial intelligence model written in a source language, the intermediate representation comprising a structured text and being independent of the source language and a target language;
obtain a reversed computation graph corresponding to a computation graph generated based on the intermediate representation, wherein nodes in the reversed computation graph represent functions related to the artificial intelligence model, and one or more directed edges in the reversed computation graph represent one or more dependencies between the functions, wherein a direction of each directed edge connecting two dependent nodes is reversed in the reversed computation graph; and
partition the reversed computation graph into sequential partition sets, such that the partition sets are executed sequentially and functions corresponding to nodes in each partition set are executed in parallel;
wherein partitioning the reversed computation graph into sequential partition sets further comprises:
calculating respective in-degrees for each node within the partition sets, wherein the in-degree of a given node represents a number of directed edges directed to the given node in the reversed computational graph;
identifying a first partition set comprising one or more nodes in the reversed computation graph having an in-degree of zero;
identifying a second partition set comprising one or more nodes in the reversed computation graph to which directed edges of the one or more nodes of the first partition set directly connect;
identifying at least a third partition set comprising one or more nodes in the reversed computation graph to which directed edges of the one or more nodes of the second partition set directly connect;
reordering the first, second and at least the third partition sets into a reversed order such that when a first function depends on a second function in accordance with the computation graph generated based on the intermediate representation, the second function is executed in the third partition set preceding the first function executed in the second partition set; and
sequentially executing respective functions of the at least the third partition set, the second partition set and the first partition set.
15. The computer program product of claim 14 , further comprising determining one or more resources needed to execute functions within a given partition set prior to executing the functions in the given partition set.
16. The computer program product of claim 15 , further comprising scheduling one or more functions in the given partition set to another partition set when at least a portion of the one or more resources needed to execute the functions within the given partition set are unavailable.
17. The computer program product of claim 14 , wherein the first function and the second function are executed by a given resource of one or more resources needed to execute functions.
18. The computer program product of claim 14 , wherein identifying the at least a third partition set comprises:
identifying at least a fourth partition set comprising one or more nodes to which directed edges of the one or more nodes of the third partition set directly connect.
19. The computer program product of claim 14 , wherein at least one of the first partition set, the second partition set and the at least a third partition set comprises multiple nodes.
20. The computer program product of claim 14 , wherein the source language is written in different source languages.