NVIDIA：Blackwell GPU MXFP8 预训练最佳实践

一、背景

笔者之前写过 FP8 训练的综述文章以及 FP4 训练和推理的综述文章，本文对其进一步补充，介绍 NVIDIA 最新的使用 MXFP8 预训练的方案。

对应的论文：[2506.08027] Recipes for Pre-training LLMs with MXFP8 [1]

二、摘要

精度缩放——即在预训练过程中使用更少的比特来表示模型参数及相关 Tensor——已成为一种在不牺牲精度前提下提升 GPU 效率的有效技术。NVIDIA 最新 Blackwell GPU 中引入 Microscaling (MX) 格式，为 Tensor 量化提供了细粒度解决方案。

尽管 MX 格式相较于其他低精度表示法有望提升数值稳定性，但在实际应用中仍需谨慎使用。本文研究表明：当采用 OCP 规范建议的舍入模式进行 LLM 预训练时，会导致模型不收敛。为此，作者提出了一种改进的舍入模式——通过采用"向无穷大舍入"方式计算缩放因子，成功实现了 8B 参数模型在 15T Token 上采用 MXFP8 格式的预训练。

PS：可能是因为采用 Hopper GPU 模拟的方式，而不是真实的 Blackwell 训练，因此论文并没有提供相应的效率提升数据。

三、引言

3.1 MXFormat

2023 年，OCP（Open Compute Project） 在 AMD, Arm, Intel, Meta, Microsoft, NVIDIA, Qualcomm 的参与下提出 Microscaling（MX）Format 规范（OCP Microscaling Formats (MX) Specification Version 1.0 [2]），主要是为了对跨硬件/软件平台可实施的新功能及格式进行标准化，有效减少软件与基础设施成本，并消除定制化解决方案带来的各类附加费用或管理负担，推动硬件性能与效率的提升。

如下图所示，MX 最主要的特点是其包含三部分内容（很类似于常见的 Per-Block 细粒度量化方式，只不过这里是制定了一个统一的规范）：

• P：规定了 d 个 bit 数据的表示（编码）方式，比如 FP8 的 E5M2 是怎么表示的。
• k：k 个元组作为一个 Block。
• X：上述 k 个元素的 Block 会对应一个共享的 Scale 值。

我们在之前的文章中提到过，即使都是 E5M2 或者 E4M3，不同公司的硬件可能采用不同的格式。比如 NVIDIA Hopper GPU 上的 E5M2 符合 IEEE 754 Style，而 E4M3 却不符合 IEEE 754 Style。如下图所示，IEEE 754 Style 的 E4M3 的范围为 [-240, 240]，而 ARM-Intel-Nvidia Style 的 E4M3 的范围是 [-448, 448]： 

在 MX 中也对上述问题进行了规范化，以 MXFP8 为例，其规定的 E4M3 和 E5M2 编码方式如下图 Table 1 和 Table 2 所示：

3.2 细粒度量化

更低的精度通常意味着更难量化，为了维持精度需要更细力度的 Scaling Factor，比如：

• FP16：早期使用 FP16 进行混合精度训练时通常整个模型一个 Scaling Factor 即可。
• FP8：在 Inference 时通常 Per-Tensor 的 Scaling Factor 即可比较好的维持精度；而 Training 时往往需要 Per-Block 或 Per-Channel，不过 Block 通常比较大，比如 128x128 或 128x1。
• FP4：需要 Per-Block 量化，并且 Block 需要比较小，比如 32 或 16。

更细粒度的量化也意味着更高的额外成本，Block 越小（粒度越细），额外成本越高。如下图所示，对于一个常见的内积操作：

• Per-Tensor 量化：需要额外执行 1 次 Scaling Factor 处理。
• Per-Block 量化：需要额外执行很多次 Scaling Factor 处理。

为了更好的解决上述问题，NVIDIA 在新的 Blackwell Tensor Core 中支持了新的 Block-Scaled 类型，原生支持 Microscaling Formats，如下图 Table 1 所示，其支持 MXFP8、MXFP6、MXFP4：

如下图所示，在 Tensor Core 计算时，可以将数据 A/B 及它们对应的 Scaling Factor A/B 一起输入，并全部在 Tensor Core 内完成。

当然，其对数据类型也有一定的要求，如下图所示：

四、MXFP8 预训练

4.1 转换 FP32 到 MXFP8

在训练的 Forward 与 Backward 过程中，模型 Weight、Activation 及 Gradient Tensor 均由 FP32 量化到 MXFP8 格式。量化后的 MXFP8 Tensor 随后存储于硬件中并执行运算。作者首先阐述了转换过程 Quantize_to_fp8(Vi/2X)。下文所述的量化方法统一适用于所有 MX 格式（包括 E4M3、E5M2、E2M3、E3M2 及 E2M1），仅 MX 数据类型存在差异。

计算 X 值：通常情况下，Tensor 中各 Block 内的大部分数值会超出目标 MX 格式的可表示范围，既可能低于最小可表示数（下溢），也可能高于最大可表示数（上溢）。为解决这一问题，需将 Block 内所有数值乘以一个 Scale 因子，使绝大多数数值被调整至可表示范围内。

该 Scale 因子 X 的计算基于 32 个（MX 规范，k）高精度输入值中的绝对最大值（amax），即amax = max(‖Vi‖); 1≤i≤32。其核心目标是将输入中的 amax 映射为 MX 格式中的最大可表示值。当输入数据包含无穷大（Infinity）或非数值（NaN）时需特殊处理：若某 Block 的 amax 为 0，则设定 X=-127，此时 X 为 2-127，且该情况下所有 Qi’ 均置为 0。

根据 OCP 规范，当 X 不为 Inf、NaN 或 0 时，X 应设定为不超过 “amax 除以 MX 格式类型最大可表示二次幂” 的最大二次幂。以 E4M3 类型为例，由于 448 是其最大幅值，故 X 的计算公式为X=floor(log₂(amax))/floor(log₂(448))。值得注意的是，OCP 规范在此计算过程中忽略了该比率浮点数尾数部分的影响。

作者观察到遵循 OCP 规范时存在精度下降现象。如下图 Figure 2 所示为两种 Token 规模（300B 和 1T）下训练的 843M 参数量 Transformer 模型的训练损失曲线。其采用两种不同配置方案：

• cfg1：所有 Tensor（Weight W、Activation A、Gradient G）均采用 E4M3 格式。
• cfg2：Weight W、Activation A 采用 E4M3 格式，Gradient G 采用 E5M2 格式。

E5M2 格式相较于 E4M3 具有约 1.8 倍的 binades 优势。鉴于 Gradient 通常具有更大的动态范围，早期的工作 [2209.05433] FP8 Formats for Deep Learning [3] 主张采用 E5M2 格式进行 Tensor 缩放。实验结果表明，在 cfg1 与 cfg2 中，使用 OCP 方法计算缩放因子均会导致训练发散（如下图 Figure 2a）或相对于 BF16 基准的损失差距扩大（如下图 Figure 2b）。

如下图 Algorithm 1 概述了作者计算 Scale 因子的方法。其核心改进在于：当处理 amax 与 MX 格式最大可表示值 destmax 的比值指数时，采用向正无穷方向的 round-up 策略（同时饱和至 UE8M0 格式的极值边界）。这与 OCP 方案形成鲜明对比，后者实质上是建议对 Scale 值执行 round-down 操作。由于高精度值 Vi 需通过 Scale 因子 2X 进行缩放，对分数项 (Vi/2X) 分母实施 round-up 操作，会更倾向于将 amax 映射至 destmax 以下；反之，OCP 方法则倾向于使 amax 超过 destmax（后续必须通过截断处理使其可表示）。作者推测 OCP 取整方法带来的饱和效应会影响模型精度。

如上图 Figure 2 所示，采用提出的舍入方案后，Gradient 位宽配置为 E4M3 的 MXFP8（蓝色曲线）与 E5M2 的 MXFP8（紫色曲线）在 300B 和 1T Token 的训练过程中，其损失曲线均与 BF16 完全重合。

FP32 数值到 MX 格式的量化过程：当缩放因子 X 确定后，Tensor Vi 通过乘以 2X 进行尺度变换，随后量化至最接近的 FP8 可表示数值（即 Quantize_to_fp8()）。该量化步骤采用“就近取偶（Round-to-nearest-ties-to-even，RN）”舍入法，且转换过程具有饱和特性——若舍入结果超出 FP8 最大值或低于最小值，则将结果截取至相应的极值。

这种转换机制在低精度 LLM 预训练中的典型应用场景是：矩阵乘积累加运算（MMA）的输出（通常以 FP32 格式存储）需要映射为 MXFP8 格式，相比存储 FP32 数值可显著节省写入带宽和存储容量。模型后续运算读取 MXFP8 数值时，相较加载 FP32 数据也能减少读取带宽消耗。此外，由于 Tensor Core 可直接处理 MX 格式输入，低精度 MMA 操作不仅能降低能耗，还能获得更高的计算吞吐量。

4.2 所有 Tensor 采用 E4M3

在 Blackwell 架构中，FP8 浮点格式包含两种变体：E4M3 与 E5M2。实验研究表明：

Weight 与 Activation 量化性能对比：采用 E4M3 格式量化 Weight 和 Activation 时展现出更优的训练收敛性。如下图 Figure 3a 所示（测试模型与 Figure 2b 相同，参数量 843M），当 Activation（紫色曲线）或 Weight（蓝色曲线）采用 E5M2 格式时，其损失函数收敛性显著差于所有 Tensor 采用 E4M3 量化方案（橙色曲线）。值得注意的是，仅 Gradient 采用 E5M2 时（黄色曲线）仍能维持较好的收敛特性。

Gradient Tensor 量化分析：E4M3 格式在 Gradient 量化中能保持与 BF16 预训练相当的损失，这一优势在参数量 ≥2B 的模型中尤为显著。如下图 Figure 3c 展示了 8B LLM（1T Token 训练）的对比结果：E4M3 Gradient 量化（橙色曲线）的最终损失值显著低于 E5M2 方案（黄色曲线），且该差距随训练 Token 数量增加而扩大。这一现象揭示了模型参数量对数值格式选择的敏感性，强调需在不同规模模型中系统评估格式的数值特性（PS：这也是为什么笔者一直提到之前很多文章只在小规模模型、数据量下做实验不够有说服力的原因）。

既往研究采用的 Tensor 级缩放 FP8 方案及 DeepSeek V3（[2412.19437] DeepSeek-V3 Technical Report [4]） 提出的粗粒度 Block-Scaled 方案均默认选用 E5M2 格式处理 Gradient Tensor（PS：论文这里表述有问题，DeepSeek V3 中其实所有 Tensor 都已经采用 E4M3 格式；此外，这里的粗粒度是相对 32 的 Block 大小而言，在 DeepSeek 中为了效率采用的是 128 或者 128x128 的 Block 大小）。本研究发现：当采用细粒度缩放（32 元素 Block）时，E4M3 格式的 17.8 个 binades 可充分满足动态范围需求。在满足动态范围前提下，量化精度成为关键因素——E4M3 每个指数区间包含 8 个量化样本，其采样密度是 E5M2（4样本/区间）的 2 倍。因此，提出的 MXFP8 预训练方案对所有三类 Tensor（Weight、Activation、Gradient）均采用 E4M3 数据类型进行量化。

MXFP8 实例化层级及训练流程：论文所有研究均采用基于语言的 Transformer 模型，未来工作将探索该方案在语音与视觉模型中的应用。研究表明，量化策略建议：

• 将模型中所有 Transformer  Block 的 QKV、Proj 以及 FFN 的 Up-proj 和 Down-proj 转换为 MXFP8 格式。
• Self-Attention 中的批量矩阵乘法（BMM1：Query-Key 点积和 BMM2：Attention Score-Value 乘积）以及 Softmax、激活函数和残差相加等运算仍保持高精度计算。
• 输入 Embedding 层和最终输出 LM-head 同样采用 BF16 或 FP16 格式。

如下图 Figure 4 所示，这种配置能最可靠地维持与 BF16 预训练相当的精度水平，论文所有实验均遵循此准则。在 MXFP 量化训练过程中，框架需为 Tensor（Weight、Activation 和 Gradient）保持两个副本：每个副本沿点积归约（dot-product reduction）轴（行与列）分别量化。Figure 4 展示了训练迭代中各 Tensor 在 Forward（FPROP）、Weight Gradient（WGRAD）和 Activation Gradient（DGRAD）计算中的使用方式。由于每个 Tensor 需以原始和转置两种形态参与运算，量化需沿行列两个独立轴向分别执行。

当前研究总结：提出的 MX Scale 因子舍入方案解决了基于 OCP 方法导致的不收敛问题，在 843M 参数模型上实现了 1T Token 训练下与 BF16 相当的精度。结合 Algorithm 1 中的 E4M3 格式及 Scale 因子计算方法，该方案可扩展至 8B 参数模型（15T Token 训练）——作者声称，这是目前采用 MXFP 格式的最大规模 LLM 预训练案例。

如下图 Figure 6 所示，W16A2.5 规模 MoE 模型，1T Token 预训练也能实现同样的效果：

4.3 15T Token MXFP8 预训练结果

作者采用 Megatron-LM 框架预训练了一个 8B 参数的 Nemotron 模型。该模型包含 32 个 Transformer Block，每个 Block 32 个 Attention Head，隐层维度为 4096，采用 GQA 且 Group 大小为 8，KV 通道数为 128，预训练阶段序列长度为 8192。共训练 15T Token，Batch Size 为 768。初始学习率设为 6×10⁻⁴，并通过 cosine decay 到 6×10⁻⁶。如下图 Table 2 为几个模型的详细配置：

采用分阶段数据混合策略进行训练：第一阶段使用促进数据多样性的混合数据集，第二阶段则转向高质量数据集（如维基百科），在训练进度达到 60% 时切换至第二阶段。此类混合策略在其他大规模预训练框架中亦有应用。

模型预训练在 3072 Hopper GPU 上完成（实验周期内缺乏支持 MX 格式的 Bloackwell 硬件平台）。通过在 Hopper GPU 上模拟 MX 格式实现：输入矩阵乘法加速器（MMA）的 Tensor 先量化为 MX 格式，在执行 BF16 MMA 运算前转换回 BF16 格式。为验证模拟方案的数值保真度，作者与 Blackwell 平台上采用真实 MXFP8 格式训练的 2B 参数 LLM 进行对比实验，确认二者输出结果完全一致。

如下图 Figure 5 展示了 8B 预训练模型的训练损失及任务级准确率。可以看出，两组下游任务的评估分数：

• MMLU 上的 5-shot 分数。
• 9 个通用 Reasoning 基准（ARC-Challenge 与ARC-Easy、Race、PIQA、Winogrande、Hellaswag、OpenBookQA、Social IQA 和 Commonsense QA）上 1-shot 分数的平均值。

主要结果如下：

• 采用 MXFP8 预训练时，模型的验证困惑度与 BF16 预训练结果持平（Figure 5 左图）。在整个预训练过程中，MXFP8 与 BF16 的验证困惑度差异始终小于 0.50%。
• Figure 5 中、右两图显示了两组下游任务的评估分数。MXFP8 训练模型的得分与 BF16 训练模型完全匹配，证明 MXFP8 可作为 LLM 预训练的有效候选方案。

MXFP8 与 FP8 对比：除 MXFP8 和 BF16 外，Figure 5 还展示了传统 FP8 精度训练同模型的任务级分数。FP8 方案采用软件模拟的分块缩放技术，通过整体 Tensor 缩放使多数 Tensor 值落入量化格式的可表示范围。遵循 [2504.03624] Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models [5] 的 FP8 预训练设置建议：模型首尾 Transformer Block 保持 BF16 精度，其余 Block 的线性层量化为 FP8，该配置适用于 20T Token 规模的 8B 和 56B 参数 LLM 预训练。但保留部分 BF16 层会影响端到端加速比，并增加预训练复杂性——需额外决策哪些层维持高精度。实验表明，MXFP8 在这两组任务上无需任何 BF16 层即可达到与 FP8 相当的精度。

MXFP8 与分块 FP8 的对比：进一步地，诸如 Deepseek-V3 等研究表明，在使用 FP8 时需要缩小 Block 规模。在此配置下，部分 Tensor 需采用 1x128 向量级缩放，而其他 Tensor 则需实施分块（如 128x128）缩放，这增加了 GEMM Kernel 函数设计的复杂度。MXFP8 的原生支持则简化了这一过程——其细粒度缩放机制提供了更优的数值鲁棒性，同时规避了小 Block 尺寸与硬件速度之间的权衡问题。

综上所述，相比于 BF16 或 FP8 预训练，MXFP8 能保持同等精度。在 GB200 Blackwell 系统上，MXFP8 的吞吐量是 BF16 的 2 倍，这使得端到端 MXFP8 预训练速度超越 BF16 预训练。与 FP8 相比，MXFP8 方案还更加简便（所有层均可量化且缩放由硬件处理），同时保持同等或更优的吞吐性能。

五、参考链接：

• [1] https://arxiv.org/abs/2506.08027
• [2] https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
• [3] https://arxiv.org/abs/2209.05433
• [4] https://arxiv.org/abs/2412.19437
• [5] https://arxiv.org/abs/2504.03624

本文转载自​​AI闲谈​​，作者：AI闲谈