MXNorm：复用 MXFP 块尺度实现高效张量归一化
机器学习
2026-03-16 v1 人工智能
神经与进化计算
摘要
矩阵乘法性能长期以来是扩缩深度学习工作负载的主要瓶颈，这刺激了设计新的加速器，使用越来越低的精度数字格式。然而，矩阵乘法性能的提高远远快于对约束和逐元素计算的性能提升，这些计算仍以更高的精度进行。本文提出 MXNorm，作为 RMSNorm 的即插即用替换方案，通过仅使用 MXFP8 转换中计算的块尺度来估计 RMS，从而将归一化所需的约束大小降低32倍。我们在 125M、1B 和 8B 参数的 Llama 3 模型预训练中验证了该近似方法，发现与使用 RMSNorm 并进行 MXFP8 矩阵乘法的基线相比，训练精度损失最小。我们还展示了仅使用 torch.compile 对 MXNorm 相对于 RMSNorm 实现了最高达 2.4 倍的实际内核加速，对应于在 MXFP8 中的 Llama 3 8B Transformer 层实现 1.3% 的加速，以及在 NVFP4 中实现 2.6% 的加速。
引用
@article{arxiv.2603.13180,
title = {MXNorm: Reusing MXFP block scales for efficient tensor normalisation},
author = {Callum McLean and Luke Y. Prince and Alexandre Payot and Paul Balança and Carlo Luschi},
journal= {arXiv preprint arXiv:2603.13180},
year = {2026}
}
备注
Preprint, Under Review. 15 pages, 12 figures