Tensor operations in AI models
A method of performing computations for artificial intelligence models may include obtaining an input tensor based on an input to an artificial intelligence model and loading the input tensor into multiple processing devices. The input tensor may be split into multiple input tensor tiles that are distributed among the processing devices such that each of the processing devices does not include an entirety of the input tensor. The method may also include performing multiple tensor operations according to the artificial intelligence model to generate multiple intermediate tensors and an output tensor, one or more of the tensor operations performed using the input tensor.
1. A method of performing computations for artificial intelligence models, the method comprising:
obtaining an input tensor based on an input to an artificial intelligence model;
loading the input tensor into a plurality of processing devices, the input tensor split into a plurality of input tensor tiles that are distributed among the plurality of processing devices such that each of the plurality of processing devices does not include an entirety of the input tensor; and
performing a plurality of tensor operations according to the artificial intelligence model to generate a plurality of intermediate tensors and an output tensor, one or more of the plurality of tensor operations performed using the input tensor and each of the plurality of intermediate tensors being split into tensor tiles distributed among the plurality of processing devices such that each of the plurality of processing devices does not include an entirety of the plurality of intermediate tensors during any of the plurality of tensor operations.
2. The method of claim 1 , further comprising iteratively performing the steps of obtaining, loading, and performing, wherein the input tensor for a subsequent iteration is the output tensor from a previous iteration.
3. The method of claim 1 , wherein the distribution of the tensor tiles of the one or more of the plurality of intermediate tensors among the plurality of processing devices is different than the distribution of the plurality of input tensor tiles among the plurality of processing devices.
4. The method of claim 1 , wherein the output tensor is split into tensor tiles distributed among the plurality of processing devices such that each of the plurality of processing devices does not include an entirety of the output tensor.
5. The method of claim 1 , wherein the plurality of input tensor tiles are distributed among the plurality of processing devices such that half of the plurality of processing devices each include half of the plurality of input tensor tiles.
6. The method of claim 1 , wherein the artificial intelligence model implements a transformer architecture.
7. The method of claim 6 , wherein the plurality of intermediate tensors include a self-attention tensor, and the self-attention tensor is split into a plurality of intermediate tensor tiles and distributed among the plurality of processing devices such that each of the plurality of processing devices includes a different sub-set of the plurality of intermediate tensor tiles and no duplication of the plurality of intermediate tensor tiles exists among the plurality of processing devices.
8. The method of claim 7 , wherein the plurality of intermediate tensors includes a projection tensor that is split into a plurality of projection tensor tiles and distributed among the plurality of processing devices in the same manner as the input tensor is distributed among the plurality of processing devices.
9. The method of claim 8 , wherein the distribution of the input tensor among the plurality of processing devices is different than the distribution of the self-attention tensor among the plurality of processing devices.
10. A system comprising:
one or more memory devices configured to store a plurality of tensors; and
a plurality of processing devices coupled to the one or more memory devices and configured to perform tensor operations on the plurality of tensors, the system configured to execute instructions to cause the system to perform operations, the operations comprising:
obtaining an input tensor based on an input to an artificial intelligence model;
loading the input tensor into the plurality of processing devices, the input tensor split into a plurality of input tensor tiles that are distributed among the plurality of processing devices such that each of the plurality of processing devices does not include an entirety of the input tensor; and
performing, using the plurality of processing devices, a plurality of tensor operations according to the artificial intelligence model to generate a plurality of intermediate tensors and an output tensor, one or more of the plurality of tensor operations performed using the input tensor and each of the plurality of intermediate tensors being split into tensor tiles distributed among the plurality of processing devices such that each of the plurality of processing devices does not include an entirety of the plurality of intermediate tensors during any of the plurality of tensor operations.
11. The system of claim 10 , wherein the operations further comprise iteratively performing the operations with the input tensor for a subsequent iteration being the output tensor from a previous iteration.
12. The system of claim 10 , wherein the distribution of the tensor tiles of the one or more of the plurality of intermediate tensors among the plurality of processing devices is different than the distribution of the plurality of input tensor tiles among the plurality of processing devices.
13. The system of claim 10 , wherein the output tensor is split into tensor tiles distributed among the plurality of processing devices such that each of the plurality of processing devices does not include an entirety of the output tensor.
14. The system of claim 10 , wherein the plurality of input tensor tiles are distributed among the plurality of processing devices such that half of the plurality of processing devices each include half of the plurality of input tensor tiles.
15. The system of claim 10 , wherein the artificial intelligence model implements a transformer architecture.
16. The system of claim 15 , wherein the plurality of intermediate tensors include a self-attention tensor, and the self-attention tensor is split into a plurality of intermediate tensor tiles and distributed among the plurality of processing devices such that each of the plurality of processing devices includes a different sub-set of the plurality of intermediate tensor tiles and no duplication of the plurality of intermediate tensor tiles exists among the plurality of processing devices.
17. The system of claim 16 , wherein the plurality of intermediate tensors includes a projection tensor that is split into a plurality of projection tensor tiles and distributed among the plurality of processing devices in the same manner as the input tensor is distributed among the plurality of processing devices.
18. The system of claim 17 , wherein the distribution of the input tensor among the plurality of processing devices is different than the distribution of the self-attention tensor among the plurality of processing devices.
19. A method of performing computations for artificial intelligence models, the method comprising:
obtaining an input tensor based on an input to an artificial intelligence model;
loading the input tensor into a plurality of separate processing devices, the input tensor split into a plurality of input tensor tiles that are distributed among the plurality of separate processing devices such that each of the plurality of separate processing devices includes a portion of the plurality of input tensor tiles but does not include an entirety of the input tensor; and
performing a plurality of tensor operations according to the artificial intelligence model to generate a plurality of intermediate tensors and an output tensor, one or more of the plurality of tensor operations performed using the input tensor.
20. The method of claim 19 , wherein at least one of the plurality of intermediate tensors is split into a plurality of intermediate tensor tiles and distributed among the plurality of separate processing devices and the distribution of the input tensor among the plurality of separate processing devices is different than the distribution of the at least one of the plurality of intermediate tensors among the plurality of separate processing devices.
21. The method of claim 20 , wherein the artificial intelligence model implements a transformer architecture and the at least one of the plurality of intermediate tensors includes a self-attention tensor.