Graph partitioning and implementation of large models on tensor streaming processors
A graph partitioning compiler partitions an AI program or model for execution on multiple TSP modules configured for accelerating deep learning workloads.
1 . A method, comprising:
partitioning, by a computing system comprising one or more computing devices, a machine learning model into a plurality of partitions for execution on a plurality of deterministic-timing processor devices each configured to perform operations at a predetermined time determined at a compile time by a compiler; and
compiling, by the computing system, each of the plurality of partitions for execution on a corresponding deterministic-timing processor device of the plurality of deterministic-timing processor devices;
wherein partitioning the machine learning model comprises:
performing at least one horizontal slicing operation configured to assign a first tensor operation of the machine learning model to a first partition and a second tensor operation of the machine learning model to a second, different partition, while keeping each of the first tensor operation and second tensor operation intact; and
performing at least one vertical slicing operation configured to divide the first tensor operation for execution on two or more deterministic-timing processor devices of the plurality of deterministic-timing processor devices.
2 . The method of claim 1 , further comprising:
determining, by the computing system, resource requirement data indicative of one or more resource requirements for each respective tensor operation of a plurality of tensor operations of the machine learning model, wherein partitioning is performed based at least in part on the resource requirement data.
3 . The method of claim 2 , wherein the one or more resource requirements comprise a number of memory addresses to be allocated to inputs and outputs of the respective tensor operation.
4 . The method of claim 3 , wherein the resource requirement data comprises data indicative of a number of memory addresses that will be unavailable for allocation to the inputs and outputs of the respective tensor operation.
5 . The method of claim 4 , wherein the number of memory addresses that will be unavailable for allocation to the inputs and outputs comprises a number of memory addresses allocated to intermediate values to be used by a following operation subsequent to the respective tensor operation.
6 . The method of claim 1 , further comprising generating, by the computing system, a graph representation of the machine learning model, wherein partitioning the machine learning model comprises partitioning the graph representation into a plurality of subgraphs.
7 . The method of claim 6 , further comprising compiling, by the computing system, the machine learning model from a first language to a second language representation, wherein the graph representation is generated based on the second language representation.
8 . The method of claim 6 , wherein partitioning the graph representation into the plurality of subgraphs comprises:
determining, for each respective node or edge of a plurality of nodes or edges of the graph representation, one or more resource usage values indicative of an amount of resources to be used to execute the respective node or edge; and
partitioning the graph representation based at least in part on the one or more resource usage values such that a peak resource usage value for each of the plurality of subgraphs does not exceed a corresponding resource availability value.
9 . The method of claim 8 , further comprising selecting, by the computing system, the two or more deterministic-timing processor devices for execution of the first tensor operation based at least in part on a first peak resource usage value associated with the first tensor operation.
10 . The method of claim 1 , wherein the second partition comprises the second tensor operation and a third tensor operation configured to receive, as input, one or more outputs of the second tensor operation, and wherein compiling each of the plurality of partitions for execution comprises compiling one or more computer-executable instructions to cause a first deterministic-timing processor device of the plurality of deterministic-timing processor devices to execute at least a portion of the second tensor operation and at least a portion of the third tensor operation.
11 . A computing system comprising one or more processors and one or more non-transitory computer-readable media storing instructions that are executable by one or more processors to cause the computing system to:
partition a machine learning model into a plurality of partitions for execution on a plurality of deterministic-timing processor devices each configured to perform operations at a predetermined time determined at a compile time by a compiler; and
compile each of the plurality of partitions for execution on a corresponding deterministic-timing processor device of the plurality of deterministic-timing processor devices;
wherein partitioning the machine learning model comprises:
performing at least one horizontal slicing operation configured to assign a first tensor operation of the machine learning model to a first partition and a second tensor operation of the machine learning model to a second, different partition, while keeping each of the first tensor operation and second tensor operation intact; and
performing at least one vertical slicing operation configured to divide the first tensor operation for execution on two or more deterministic-timing processor devices of the plurality of deterministic-timing processor devices.
12 . The computing system of claim 11 , wherein the instructions further cause the computing system to:
determine resource requirement data indicative of one or more resource requirements for each respective tensor operation of a plurality of tensor operations of the machine learning model, wherein partitioning is performed based at least in part on the resource requirement data.
13 . The computing system of claim 12 , wherein the one or more resource requirements comprise a number of memory addresses to be allocated to inputs and outputs of the respective tensor operation.
14 . The computing system of claim 13 , wherein the resource requirement data comprises data indicative of a number of memory addresses that will be unavailable for allocation to the inputs and outputs of the respective tensor operation.
15 . The computing system of claim 14 , wherein the number of memory addresses that will be unavailable for allocation to the inputs and outputs comprises a number of memory addresses allocated to intermediate values to be used by a following operation subsequent to the respective tensor operation.
16 . The computing system of claim 11 , wherein the instructions further cause the computing system to generate a graph representation of the machine learning model, wherein partitioning the machine learning model comprises partitioning the graph representation into a plurality of subgraphs.
17 . The computing system of claim 16 , wherein the instructions further cause the computing system to compile the machine learning model from a first language to a second language representation, wherein the graph representation is generated based on the second language representation.
18 . The computing system of claim 16 , wherein partitioning the graph representation into the plurality of subgraphs comprises:
determining, for each respective node or edge of a plurality of nodes or edges of the graph representation, one or more resource usage values indicative of an amount of resources to be used to execute the respective node or edge; and
partitioning the graph representation based at least in part on the one or more resource usage values such that a peak resource usage value for each of the plurality of subgraphs does not exceed a corresponding resource availability value.
19 . The computing system of claim 18 , wherein the instructions further cause the computing system to select the two or more deterministic-timing processor devices for execution of the first tensor operation based at least in part on a first peak resource usage value associated with the first tensor operation.
20 . One or more example non-transitory computer-readable media storing instructions that are executable by one or more processors to cause a computing system to:
partition a machine learning model into a plurality of partitions for execution on a plurality of deterministic-timing processor devices each configured to perform operations at a predetermined time determined at a compile time by a compiler; and
compile each of the plurality of partitions for execution on a corresponding deterministic-timing processor device of the plurality of deterministic-timing processor devices;
wherein partitioning the machine learning model comprises:
performing at least one horizontal slicing operation configured to assign a first tensor operation of the machine learning model to a first partition and a second tensor operation of the machine learning model to a second, different partition, while keeping each of the first tensor operation and second tensor operation intact; and
performing at least one vertical slicing operation configured to divide the first tensor operation for execution on two or more deterministic-timing processor devices of the plurality of deterministic-timing processor devices.