Shared memory parallel execution of neural networks

1 . A method for performing a collective compute operation in a system having M number of processing nodes, wherein the M number of processing nodes are implemented in (M/N) number of acceleration engines each having N number of processing nodes, and each processing node includes a high bandwidth memory (HBM), a state buffer, and a tensor processing block, the method comprising:

partitioning a set of weight parameters into M number of sharded portions of weight parameters to provide each of the M number of processing nodes with a sharded portion of weight parameters;

partitioning a batch of input data into M number of minibatches of input data to provide each of the M number of processing nodes with a minibatch of input data;

storing each sharded portion of weight parameters and each minibatch of input data in the HBM of their corresponding processing node in the system; and

for each of the processing nodes in each acceleration engine:

allocating an amount of space for (M/N) number of sharded portions of weight parameters in the HBM;

obtaining a sharded portion of weight parameters from a corresponding process node of each of the acceleration engines to gather (M/N) number of sharded portions of weight parameters in the HBM of the processing node;

for each of the (M/N) number of sharded portions of weight parameters stored in the HBM of the processing node, multicasting the sharded portion of weight parameters to the state buffer of each processing node within the acceleration engine to provide each state buffer with the set of weight parameters;

loading the minibatch of input data from the HBM into the state buffer of the processing node;

loading the set of weight parameters and the minibatch of input data from the state buffer into the tensor processing block of the processing node; and

performing computations on the set of weight parameters and the minibatch of input data in the tensor processing block.

2 . The method of claim 1 , wherein multicasting the sharded portion of weight parameters within each acceleration engine is performed by a set of one or more direct memory access (DMA) engines in the acceleration engine.

3 . The method of claim 1 , wherein the collective compute operation is managed by a runtime driver of the system based on the N number of processing nodes in each acceleration engine and the (M/N) number of acceleration engines in the system.

4 . The method of claim 1 , wherein the set of weight parameters corresponds to a layer of a neural network model.

5 . A method comprising:

performing in each of N number of processing nodes of an integrated circuit device that is part of a system of M number of processing nodes implemented in (M/N) number of integrated circuit devices:

allocating an amount of space in a shared memory of a processing node for a collective compute operation operating on M number of sharded portions of parameters, wherein the amount of space allocated in the shared memory corresponds to (M/N) number of sharded portions of parameters;

storing a sharded portion of parameters associated with the processing node in the allocated space of the shared memory of the processing node;

obtaining (M/N)−1 number of sharded portions of parameters from other integrated circuit devices to store (M/N) number of sharded portions of parameters in the allocated space of the shared memory of the processing node; and

for each of the (M/N) number of sharded portions of parameters stored in the shared memory of the processing node, multicasting the sharded portion of parameters to a state buffer of each of the N number of processing nodes in the integrated circuit device to provide each state buffer with the M number of sharded portions of parameters for the collective compute operation.

6 . The method of claim 5 , wherein multicasting the sharded portion of parameters includes reading the sharded portion of parameters once from the shared memory, and writing the sharded portion of parameters to each state buffer in the integrated circuit device.

7 . The method of claim 6 , wherein multicasting the sharded portion of parameters is performed by a set of one or more direct memory access (DMA) engines of the integrated circuit device.

8 . The method of claim 5 , wherein the M number of sharded portions of parameters correspond to weight values of a layer of a neural network model.

9 . The method of claim 8 , further comprising:

loading a minibatch of input data into the state buffer of the processing node;

loading the weight values and the minibatch of input data from the state buffer into a tensor processing block of the processing node; and

performing computations on the weight values and the minibatch of input data.

10 . The method of claim 9 , wherein the minibatch of input data is one of a plurality of minibatches of input data partitioned from a batch of input data, and each processing node of the integrated circuit device processes one of the minibatches of input data in parallel with the other processing nodes of the integrated circuit device.

11 . The method of claim 5 , wherein memory allocation and data transfers for the collective compute operation are managed by a runtime driver of the integrated circuit device.

12 . The method of claim 11 , wherein the runtime driver is configured with the values M and N to enable the runtime driver to allocate space in the shared memory for the collective compute operation.

13 . The method of claim 5 , wherein the collective compute operation is an all-gather operation or an all-reduce operation.

14 . A non-transitory computer readable medium having stored therein instructions that, when executed by one or more processors, cause the one or more processors to execute a runtime driver, the runtime driver performing operations in each of N number of processing nodes of an integrated circuit device that is part of a system of M number of processing nodes implemented in (M/N) number of integrated circuit devices, the operations including:

allocating an amount of space in a shared memory of a processing node for a collective compute operation operating on M number of sharded portions of parameters, wherein the amount of space allocated in each shared memory corresponds to (M/N) number of sharded portions of parameters;

storing a sharded portion of parameters associated with the processing node in the allocated space of the shared memory of the processing node;

obtaining (M/N)−1 number of sharded portions of parameters from other integrated circuit devices to store (M/N) number of sharded portions of parameters in the allocated space of the shared memory of the processing node; and

for each of the (M/N) number of sharded portions of parameters stored in the shared memory of the processing node, multicasting the sharded portion of parameters to a state buffer of each of the N number of processing nodes in the integrated circuit device to provide each state buffer with the M number of sharded portions of parameters for the collective compute operation.

15 . The non-transitory computer readable medium of claim 14 , wherein multicasting the sharded portion of parameters includes reading the sharded portion of parameters once from the shared memory, and writing the sharded portion of parameters to each of N number of state buffers in the integrated circuit device.

16 . The non-transitory computer readable medium of claim 14 , wherein the operations further include:

loading a minibatch of input data into the state buffer of the processing node;

loading weight values and the minibatch of input data into a tensor processing block of the processing node; and

performing computations on the weight values and the minibatch of input data.

17 . The non-transitory computer readable medium of claim 14 , wherein the collective compute operation is an all-gather operation or an all-reduce operation.

18 . An integrated circuit device comprising:

N number of processing nodes, wherein each processing node includes:

a shared memory accessible by each of the processing nodes;

a set of direct memory access (DMA) engines;

a state buffer; and

a tensor processing block,

wherein each of the processing nodes is operable to:

allocate an amount of space in the shared memory of the processing node for a collective compute operation operating on M number of sharded portions of parameters, wherein M is greater than N, and wherein the amount of space allocated in the shared memory corresponds to (M/N) number of sharded portions of parameters;

store a sharded portion of parameters associated with the processing node in the allocated space of the shared memory of the processing node;

obtain (M/N)−1 number of sharded portions of parameters from other integrated circuit devices to store (M/N) number of sharded portions of parameters in the allocated space of the shared memory of the processing node; and

for each of the (M/N) number of sharded portions of parameters stored in the shared memory of the processing node, perform DMA accesses to multicast the sharded portion of parameters to a state buffer of each of the N number of processing nodes to provide each state buffer with the M number of sharded portions of parameters for the collective compute operation.

19 . The integrated circuit device of claim 18 ,

wherein each of the processing nodes is further operable to load the M number of sharded portions of parameters and a minibatch of input data into the tensor processing block of the processing node; and

wherein the integrated circuit device is operable to perform computations on N number of minibatches of input data with the M number of sharded portions of parameters in parallel.

20 . The integrated circuit device of claim 18 , wherein the collective compute operation is an all-gather operation or an all-reduce operation.