MXNorm: Reusing MXFP block scales for efficient tensor normalisation

In recent years, deep learning has made significant progress in natural language processing, computer vision, and scientific fields such as molecular biology. This progress has been enabled by leaps in AI accelerator capabilities, particularly in the acceleration of low-precision matrix multiplications. Over the past eight years, GPU acceleration of low-precision matrix multiplications has improved by 80 times, allowing researchers and practitioners to scale pre-training of transformer-style neural networks and learn from vast pools of unlabelled data. However, as matrix multiplication becomes less of a bottleneck to throughput, other components of model architectures emerge as new bottlenecks. Indeed, other aspects of AI accelerators have not kept pace with improvements in matrix multiplication throughput. For example, elementwise operations and reductions are limited by memory bandwidth and CUDA core throughput in GPUs; however, these have only improved by factors of 8.9x and 5.1x, respectively, over the past eight years. Furthermore, this gap is set to widen with upcoming GPU architecture releases.

While the improvement in matrix multiplication performance has addressed the bottleneck in scaling deep learning workloads, the performance of reductions and elementwise computations has lagged, becoming a new bottleneck. These operations are limited by memory bandwidth and CUDA core throughput, which have only improved by factors of 8.9x and 5.1x, respectively, over the past eight years. Furthermore, this gap is set to widen with upcoming GPU architecture releases. In some cases, these operations can be hidden by overlapping them with matrix multiplications, but others require too much memory to pipeline in practice. Therefore, the community needs to consider new building blocks that contribute a smaller overhead.

The core innovation of MXNorm lies in its integration of normalization with the MX quantization process, reusing MXFP8 block scales to estimate RMS. This method significantly reduces redundant computation overhead, enhancing training efficiency and speed. Specifically:

• �� MXNorm estimates RMS by reusing block scales calculated during MXFP8 quantization, achieving efficient normalization and reducing the reduction size requirement by 32 times.

• �� MXNorm serves as a replacement for RMSNorm and was validated on pre-training Llama 3 models, showing minimal training accuracy loss compared to a baseline using RMSNorm with MXFP8 matrix multiplications.

• �� Practical kernel speedups using torch.compile demonstrated MXNorm to be up to 2.4x faster than RMSNorm, translating to a 1.3% speedup in Llama 3 8B transformer layers in MXFP8 and a 2.6% speedup in NVFP4.

The implementation of MXNorm involves the following steps:

• �� First, MXNorm estimates RMS by reusing block scales calculated during MXFP8 quantization, reducing the reduction size requirement by 32 times.

• �� Then, MXNorm was validated on pre-training Llama 3 models with 125M, 1B, and 8B parameters. Experimental results showed minimal training accuracy loss compared to a baseline using RMSNorm with MXFP8 matrix multiplications.

• �� Finally, practical kernel speedups using torch.compile demonstrated MXNorm to be up to 2.4x faster than RMSNorm, translating to a 1.3% speedup in Llama 3 8B transformer layers in MXFP8 and a 2.6% speedup in NVFP4.

The experimental design involved validating the performance of MXNorm on pre-training Llama 3 models with 125M, 1B, and 8B parameters. The baseline used was RMSNorm with MXFP8 matrix multiplications. The experiments utilized torch.compile to implement kernel acceleration for MXNorm and were tested on actual hardware. Additionally, a learning rate sensitivity test was conducted to evaluate the stability and performance of MXNorm under different learning rates. The results showed minimal training accuracy loss and significant speed improvements in kernel acceleration.

The experimental results demonstrated that MXNorm performed excellently in pre-training Llama 3 models. On 125M, 1B, and 8B parameter models, MXNorm showed minimal training accuracy loss compared to RMSNorm. Additionally, using torch.compile for kernel acceleration, MXNorm was up to 2.4x faster than RMSNorm, translating to a 1.3% speedup in Llama 3 8B transformer layers in MXFP8 and a 2.6% speedup in NVFP4. On the 8B parameter model, MXNorm(p=2) achieved a training loss comparable to RMSNorm, with losses of 2.126 and 2.132 respectively, while MXNorm(p=1) had a worse final loss of 2.175.

MXNorm's application scenarios include pre-training large language models, particularly in situations requiring efficient normalization and reduced computational overhead. MXNorm allows for significant speed improvements without substantial accuracy loss. Additionally, MXNorm demonstrated excellent performance in practical kernel acceleration, making it suitable for high-performance computing scenarios such as natural language processing and computer vision.

While MXNorm performed excellently in experiments, it also has some limitations. On larger models (e.g., 8B parameters), MXNorm showed slightly higher training loss than RMSNorm, indicating potential accuracy loss in certain scenarios. Additionally, MXNorm's stability is poorer at higher learning rates, with a tendency to exhibit loss spikes. These issues need to be addressed in future research.