Pretraining Large Language Models with MXFP4 on Native FP4 Hardware

发表时间: 2026-05 · arXiv:2605.09825 (AMD)

作者/机构: Musa Cim¹, Poovaiah Palangappa², Miro Hodak², Ravi Dwivedula², Meena Arunachalam², Mahmut Taylan Kandemir¹
¹The Pennsylvania State University, ²Advanced Micro Devices, Inc.

A1 主要贡献

本文旨在探究一个核心问题：为何即使前向激活（forward activations）和激活梯度（activation gradients）保持稳定，大型语言模型的全流程FP4训练也常常会发散。为了回答这个问题，研究者们进行了一项关于在Transformer训练中使用MXFP4量化的受控研究。

研究目标与方法：
研究的核心目标是，在保持所有其他因素固定的情况下，通过逐步在前向传播（Fprop）、激活梯度（Dgrad）和权重梯度（Wgrad）中启用MXFP4，来识别导致训练不稳定的关键环节，并验证各种稳定策略的有效性。其动机在于，原生支持FP4的硬件（如AMD Instinct MI355X GPU）上的MXFP4 GEMM（通用矩阵乘法）操作相比FP8能提供更高的吞吐量和更低的内存带宽压力，但前提是必须解决其可能带来的收敛变慢甚至训练发散的问题。

核心发现：
本文通过在原生支持MXFP4的AMD Instinct MI355X GPU上对Llama 3.1–8B模型进行C4数据集上的预训练，得出以下三个关键结论：

• 定位不稳定性来源 (Where)：通过逐步启用MXFP4的实验揭示，训练收敛性的恶化主要由权重梯度（Wgrad）的量化导致。相比之下，仅在前向传播（Fprop）和激活梯度（Dgrad）中引入MXFP4对收garan't收敛所需token数的影响较小，训练过程相对稳定。

• 解释不稳定性原因 (Why)：当Wgrad被量化后，随机性干预措施（如随机舍入和随机Hadamard旋转）未能稳定训练过程。这表明，简单地增加随机性并不能解决问题，反而可能因为引入额外噪声而放大了量化误差，而不是像预期的那样缓解离群值效应。

• 提供稳定化方案 (How)：在随机方法失败的相同受控环境下，确定性的Hadamard旋转成功地恢复了全流程MXFP4训练的稳定性。这一结果表明，FP4训练的不稳定性是由敏感梯度路径上结构化的微缩放（micro-scaling）误差驱动的，而非随机性不足。因此，确定性Hadamard旋转不仅是一种有效的稳定器，也成为一种隔离和分析FP4训练中稳定结构的实用探针。

所有实验均在原生硬件上完成，避免了软件仿真的开销和潜在偏差，为研究FP4量化效应提供了可靠的基础。

A2 方法细节

A.1 MXFP4与微缩放（Micro-scaling）

• MXFP4的微缩放机制：标准的整数 quantization 通常对整个张量使用单一的缩放因子。与此不同，MXFP4采用了微缩放（micro-scaling）技术，即为数据的小块（例如，32个元素组成的块）定义一个共享的指数。具体来说，如果一个向量 $x$ 被划分为多个块，对于其中一个块 $x_B$，其重建值的计算方式如下：

  

  在这个公式中，$E_{shared}$ 是该数据块中所有元素指数的最大值【7，Microscaling data formats for deep learning，2023，arXiv】，而 $Q_{FP4}$ 表示将数值舍入到最接近的可表示的4位浮点值的操作。

A.2 Hadamard变换架构

• 采用确定性Hadamard变换：为了提升收敛稳定性并减小离群值的影响，本研究采用了一种确定性的Hadamard变换【9，Outlier suppression: Pushing the limit of low-bit transformer language models，2022，Advances in Neural Information Processing Systems】。尽管随机化的Hadamard旋转（例如，随机符号翻转）被广泛使用，但本研究的实验发现，在这种设置下，随机符号可能会对收敛稳定性产生负面影响。因此，研究最终采用了确定性的变换。

• Hadamard矩阵及其作用：Hadamard矩阵 $H$ 是一个其元素仅为+1或-1的正交矩阵。应用 $H$ 变换可以被理解为在特征空间中进行一次旋转，它能将集中的能量（通常由离群值导致）分散到更多的维度上。在具体的实现中，研究者使用了分块（tiled）布局，并选择了16x16大小的Hadamard变换（H16），因为发现在目标硬件架构上，H16相比H32能提供更高的内核吞吐量，并且经验证能够保持稳定的收敛。

• Hadamard旋转的注入流程：研究流程在GEMM操作之前注入Hadamard旋转（如图3所示）。由于该变换是线性和正交的，它在不改变底层计算结果的同时，改变了进入量化器的数据分布。具体操作上，研究者使用AMD的ROCm Transformer Engine，将Transformer线性层中使用的FP8 GEMM内核替换为MXFP4 GEMM内核（可选择性地应用于Fprop、Dgrad和/或Wgrad）。这些内核通过标准的Transformer线性模块在训练栈中被调用。

B 架构图

C 数学抵消证明

• 利用正交性进行抵消：为了确保Hadamard变换不改变线性层的算术结果，研究利用了其正交属性：

  

  其中 $I$ 是单位矩阵。

C.1 前向传播 (Fprop)

• 前向传播中的变换与抵消：在一个标准的线性层中，输出 $Y = XW^T$。研究者将Hadamard变换应用于两个输入：$\tilde{X} = XH$ 和 $\tilde{W} = WH$。然后，计算过程如下：

  

  因此，Hadamard因子相互抵消，在提升量化鲁棒性的同时，保持了与原始计算完全相同的输出。

C.2 后向传播 - 激活梯度 (Dgrad)

• 激活梯度计算中的变换与抵消：对于输入梯度，标准操作是 $\nabla X = \nabla Y W$。在应用变换后，输入变为 $\tilde{G} = (\nabla Y)H$ 和 $\tilde{W}_{rot} = W^T H$。计算过程如下：

  

C.3 后向传播 - 权重梯度 (Wgrad)

• 权重梯度计算中的变换与抵消：对于权重梯度，标准操作是 $\nabla W = (\nabla Y)^T X$。在研究流程中，变换后的输入为 $\tilde{G}^T = (\nabla Y)^T H$ 和 $\tilde{X}^T = X^T H$。计算过程如下：

  

D 1D与2D量化策略

• 量化策略探索：研究者探索了用于激活和梯度的一维（1D）和二维（2D）块量化策略。

• 2D块量化 (32×32)：该策略将张量划分为32×32的区域，并为每个$C_{in} \times C_{out}$块使用一个E8M0格式的逐块缩放因子。这种方法计算效率高，对转置操作友好，并且能够在保持良好精度的同时利用硬件矢量化。

• 1D行式量化 (1×32)：该策略与Transformer中的逐令牌（token-wise）处理方式保持一致。它对每32个连续元素计算一个逐令牌的E8M0缩放因子，从而保留了令牌级别的信息，非常适用于激活值的量化。

• 1D列式量化 (32×1)：该策略通过对每32个元素计算一个逐$C_{out}$的缩放因子，为权重提供了细粒度的逐通道（per-channel）缩放。然而，由于需要转置操作和不友好的内存访问模式，这种方法的计算成本较高。

PROTECTED_IMAGE_8____PROTECTED_IMAGE_9

A4 实验环境

• 模型与任务：实验采用 Llama 3.1–8B 模型，在 C4 数据集上进行预训练，遵循 MLPerf LLM 预训练的设置。
• 收敛标准：定义收敛为模型的验证集困惑度（validation perplexity）首次达到或低于3.3。
• 硬件配置：实验在 AMD Instinct MI355X GPU 上进行，该硬件原生支持FP4张量计算，无需软件仿真。
• 软件配置：通过 AMD 的 ROCm Transformer Engine 调用 MXFP4 GEMM 内核，并将其集成到标准的 Transformer 线性模块中进行训练。
• 评估指标：主要评估指标是达到收敛目标所需的训练token数量，并将其与FP8基线进行比较，计算出token开销（即相对增加的token百分比）。

A4 实验结果

实验假设

实验旨在验证四个核心假设：
1. (H1) 逐步启用MXFP4会揭示，收敛性恶化主要由量化Wgrad（权重梯度）引起，其影响远大于Fprop（前向传播）或Dgrad（激活梯度）。
2. (H2) 随机方法（如随机舍入或随机Hadamard）提供的益处有限，因为它们只增加了噪声，而未能解决由离群值驱动的MXFP4微缩放误差。
3. (H3) 确定性Hadamard旋转能够减少进入MXFP4块的离群值的影响，从而稳定全流程优化。
4. (H4) 在更多GEMM路径中启用MXFP4会提高训练步吞吐量，并且当token开销很小时，这能转化为端到端的训练加速。

分阶段MXFP4启用与稳定策略评估

为隔离导致收敛下降的组件，实验逐步在训练流程中启用MXFP4，并评估了不同稳定策略的效果。

• 实验结果：如表1所示，与FP8基线相比：

  • 仅对Fprop应用MXFP4，导致8-9%的token开销。
  • 扩展到Fprop + Dgrad，token开销相似，为10-11%。
  • 当在整个流程（Fprop + Dgrad + Wgrad）中启用MXFP4时，token开销急剧上升至26-27%。这证实了Wgrad量化是收敛恶化的主要驱动因素（验证H1）。
  • 随机舍入和随机Hadamard在部分启用时表现尚可，但在全流程应用时均无法收敛（验证H2）。
  • 相比之下，对全流程MXFP4应用确定性Hadamard变换（H16）成功恢复了稳定的训练，并将token开销降低至8-9%，与FP8基线接近（验证H3）。
  • 此外，实验发现H16（16x16矩阵）内核比H32（32x32）快8%。

• 关键理解 #1：逐步启用MXFP4的实验表明，收敛性下降主要由量化Wgrad引起。确定性Hadamard能够实现稳定的全流程MXFP4训练；而当Wgrad被量化时，随机技术会因引入更高的量化误差而失败。

Table 1: 分阶段MXFP4启用与稳定策略。Token开销是相对于FP8基线（验证困惑度3.3）而言。“不收敛”表示该运行发散或在延长训练后仍未能达到困惑度3.3的目标。

内核吞吐量：H16比H32快8%（1.08倍 vs 1.00倍）。

训练轨迹可视化

图1和图2展示了不同配置下的训练轨迹。

• 实验结果：
  • 全流程MXFP4（无稳定器）的收敛速度明显慢于FP8基线，且稳定性较差。
  • 而添加了确定性Hadamard变换的全流程MXFP4，其训练曲线与FP8基线紧密贴合，稳定地达到了MLPerf的目标困惑度3.3。

PROTECTED_IMAGE_10____PROTECTED_IMAGE_11

端到端训练效率

在证明了确定性Hadamard可以恢复收敛稳定性后，实验进一步评估了FP4带来的实际效率提升。

• 实验结果：如表2所示，在MI355X GPU上，稳定的MXFP4+H16配置相比FP8基线：

  • 训练步吞吐量提高了1.21倍。
  • 由于token开销仅为8%，最终的端到端吞吐量（Tokens/s）提高了1.12倍，实现了净加速（验证H4）。

• 关键理解 #2：FP4的实际加速是稳定性门控（stability-gated）的。如果没有稳定Wgrad量化，那么因收敛变慢或发散而增加的token需求会完全抵消FP4在硬件计算上的吞吐量优势。

Table 2: 端到端训练效率（MXFP4 + H16 vs. FP8基线）。

该加速的关键在于，降低激活/权重的位宽减小了内存带宽压力，这在大批量训练中是一个主要瓶颈。而本方法在实现这一点的同时，还稳定了全流程的训练。

A7 补充细节

• Wgrad的敏感性和稳定性：本研究的结果表明，在全流程MXFP4训练中，Wgrad（权重梯度）是最敏感的环节。稳定FP4预训练的关键在于控制微缩放误差，而不是注入额外的随机性。

• 方法的非通用性：同时，本文强调这个稳定方案并非“万金油”。FP4的训练行为可能因不同的模型、数据集和适应方法（如微调）而异。

• 关键理解 #3：MXFP4的训练行为高度依赖于具体设置。一个在全预训练（如本文的MLPerf C4数据集上的Llama 3.1–8B）中有效的稳定策略，未必能推广到其他模型或微调任务中。因此，在没有进一步证据的情况下，FP4的训练配方不应被视为普遍适用。

A5 结论

本文通过一项对Transformer线性层中用MXFP4 GEMM内核替代FP8的受控研究，揭示了权重梯度（Wgrad）的量化是导致收敛性下降的主要因素。对于在MLPerf C4数据集上预训练的Llama 3.1–8B模型，在所有测试的干预措施中，确定性Hadamard旋转是唯一能够恢复全流程MXFP4训练稳定性，并同时提升端到端效率的方法。

总的来说，本研究结果为FP4大型语言模型训练提供了一个实用的原则：训练的稳定性取决于控制最敏感梯度路径中的微缩放误差，而不是增加随机性。

方法细节中的引用文献

1. 【7】Microscaling data formats for deep learning (2023, arXiv) - Bita Darvish Rouhani, et al.

   • 引用位置: 附录A.1 MXFP4与微缩放。
   • 引用内容: 描述MXFP4中微缩放块的共享指数 $E_{shared}$ 是块内指数的最大值。

2. 【9】Outlier suppression: Pushing the limit of low-bit transformer language models (2022, Advances in Neural Information Processing Systems) - Xiuying Wei, et al.

   • 引用位置: 附录A.2 Hadamard变换架构。
   • 引用内容: 引用该文献作为使用Hadamard变换来改善收敛稳定性和减少离群值影响的先例。