MXNorm: Reusing MXFP block scales for efficient tensor normalisation

本文提出了 MXNorm，一种通过复用 MXFP8 矩阵乘法中已计算的块缩放因子来估计均方根（RMS）的归一化方法，该方法在 Llama 3 模型预训练中保持了训练精度，同时显著减少了归一化所需的约简操作并提升了内核执行速度。

要点

这篇论文介绍了一种名为 MXNorm 的新方法，旨在解决人工智能（AI）大模型训练中的一个“速度瓶颈”问题。

为了让你轻松理解，我们可以把训练一个大型 AI 模型（比如 Llama 3）想象成指挥一支庞大的交响乐团演奏一首复杂的交响曲。

1. 背景：乐团里的“速度差”

在过去几年里，AI 芯片（就像乐团的指挥台）变得非常强大，特别是处理矩阵乘法（也就是乐团里成千上万个乐手同时演奏主旋律）的速度提升了 80 倍！这就像给乐团装上了火箭推进器。

但是，乐团里还有一些辅助工作，比如“检查音准”和“调整音量平衡”（在 AI 里叫归一化，具体是 RMSNorm）。这些辅助工作原本是为了确保音乐（数据）不会太吵或太轻，保持整体和谐。

问题出在哪里？
虽然主旋律演奏得飞快，但“检查音准”和“调整音量”的速度却提升得很慢（只提升了约 5-9 倍）。这就好比：

乐手们已经能像闪电一样演奏，但指挥家还在慢吞吞地拿着尺子一个个去量音高。结果，整个乐团的节奏被这些慢吞吞的“检查”环节拖累了，导致整体效率上不去。

2. 核心创新：MXNorm（“复用”的智慧）

为了解决这个问题，Graphcore 的研究团队提出了 MXNorm。

原来的做法（RMSNorm）：
在把乐谱（数据）交给乐手之前，指挥家必须先重新计算一遍每个乐手的音量平均值（计算 RMS），然后再调整。这就像每次演奏前，都要重新拿尺子量一遍，非常耗时。

MXNorm 的做法（“偷梁换柱”）：
现在的 AI 芯片为了跑得更快，会把乐谱压缩成一种更小的格式（叫 MXFP 格式）。在压缩乐谱的过程中，芯片已经顺便计算了每个小乐队的“最大音量”（Block Scales，块尺度）。

MXNorm 的聪明之处在于：它不再重新计算音量平均值，而是直接“借用”压缩过程中已经算好的“最大音量”数据，来估算平均值。

生活中的类比：

• 传统方法：你要统计一个班级学生的平均身高。你必须把全班 50 个人的身高都量一遍，加起来除以 50。
• MXNorm 方法：你发现班里已经有人为了买校服，把每 10 个人的“最高身高”都记下来了。你不需要重新量所有人，直接把这 5 个“最高身高”拿来算个平均，就能非常准确地估算出全班的平均身高。
  • 结果：你省去了 90% 的测量工作（减少了 32 倍的数据处理量），而且算出来的结果几乎一样准。

3. 实验结果：快且稳

研究人员在 Llama 3 模型（从 1.25 亿参数到 80 亿参数）上测试了这种方法：

1. 速度提升：

   • 在单个计算步骤上，MXNorm 比传统方法快了 2.4 倍。
   • 在整个大模型推理中，整体速度提升了 1.3% 到 2.6%。
   • 比喻：这就像让交响乐团在保持音准完美的前提下，把整首曲子的演奏时间缩短了，而且不需要更换任何乐器（硬件），只需要换一种指挥技巧（软件算法）。

2. 稳定性：

   • 他们发现，如果用简单的“算术平均”（p=1）来估算，大模型训练容易“翻车”（出现数据爆炸，就像乐团突然有人尖叫，导致整首曲子跑调）。
   • 但如果用“平方平均”（p=2，即 MXNorm 的推荐方案），就能像 RMSNorm 一样稳定，完美地防止了“尖叫”，让大模型训练过程平稳顺利。

4. 总结：为什么这很重要？

这篇论文的核心思想是**“物尽其用，拒绝重复劳动”**。

在 AI 算力越来越强的今天，单纯靠堆硬件已经不够了。MXNorm 告诉我们，通过巧妙地复用计算过程中已经产生的数据（块尺度），我们可以消除那些拖慢速度的“小瓶颈”。

一句话总结：
MXNorm 就像是一个聪明的指挥家，他不再拿着尺子重新测量每个乐手，而是直接利用乐手们为了买衣服已经量好的“最高身高”数据，瞬间算出了平均身高。这让整个乐团（AI 大模型）能跑得更快、更稳，而且不需要花一分钱买新设备。

技术摘要

MXNorm：利用 MXFP 块尺度实现高效张量归一化技术总结

1. 研究背景与问题 (Problem)

随着深度学习工作负载的扩展，矩阵乘法（MatMul）的性能瓶颈已逐渐被新一代 AI 加速器（如 NVIDIA H100, GB200）通过低精度格式（如 FP8, FP4）的优化所解决。然而，**归一化（Normalization）和逐元素计算（Elementwise operations）**的性能提升并未跟上矩阵乘法的步伐。

• 性能失衡：过去八年，GPU 的低精度矩阵乘法吞吐量提升了约 80 倍，但内存带宽和 CUDA 核心吞吐量仅提升了约 9 倍和 5 倍。
• 归一化瓶颈：在 Pre-Norm 架构（如 Llama 系列）中，RMSNorm 通常在量化之前执行。传统的 RMSNorm 需要计算整个隐藏维度的均方根（RMS），这涉及大量的归约（Reduction）操作，成为低精度训练中的新瓶颈。
• 现有方案局限：现有的加速归一化尝试（如异步计算或仅使用部分向量元素）往往会导致精度损失或无法处理异常值（Outliers），影响训练稳定性。

核心问题：如何在保持训练稳定性的前提下，消除归一化步骤中的冗余计算，使其与 MXFP（Microscaling Floating Point）量化过程融合，从而减少归约操作并提升推理/训练速度？

2. 方法论 (Methodology)

作者提出了 MXNorm，这是一种 RMSNorm 的即插即用（drop-in）替代方案。其核心思想是**复用 MXFP 量化过程中已经计算出的块尺度（Block Scales）**来估计 RMS，从而避免单独进行全维度的归约计算。

2.1 核心原理

在 MXFP 量化中，张量被划分为大小为 $B$ 的块，每个块计算一个 8 位的块尺度（Block Scale），通常基于该块的绝对最大值（Block Absmax）。

• 观察：MX 量化和 RMSNorm 都需要沿隐藏维度收集统计信息。
• 理论依据：作者证明了对于服从特定分布（如高斯分布）的张量，块的绝对最大值的广义 $p$-均值（Generalized $p$-mean）与整个张量的 RMS 之间存在一个恒定的比例关系。
  • 公式：$\tilde{\rho} \approx c(p, B) \cdot (\frac{1}{K} \sum_{k=1}^K \tilde{m}_k^p)^{-1/p}$
  • 其中 $\tilde{m}_k$ 是第 $k$ 个块的绝对最大值，$c(p, B)$ 是依赖于块大小和 $p$ 值的修正系数。

2.2 MXNorm 算法流程

1. 输入：激活张量 $X$。
2. 块统计：计算每个块的绝对最大值 $m_{tk} = \max_b |X_{tkb}|$。
3. RMS 估计：利用块最大值的 $p$-均值（$p=1$ 为算术平均，$p=2$ 为均方根）乘以修正系数 $c(p, B)$ 来估计全局逆 RMS $\tilde{\rho}$。
4. 融合量化与归一化：
   • 直接使用估计的 $\tilde{\rho}$ 对原始数据进行缩放。
   • 计算新的块尺度 $\tilde{S}$ 和量化值 $\tilde{V}$，直接输出 MXFP 格式。
   • 优势：只需一次统计收集（Pass），将归一化和量化融合在一个内核中。

2.3 梯度处理与参数融合

• 梯度计算：为了保持训练平滑，MXNorm 使用 RMSNorm 的梯度计算作为直通估计器（Straight-Through Estimator）。
• 增益参数（Gain）：由于在 MX 格式上直接进行仿射变换（乘以 $\gamma$）较复杂，作者利用线性运算的结合律，将归一化增益 $\gamma$ 融合到后续的线性层权重矩阵中（即 $W_{fused} = W \cdot \gamma$），从而在推理时消除额外的乘法操作。

3. 关键贡献 (Key Contributions)

1. MXNorm 架构设计：提出了一种无需额外超参数、可直接替换 RMSNorm 的方案，利用 MXFP 量化中已有的块尺度信息来近似 RMS。
2. 理论证明：证明了块绝对最大值的广义均值与 RMS 之间的渐近收敛性，并推导了输出值的上界。
3. 稳定性分析：深入分析了不同 $p$ 值（$p=1$ vs $p=2$）对训练稳定性的影响。发现 $p=2$（基于均方根的估计）能提供比 $p=1$ 更紧的输出上界，从而有效抑制训练过程中的损失尖峰（Loss Spikes）和异常值爆炸。
4. 性能优化：实现了归一化与量化的内核融合，显著减少了归约操作的数据量（减少 32 倍），并优化了内存访问模式。

4. 实验结果 (Results)

作者在 Llama 3 模型（125M, 1B, 8B 参数）上进行了预训练验证，并在 Graphcore GB200 硬件上进行了性能测试。

4.1 训练稳定性与精度

• 小模型（125M, 1B）：MXNorm ($p=1$ 和 $p=2$) 与基线 RMSNorm 在最佳学习率下的训练损失几乎没有差异。
• 大模型（8B）：
  • $p=1$ 失败：使用算术平均（$p=1$）估计 RMS 导致训练不稳定，出现明显的损失尖峰，最终损失显著高于基线。
  • $p=2$ 成功：使用均方根（$p=2$）估计 RMS 能够完美复现 RMSNorm 的训练曲线和最终损失（2.126 vs 2.132）。
  • 下游任务：在 OLMES 基准测试中，MXNorm ($p=2$) 的表现与 RMSNorm 相当，在 10 个基准中有 5 个表现更优。
• 原因分析：$p=2$ 提供了更紧的输出上界（$O(\sqrt{K})$），而 $p=1$ 的上界较宽（$O(K)$），导致 $p=1$ 无法有效抑制训练早期的异常值，从而引发梯度爆炸。

4.2 性能加速

• 内核速度：在 GB200 上，仅使用 torch.compile 优化，MXNorm 相比 "RMSNorm + MXCast" 的组合实现了最高 2.4 倍 的内核加速。
• 平均加速：
  • MXFP8 内核平均加速 41.7%。
  • NVFP4 内核平均加速 31.2%。
• 端到端影响：在 Llama 3 8B 的 Transformer 层中，MXNorm 带来了 1.3% (MXFP8) 和 2.6% (NVFP4) 的整体速度提升。这证明了随着精度降低，优化非矩阵乘法操作的重要性日益凸显。

5. 意义与展望 (Significance)

• 解决新瓶颈：MXNorm 有效地解决了低精度训练中归一化操作的性能瓶颈，填补了矩阵乘法加速与归一化/归约操作加速之间的差距。
• 工程落地性强：作为一种“即插即用”的替代方案，它不需要改变模型架构，且能直接利用现有的 MXFP 硬件特性，无需额外的软件工程负担。
• 理论指导实践：研究揭示了归一化层输出上界对训练稳定性的关键作用，为未来设计更激进的量化归一化方案提供了理论依据。
• 扩展性：该方法不仅适用于 MXFP，还可推广至其他基于块最大值的量化格式（如 INT2, 三值量化等），具有广泛的适用性。

总结：MXNorm 通过巧妙地复用量化过程中的中间统计量，成功实现了归一化与量化的融合，在保持 LLM 训练稳定性和精度的同时，显著提升了计算效率，是迈向超低精度大模型训练的重要一步。