Sparse machine learning acceleration
To reduce the storage size of weight tensors and speed up loading of weight tensors from system memory, a compression technique can be employed to remove zero values from a weight tensor before storing the weight tensor in system memory. A sparsity threshold can be enforced to achieve a compression ratio target by forcing small weight values to zero during training. When the weight tensor is loaded from system memory, a direct memory access (DMA) engine with an in-line decompression unit can decompress the weight tensor on-the-fly. By performing the decompression in the DMA engine, expansion of the weight values back to the original weight tensor size can be carried out in parallel while other neural network computations are being performed by the processing unit.
1. A computer-implemented method comprising:
performing a training operation on a neural network model to generate a first weight tensor;
compressing, by a direct memory access (DMA) engine, the first weight tensor to meet a first sparsity criteria by:
replacing a subset of non-zero weight values in the first weight tensor with zero values to generate a clipped weight tensor having a first sparsity corresponding to the first sparsity criteria for the neural network model; and
generating a first compressed weight tensor by removing zero values from the clipped weight tensor;
performing a subsequent stage of the training operation using the first compressed weight tensor to generate a second weight tensor;
compressing, by the DMA engine, the second weight tensor into a second compressed weight tensor having a second sparsity corresponding to a second sparsity criteria, the second sparsity having a higher density than the first sparsity;
storing the second compressed weight tensor with the compression header having information on locations of zero values in the second weight tensor in system memory;
retrieving the second compressed weight tensor with the compression header from the system memory by the DMA engine according to a memory descriptor placed in a descriptor queue to transfer a data set, the memory descriptor containing an address and an operation type that indicates the data set includes compressed data;
in response to the operation type of the memory descriptor indicating the data set includes compressed data, using the compression header to decompress the second compressed weight tensor into a regenerated weight tensor in the DMA engine;
loading the regenerated weight tensor into a systolic array of a neural network accelerator using the DMA engine; and
performing neural network computations in the systolic array of the neural network accelerator using the regenerated weight tensor for further training or inference,
wherein the systolic array of the neural network accelerator is performing other neural network computations including matrix multiplication operations while the DMA engine is performing compression and decompression operations.
2. The computer-implemented method of claim 1 , wherein replacing the subset of non-zero weight values in the first weight tensor includes:
comparing a non-zero weight value with a threshold value;
determining that a magnitude of the non-zero weight value is below the threshold value; and
replacing the non-zero weight value with a zero value.
3. The computer-implemented method of claim 1 , further comprising quantizing weight values to generate the first weight tensor.
4. A direct memory access (DMA) engine comprising:
a DMA controller configured to perform weight tensor compression operations according to a sparsity criteria during a training process, wherein the sparsity criteria is adjusted as training progresses;
a descriptor queue configured to store memory descriptors for the DMA engine to process, each memory descriptor providing an address and an operation type for a data transfer;
an input buffer configured to store a data set read from system memory according to a memory descriptor retrieved from the descriptor queue;
an output buffer configured to output a weight tensor to a neural network accelerator; and
a decompression circuit block coupled between the input buffer and the output buffer, wherein the decompression circuit block is configured to:
determine, based on the operation type of the memory descriptor retrieved from the descriptor queue, that the data set stored in the input buffer includes a compressed weight tensor corresponding to the weight tensor for a neural network model, wherein the compressed weight tensor contains a set of non-zero weight values;
determine a size of the weight tensor;
decompress the compressed weight tensor by inserting zero weight values into the set of non-zero weight values; and
provide the weight tensor to the output buffer,
wherein a systolic array of the neural network accelerator is configured to perform other neural network computations including matrix multiplication operations while the DMA engine is performing compression and decompression operations.
5. The direct memory access engine of claim 4 , wherein the data set includes a compression header that is used by the decompression circuit block to decompress the compressed weight tensor.
6. The direct memory access engine of claim 5 , wherein the compression header includes index information indicating locations in the weight tensor to insert the zero weight values.
7. The direct memory access engine of claim 6 , wherein the index information is implemented as a binary bit map containing a number of bits being equal to a number of elements in the weight tensor.
8. The direct memory access engine of claim 5 , wherein the compression header includes a length value indicating a size of the weight tensor.
9. The direct memory access engine of claim 4 , wherein the DMA controller is configured to adjust the sparsity criteria to increase sparsity as training progresses.
10. The direct memory access engine of claim 4 , wherein the weight values in the weight tensor are quantized weight values.
11. The direct memory access engine of claim 4 , wherein the weight tensor includes zero values that replaced non-zero values having a magnitude that is below a threshold value.
12. The direct memory access engine of claim 4 , wherein the sparsity criteria is expressed as a maximum number of non-zero values allowed for the weight tensor.
13. A method comprising:
performing, by a direct memory access (DMA) engine, weight tensor compression operations according to a sparsity criteria during a training process of a neural network model, wherein the sparsity criteria is adjusted to increase sparsity as training progresses;
storing, in system memory, a compressed weight tensor corresponding to a weight tensor generated by the training process;
receiving, by the DMA engine, a memory descriptor placed in a descriptor queue to transfer data from the system memory to a neural network accelerator, wherein the memory descriptor contains an address and an operation type that indicates the data set includes the compressed weight tensor;
reading, by the DMA engine, a data set from the system memory;
determining, by the DMA engine based on the operation type of the memory descriptor, that the data set includes the compressed weight tensor corresponding to the weight tensor for the neural network model, the compressed weight tensor including a set of non-zero weight values;
determining, by the DMA engine, a size of the weight tensor;
inserting, by the DMA engine, zero weight values in the set of non-zero weight values to decompress the compressed weight tensor; and
writing, by the DMA engine, the weight tensor to a buffer in the neural network accelerator,
wherein the weight tensor compression operations and the decompression of the compressed weight tensor are performed by the DMA engine while a systolic array of the neural network accelerator is performing other neural network computations including matrix multiplication operations.
14. The method of claim 13 , wherein the data set includes a compression header containing index information indicating locations of the zero weight values in the weight tensor.
15. The method of claim 14 , wherein the compression header includes a length field indicating the size of the weight tensor.
16. The method of claim 13 , wherein the compressed weight tensor is compressed by replacing a subset of non-zero weight values in the weight tensor with zero values.
17. The method of claim 16 , wherein the number of non-zero weight values in the weight tensor being replaced by zero values is a predetermined number of lowest magnitude non-zero weight values.
18. The method of claim 16 , wherein the non-zero weight values being replaced by zero values include all non-zero weight values in the weight tensor having a magnitude that is below a threshold value.
19. The method of claim 18 , further comprising adjusting the threshold value to meet the sparsity criteria of the weight tensor.