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.

速读

一句话结论 本研究在原生支持 FP4 张量计算的硬件上，通过在矩阵乘法前注入确定性 Hadamard 旋转，成功解决了大语言模型全流程 MXFP4 量化训练发散的问题，在保证收敛的前提下实现了端到端训练吞吐量的净加速。

要解决什么问题 原有的大语言模型训练在向更低精度的 FP4 格式迁移时，面临着收敛变慢甚至完全发散的严重卡点。虽然原生 FP4 硬件能大幅降低大批量训练中的内存带宽压力并提升吞吐量，但全流程使用 MXFP4（一种采用微缩放机制的 4 位浮点格式）会导致训练极不稳定。通过分阶段启用量化的受控实验发现，这种不稳定性并非来自前向传播或激活梯度的量化，而是由权重梯度量化时产生的结构化微缩放误差驱动的。传统的随机性干预措施（如随机舍入或随机符号翻转）不仅无法缓解离群值带来的影响，反而会引入额外噪声，进一步放大微缩放误差，导致模型彻底无法收敛。因此，如何控制最敏感的权重梯度路径上的微缩放误差，而不是盲目增加随机性，成为解锁 FP4 硬件加速潜力的关键。

怎么做的 核心思路是在 GEMM 操作之前注入确定性的 Hadamard 变换，通过在特征空间中进行正交旋转，将由离群值导致的集中能量均匀分散到更多维度上，从而降低进入量化器的数据分布对微缩放误差的敏感度。MXFP4 格式不同于传统整数量化，它采用了微缩放机制，即为一小块数据（如 32 个元素）共享一个指数，其重建值定义为： $$x_{B, reconstructed} = 2^{E_{shared}} \times Q_{FP4}(x_B / 2^{E_{shared}})$$ 其中 $E_{shared}$ 是该数据块内所有元素指数的最大值，$Q_{FP4}$ 表示舍入到最接近的 4 位浮点值的操作。为了在不改变底层线性层算术结果的前提下压制离群值，该方法利用了 Hadamard 矩阵 $H$（元素仅为 +1 或 -1 的正交矩阵）的性质： $$H H^T = I$$ 关键设计由前向传播和反向传播中的数学抵消机制构成。在前向传播中，对输入和权重分别应用变换 $\tilde{X} = XH$ 和 $\tilde{W} = WH$，计算输出时 Hadamard 因子自动抵消： $$Y = \tilde{X} \tilde{W}^T = X H H^T W^T = X W^T$$ 在反向传播计算激活梯度时，对输出梯度和转置权重应用变换 $\tilde{G} = (\nabla Y)H$ 和 $\tilde{W}_{rot} = W^T H$，计算结果同样完美抵消： $$\nabla X = \tilde{G} \tilde{W}_{rot}^T = (\nabla Y) H H^T W = \nabla Y W$$ 在最关键的权重梯度计算中，对转置的输出梯度和转置的输入应用变换 $\tilde{G}^T = (\nabla Y)^T H$ 和 $\tilde{X}^T = X^T H$，抵消过程为： $$\nabla W = \tilde{G}^T \tilde{X} = (\nabla Y)^T H H^T X = (\nabla Y)^T X$$ 通过这种确定性的正交旋转，该方法成功绕开了权重梯度量化时的结构化误差卡点，既不引入破坏性的随机噪声，又能在量化前平滑数据分布，从而稳定整个训练流程。

效果如何 实验在原生支持 FP4 张量计算的 AMD Instinct MI355X GPU 上进行，避免了软件仿真的偏差。研究采用 Llama 3.1–8B 模型，在 C4 数据集上遵循 MLPerf 预训练设置进行测试，目标是达到 3.3 的验证集困惑度。对比基线为全流程 FP8 训练（代表稳定收敛的标准路线），以及引入随机舍入或随机 Hadamard 旋转的全流程 MXFP4 训练（代表试图通过增加随机性缓解量化误差的传统路线）。结果显示，当在全流程中直接启用 MXFP4 时，收敛所需的 token 开销相比 FP8 基线急剧上升了 26% 至 27%；而采用随机舍入或随机旋转的基线方法则完全发散。相比之下，采用 16x16 确定性 Hadamard 变换的全流程 MXFP4 训练成功恢复了稳定性，将 token 开销降低至 8% 到 9%。在端到端效率方面，由于降低位宽缓解了内存带宽压力，该方法使训练步吞吐量提升了 1.21 倍，扣除 8% 的 token 开销后，最终端到端吞吐量实现了 1.12 倍的净加速。此外，16x16 分块矩阵的内核吞吐量比 32x32 分块快 8%。作者也承认了该方法的局限性：这并非万能方案，MXFP4 的训练行为高度依赖具体设置，在全量预训练中有效的稳定策略未必能直接推广到其他模型结构或微调任务中。

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$，其重建值的计算方式如下：

    公式1：MXFP4重建值计算
    公式1：MXFP4重建值计算

    在这个公式中，$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 架构图

图3：用于前向和后向传播的Hadamard变换MXFP4架构。GEMM内核的输入由H进行旋转，并通过HHT=I在矩阵乘法过程中抵消旋转。
图3：用于前向和后向传播的Hadamard变换MXFP4架构。GEMM内核的输入由H进行旋转，并通过HHT=I在矩阵乘法过程中抵消旋转。

C 数学抵消证明

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

    公式2：Hadamard矩阵的正交性
    公式2：Hadamard矩阵的正交性

    其中 $I$ 是单位矩阵。

C.1 前向传播 (Fprop)

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

    公式3：前向传播中的数学抵消
    公式3：前向传播中的数学抵消

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

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

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

    公式4：激活梯度计算中的数学抵消
    公式4：激活梯度计算中的数学抵消

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

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

    公式5：权重梯度计算中的数学抵消
    公式5：权重梯度计算中的数学抵消

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变换来改善收敛稳定性和减少离群值影响的先例。