High performance layer normalization for large models

Abstract

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.