High performance layer normalization for large models
As general matrix multiply (GEMM) bottlenecks are ameliorated by tensor parallelism that is distributed to several processors, layer normalization (LN) surfaces as a latent bottleneck as it is not amenable to distribution. LN performance is linear to embedding size, which is extremely large in some AI models. Moreover, aggressive tiling prevents the use of internal pipelining. The disclosed implementation addresses this issue, composing LN from simpler operations and this composition is amenable to pipelining, facilitating efficient implementation of large AI models (e.g., GPTs). In both forward and backward propagation, the pipeline is stretched longer with improved balance across stages. This strategy improves throughput for larger batch-sizes as the workload benefits from pipelining operations for better performance. Furthermore, avoiding stochastic rounding further improves performance. In addition, LayerNorm checkpoints facilitate efficient computation of gradients during backward propagation.
1. A computer-implemented method for normalizing data in a reconfigurable dataflow processor, the method comprising:
partitioning normalization operations into multiple compute stages and intervening buffering stages in a reconfigurable dataflow processor;
configuring the multiple compute stages and the intervening buffering stages; and
processing data (X) using the multiple compute stages and the intervening buffering stages.
2. The computer-implemented method of claim 1 , wherein:
the compute stages comprise a centering stage that computes X−μ.
3. The computer-implemented method of claim 1 , wherein:
the compute stages comprise a normalization stage that computes (X−μ)/σ.
4. The computer-implemented method of claim 1 , wherein:
the compute stages comprise a shift-add stage that computes γ*((X−μ)/σ)+β.
5. The computer-implemented method of claim 1 , further including:
conducting layer normalization computations in a forward-propagation pass.
6. The computer-implemented method of claim 5 , further including:
saving selected computed values (xHat) from the layer normalization computations.
7. The computer-implemented method of claim 6 , further including:
reusing the selected computed values in a backward-propagation pass.
8. The computer-implemented method of claim 7 , wherein:
the selected computed values are computed according to the equation xHat=(X−μ)/σ.
9. A non-transitory computer-readable storage medium storing computer program instructions that, when executed on a processor, perform operations comprising:
partitioning normalization operations into multiple compute stages and intervening buffering stages in a reconfigurable dataflow processor;
configuring the multiple compute stages and the intervening buffering stages; and
processing data (X) using the multiple compute stages and the intervening buffering stages.
10. The non-transitory computer-readable storage medium of claim 9 , wherein:
the compute stages comprise a centering stage that computes X−μ.
11. The non-transitory computer-readable storage medium of claim 9 , wherein:
the compute stages comprise a normalization stage that computes (X−μ)/σ.
12. The non-transitory computer-readable storage medium of claim 9 , wherein:
the compute stages comprise a shift-add stage that computes γ*((X−μ)/σ)+β.
13. The non-transitory computer-readable storage medium of claim 9 , further comprising:
conducting layer normalization computations in a forward-propagation pass.
14. The non-transitory computer-readable storage medium of claim 13 , further comprising:
saving selected computed values (xHat) from the layer normalization computations.
15. The non-transitory computer-readable storage medium of claim 14 , further comprising:
reusing the selected computed values in a backward-propagation pass.
16. The non-transitory computer-readable storage medium of claim 15 , wherein:
the selected computed values are computed according to the equation xHat=(X−μ)/σ.
17. A system comprising one or more processors coupled to a memory device, the memory device to store computer program instructions that are executable by the one or more processors to perform operations comprising:
partitioning normalization operations into multiple compute stages and intervening buffering stages in a reconfigurable dataflow processor;
configuring the multiple compute stages and the intervening buffering stages; and
processing data (X) using the multiple compute stages and the intervening buffering stages.
18. The system of claim 17 , wherein:
the compute stages comprise a centering stage that computes X−μ.
19. The system of claim 17 , wherein:
the compute stages comprise a normalization stage that computes (X−μ)/σ.
20. The system of claim 17 , wherein:
the compute stages comprise a shift-add stage that computes γ*((X−μ)/σ)+β.