用块尺度重用提升归一化效率，32倍减少计算量
MXNorm: Reusing MXFP block scales for efficient tensor normalisation
- 复用MXFP8的块尺度估算RMS，替代传统归一化
- 在Llama 3 125M~8B模型上精度损失极小，性能提升达2.4倍
- 适合追求低精度推理加速的开发者和硬件优化者
矩阵乘法性能长期是深度学习扩展的瓶颈，推动了低精度数值格式的发展。然而，矩阵乘法性能的提升远超规约和逐元素计算的改进，后者仍以高精度执行。本文提出MXNorm，一种可直接替换RMSNorm的方法，仅利用MXFP8量化过程中已计算的块尺度来估计RMS，使归一化所需的规约规模减少32倍。我们在125M、1B和8B参数的Llama 3模型预训练中验证该方法，与使用MXFP8矩阵乘法的RMSNorm基线相比，训练精度损失极小。此外，仅通过torch.compile即实现高达2.4倍的内核加速，在Llama 3 8B Transformer层中带来1.3%的提速，在NVFP4下实现2.6%的加速。
原文摘要 · Abstract (English)
Matrix multiplication performance has long been the major bottleneck to scaling deep learning workloads, which has stimulated the design of new accelerators that use increasingly low-precision number formats. However, improvements in matrix multiplication performance have far outstripped improvements in performance on reductions and elementwise computations, which are still being performed in higher precision. In this work, we propose MXNorm, a drop-in replacement for RMSNorm that estimates the RMS using only the block scales calculated as part of the MXFP8 cast and enables a 32x decrease in the size of reduction needed for normalization. We validate our approximation method on pre-training of Llama 3 models of 125M, 1B and 8B parameters, finding minimal loss of training accuracy compared to a baseline using RMSNorm with MXFP8 matmuls. We also show practical kernel speedups using only torch.compile of up to 2.4x for MXNorm over RMSNorm, corresponding to a 1.3% speedup in Llama 3 8B transformer layers in MXFP8 and a 2.6% speedup in NVFP4.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。