FP8-Flow-MoE: A Casting-Free FP8 Recipe without Double Quantization Error

作者/机构: Fengjuan Wang, Zhiyi Su, Xingzhu Hu, Cheng Wang, Mou Sun† Zhejiang Lab

A1 主要贡献

核心问题: 训练大型混合专家(MoE)模型因其极高的计算和内存需求而成本高昂。尽管低精度训练(如FP8)有望加速计算并减少内存占用,但现有实现仍严重依赖于以BF16为主的数据流,其中包含频繁的量化-反量化(Q/DQ)转换。这些冗余的类型转换削弱了FP8的理论效率。此外,若简单地通过在FP8中保持整个数据流来移除这些转换,会引入“双重量化误差”:沿不同维度量化的张量会累积不一致的缩放因子,从而降低数值稳定性。

研究目标: 本文旨在解决现有FP8 MoE训练中的效率瓶颈和数值稳定性问题,提出一种系统性的、以FP8为中心的数据流训练方案,以充分发挥FP8格式的性能优势,同时避免双重量化误差,确保训练的稳定收敛。

创新点:
1. 避免FP8数据流中的双重量化误差:通过引入一种感知缩放的转置算子(scaling-aware transpose operator),可以直接在行式和列式布局的FP8张量之间进行转换,同时保持缩放一致性。该设计消除了跨算子进行反量化再量化的循环,有效避免了FP8数据流中的双重量化误差。
2. 提出无类型转换的FP8训练方案及原生FP8算子套件:本文提出了FP8-Flow-MoE方案,整个MoE阶段均在FP8下运行,除了入口点外没有显式的类型转换。与简单的算子替换方法相比,类型转换操作的数量从12次减少到2次,标志着从以BF16为中心到以FP8为中心的计算范式转变。为支持该方案,本文还提供了一套原生FP8算子,其中一些已合并到上游仓库(如NVIDIA/TransformerEngine),其余将作为开源实现的一部分发布。
3. 对FP8训练进行全面的系统级分析:本文深入研究并比较了FP8和BF16在MoE组件中的表现,分析了算子层面和量化相关的影响。在671B参数的MoE模型上的评估表明,与BF16基线和朴素的FP8算子替换方法相比,FP8-Flow-MoE实现了高达21%的吞吐量提升和每GPU 7.4 GB的内存减少。而朴素的FP8替换方法仅带来3%的增益,内存节省微乎其微甚至为负。

A3 背景知识与相关工作

2.1 低精度训练和FP8量化

低精度计算的演进。低精度计算已成为高效大规模模型训练的关键技术。早期的FP16和BF16混合精度方案在保持收敛性的同时显著降低了训练成本【11,Mixed precision training,2018,International Conference on Learning Representations (ICLR)】【18,Training deep neural networks with 8-bit floating point numbers,2018,Advances in neural information processing systems】。近年来,FP8格式因其能在现代加速器上将算术吞吐量翻倍而备受关注【13,Nvidia hopper architecture whitepaper,2022,Technical Whitepaper】【12,FP8 Formats for Deep Learning,2022,arXiv e-prints】【3,Deepseek-v3 technical report,2024,arXiv preprint arXiv:2412.19437】。FP8通常遵循IEEE 754约定,但定义了多种变体,最著名的是E4M3和E5M2,它们在精度和动态范围之间做出了不同的权衡【18,Training deep neural networks with 8-bit floating point numbers,2018,Advances in neural information processing systems】。E4M3提供更高的精度,而E5M2提供更宽的数值范围。另一种变体UE8M0用于编码2的幂,通常用于缩放因子而非激活值。

FP8量化的缩放策略。将高精度张量(如FP32或BF16)转换为FP8需要量化。缩放策略定义了量化方案,并直接影响稳定性和准确性。早期工作依赖于逐张量(per-tensor)缩放【11,Mixed precision training,2018,International Conference on Learning Representations (ICLR)】,当张量值跨越多个数量级时,这种方法往往无法充分利用FP8的范围。后来的研究提出了更细粒度的缩放,如逐块(per-block)或逐通道(per-channel)量化【16,Fp8-lm: Training fp8 large language models,2023,arXiv preprint arXiv:2310.18313】,这些技术在保持效率的同时,为小数值保留了准确性,并已成为近期FP8训练系统的标准。

2.2 FP8训练系统和MoE架构

软硬件对FP8的推动。FP8的采用得益于硬件和软件的共同进步。NVIDIA的Hopper架构引入了原生的FP8 Tensor Core支持【13,Nvidia hopper architecture whitepaper,2022,Technical Whitepaper】,而TransformerEngine为FP8量化提供了开源的算子级支持【14,Nvidia transformer engine documentation,2023,https://github.com/NVIDIA/ TransformerEngine】。早期的规模化研究,包括模拟的FP8 GPT-3训练【12,FP8 Formats for Deep Learning,2022,arXiv e-prints】和Inflection-2【8,Inflection-2: The next step up,2023】模型,证明了FP8可以保持与BF16相当的收敛性。DeepSeekV3【3,Deepseek-v3 technical report,2024,arXiv preprint arXiv:2412.19437】进一步在万亿参数规模上验证了FP8训练,并发布了两个关键库——用于分组GEMM的DeepGEMM【5,Deepgemm: clean and efficient fp8 gemm kernels with fine-grained scaling,2025】和用于跨节点通信的DeepEP【4,Deepep: an efficient expert-parallel communication library,2025】,为混合专家(MoE)模型建立了一个强大的FP8基线。

现有FP8实现的局限。然而,这些实现仍然遵循以BF16为主的数据流,需要在GEMM边界进行频繁的量化-反量化(Q/DQ)转换【20,Insights into deepseek-v3: Scaling challenges and reflections on hardware for ai architectures,2025】。这种冗余的类型转换限制了FP8的实际性能并增加了内存流量。反之,移除这些转换会引入双重量化误差,因为沿不同维度(行式与列式)量化的张量会累积不一致的缩放因子【6,Gpt3.int8(): 8-bit matrix multiplication for transformers at scale,2022,Advances in Neural Information Processing Systems】。这种在计算效率和数值稳定性之间的权衡是当前FP8系统面临的主要挑战。

2.3 超大规模MoE训练中的系统级挑战

LLM与MoE的规模化挑战。大型语言模型(LLM)的快速扩展【1,Language models are few-shot learners,2020,Advances in neural information processing systems】【10,Llama 2: Open Foundation and Fine-Tuned Chat Models,2023,arXiv e-prints】已将训练基础设施推向极限。张量、流水线和数据并行等并行化策略【17,DeepSpeedMoE: Advancing mixture-of-experts inference and training to power next-generation AI scale,2022,Proceedings of the 39th International Conference on Machine Learning】,以及混合专家(MoE)架构【9,{GS}hard: Scaling giant models with conditional computation and automatic sharding,2021,International Conference on Learning Representations】【7,Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity,2022,Journal of Machine Learning Research】,已成为万亿参数训练的必要条件。MoE通过将每个token路由到一小部分专家来减少活跃计算,从而在不按比例增加成本的情况下实现高模型容量。然而,其稀疏和不规则的计算模式引入了新的瓶颈:all-to-all专家通信、激活值填充和动态路由开销【17,DeepSpeedMoE: Advancing mixture-of-experts inference and training to power next-generation AI scale,2022,Proceedings of the 39th International Conference on Machine Learning】【3,Deepseek-v3 technical report,2024,arXiv preprint arXiv:2412.19437】。将FP8精度集成到此类工作流中尤其具有挑战性,因为每个通信或专家边界通常需要额外的量化和数据重格式化。这些因素将计算分割成许多小算子,降低了算术强度,阻碍了FP8的理论优势在端到端MoE训练中得到充分实现。

FP8-Flow-MoE的动机。这些挑战促使我们设计了FP8-Flow-MoE——一个量化一致、无类型转换的FP8数据流,旨在消除冗余的Q/DQ操作,并在整个MoE流水线中保持数值稳定性。

A2 方法细节

MoE模型中FP8量化的研究背景。已有大量研究将FP8量化应用于基于Transformer的模型训练,特别是在BERT、GPT等密集模型上,取得了显著的性能提升且精度损失可忽略。本工作的重点是另一类模型——混合专家(MoE)模型,它因其高容量和推理时仅激活一小部分专家的低资源需求而成为性能最强的LLM之一。

3.1 感知缩放的转置

MoE模块中FP8量化的挑战。在MoE模块中完全采用FP8精度的关键挑战之一是分组线性运算需要两种不同的量化布局。具体而言,激活值(及其梯度)以行式(row-wise)量化格式被消费(在图2中称为Fprop和Dgrad),而权重梯度计算则需要列式(column-wise)量化的输入(Wgrad)。然而,之前的all-to-all通信(dispatch)通常只传输一种格式,通常是行式量化数据,以利用低精度格式带来的通信减少优势。

双重量化误差的产生。因此,当分组线性层随后需要列式量化数据时,一个朴素的实现必须首先进行“反量化→转置→重新量化”的操作,这引入了两轮(反)量化。这个过程导致了我们所称的“双重量化误差”(double quantization error)。其根本原因是,一维的逐token量化使用的缩放因子是基于连续数据段(例如每128个元素一个缩放因子)计算的,而这两种布局的缩放因子是不同的。由于FP8量化将缩放后的值映射到由指数和尾数位决定的离散可表示网格上,使用不同缩放因子进行两次独立的量化,实际上是将值两次重新映射到不重叠的离散集合中。形式上,双重量化误差定义为:

$$E = Q_{col}(D(Q_{row}(X))) - Q_{col}(X)$$

其中,$Q_{row}(·)$ 和 $Q_{col}(·)$ 分别是行式和列式量化算子,$D(·)$ 是反量化算子。行式和列式量化都是以逐瓦片(per-tile)的方式进行的,每个瓦片包含128个连续元素。换句话说,缩放因子的计算方式为:

$$s = \frac{\max_{0 \leq i < 128} \|x_i\|}{448}$$

其中448是FP8 E4M3格式可表示的最大数值。缩放后,量化算子实质上是将 $x_i$ 映射到最近的离散值:

$$Q_{row}(x_i) = \text{round}(\frac{x_i}{s})$$

反量化算子通过乘以缩放因子将FP8数据恢复到其原始范围,但舍入误差仍然存在:

$$D(Q_{row}(x_i)) = \text{round}(\frac{x_i}{s}) \cdot s$$

同布局量化不引入额外误差。根据这个表示,很明显,再次进行行式量化不会引入任何额外误差,因为相同的1x128元素再次被用于分块量化,这会得到相同的最大值,最终保持s不变。一个数值会被严格地映射到相同的离散值,形式上表示为:

$$Q_{row}(D(Q_{row}(x_i))) = \text{round}(\frac{\text{round}(\frac{x_i}{s}) \cdot s}{s}) \cdot s$$ $$= \text{round}(\text{round}(\frac{x_i}{s})) \cdot s$$
公式
公式

$=Q_{row}(x_i)$

这里我们假设使用确定性的舍入算法,例如四舍五入到最近值(RtN),并且对同一个数进行两次舍入不会改变其值。

不同布局量化导致误差放大。然而,如果第二次量化是按列进行的,其缩放因子 $s'$ 将根据转置后的128元素块计算得出,这通常与s不同。我们有:

$$Q_{col}(D(Q_{row}(x_i))) = \text{round}(\frac{\text{round}(x_i/s) \cdot s}{s'}) \cdot s'$$

这里的两个舍入算子不能合并,因为它们作用于不同的值。然而,如果将缩放因子限制为2的幂,这种误差放大可以得到缓解。在这种情况下,只要不发生上溢或下溢,两个量化域(行式和列式)之间的重新缩放仅涉及调整指数位。形式上,让我们假设:

$$s=2^T, s'=2^{T'}, \text{ where } T \in \mathbb{N}, T' \in \mathbb{N}.$$

基于2的幂次缩放因子的误差消除。在行式量化时,一个元素的值被离散化为:

$$Q_{row}(x) = (-1)^{SN} \cdot 2^{E-7} \cdot (1 + M/8) \cdot 2^T$$

其中SN、E和M分别是其FP8编码的符号位、指数位和尾数位,7和8是由FP8 E4M3定义得出的常数。当应用列式量化时,相同的值被离散化为不同的表示:

$$Q_{col}(x) = (-1)^{SN'} \cdot 2^{E'-7} \cdot (1 + M'/8) \cdot 2^{T'}$$

我们有:

由于T和T'都是自然数,存在另一个自然数D,使得 $D = T' - T$。注意到如果我们设置:

$$\begin{aligned} \begin{aligned} SN' &= SN \\ E' &= E - D \\ M' &= M \end{aligned} \end{aligned}$$

,我们有

$$Q_{col}(x) = (-1)^{SN} \cdot 2^{E-D-7} \cdot (1 + M/8) \cdot 2^{D+T}$$ $$=(-1)^{SN} \cdot 2^{E-D-7+D+T} \cdot(1+M / 8)$$ $$=(-1)^{SN} \cdot 2^{E-7} \cdot (1 + M/8) \cdot 2^{T}$$ $$=(-1)^{SN} \cdot 2^{E-7+T} \cdot (1+M/8)$$

$= Q_{row}(x)$

直接转置算子(Direct Transpose)的设计。基于这一洞察,我们进一步将一个转置块内的所有缩放因子对齐到最大的一个(以避免溢出),并仅通过指数操作直接在不同布局之间进行转换。换句话说,我们可以通过简单地修改FP8编码的指数位,将行式量化的FP8张量转换为列式量化的张量,而无需任何反量化或重新量化。这一原则是我们直接转置算子的基础,如算法1所示。

算法1 感知缩放的FP8转置
算法1的伪代码

直接转置算子的性能验证。为了验证该算子的效率优势,我们在不同输入维度上进行了数值实验。对于每种情况,我们比较了两种获取列式量化数据的策略:
(1) 朴素转换:反量化 → 转置 → 量化的方法。
(2) 直接转置(本文方法):不进行(反)量化,仅进行指数位操作。

延迟结果如图1所示,很明显,我们的直接转置算子在所有张量形状上都实现了2到3倍的加速,证实了其在计算效率上的有效性。

图1:不同量化策略的延迟比较。
图1:不同量化策略的延迟比较。

3.2 无类型转换的FP8数据流

MoE架构的数据流。MoE架构已成为一种在不按比例增加计算足迹的情况下扩展模型容量的有效方法。在每个MoE层内,一个路由器动态地将输入token分配给一小部分专家,每个专家通常由一个两层前馈网络组成。从数据流的角度来看,单个MoE层的计算可分解为以下阶段:路由(routing)、分发(dispatch)、置换(permutation)、专家计算(expert computation)、逆置换(unpermutation)和组合(combination)。专家计算本身通常包括由非线性激活(如SwiGLU)分隔的两个分组线性投影。在规模化实现中,这些投影作为分组GEMM操作执行,以利用跨多个专家的硬件效率。

图2:一个MoE模块的计算图
图2:一个MoE模块的计算图

传统BF16训练的数据流。在传统的基于BF16的训练中,端到端数据流涉及跨这些阶段的多种数据格式转换。Token首先通过all-to-all通信原语进行路由和分发,然后进行置换和打包,形成连续的专家批次,以进行高效的GEMM计算。每个专家执行第一个分组线性变换,接着是激活、第二个分组线性变换,最后是逆置换和组合步骤以聚合输出。这个结构定义了MoE层计算图的核心。

现有FP8集成策略的局限性。当将此流程扩展到FP8计算时,现有框架主要采用保守的集成策略。例如,NVIDIA的Transformer Engine实现了一个分块(blockwise)的FP8方案,其中FP8的量化和计算严格限制在分组线性模块内部。这种设计允许在矩阵乘法(分组GEMM)中利用FP8加速,实现了比BF16更高的算术吞吐量,但将所有通信和数据移动阶段——特别是分发all-to-all操作——保留在更高精度的格式中。因此,虽然GEMM算子受益于FP8张量核,但整个MoE层的性能仍然受到高精度通信和频繁格式转换的限制。

DeepSeek-V3 FP8方案的分析。另一个代表性的实现是DeepSeek-V3训练系统,它展示了成功的规模化FP8 MoE训练并取得了最先进的性能。尽管完整的训练框架未完全开源,但其公开的描述和发布的库——特别是DeepGEMM和DeepEP——表明专家计算和专家间通信都使用了定制的高性能FP8算子进行加速。然而,整个数据流仍然存在大量的量化开销。我们的分析表明,这种方法在单个MoE前向-反向传播过程中引入了多达十二次量化/反量化操作。此外,在权重梯度计算(wgrad)阶段,FP8激活值通常是行式量化、反量化,然后再次以列式方式量化。由于这些是逐token的量化方案,其缩放沿特定方向进行,重复的量化步骤会放大舍入和缩放误差(“双重量化误差”),对数值稳定性和收敛性产生负面影响。

FP8-Flow-MoE数据流设计。为了在保持数值稳定性的同时充分利用FP8的计算优势,我们设计了一个以FP8为中心的数据流,称为FP8-Flow-MoE,它最小化了格式转换,并尽可能在整个MoE专家路径中保持FP8表示。在此设计中,输入激活和输出梯度在整个前向和反向传播过程中都保持为FP8格式,除了两个特定的边界:
1. 第一个分组线性层的输出和激活函数之间。
2. 第二个分组线性层和反向传播中的组合/分发操作之间。

保留BF16的考量。这两个例外对应于FP8算术的基本数值挑战。第一个涉及归约操作(如求和),在有限的指数范围内极易发生溢出;第二个涉及非线性变换(如GELU或ReLU),可能会放大微小的量化误差。由于这两个计算在计算图中是相邻的,我们局部保留BF16精度以确保稳定性,同时将所有其他张量严格保持在FP8格式。这种设计产生了一个近乎连续的FP8数据流,并与先前的MoE中FP8实现相比,大幅减少了量化边界。

3.3 高性能FP8算子

高性能FP8算子的必要性。为了使所提出的以FP8为中心的计算流实用且高效,我们进一步开发了一套全面的支持算子,这些算子专门针对MoE层内的数据依赖和数值模式进行了定制。这些算子解决了在大规模MoE训练中应用FP8时出现的关键挑战,包括累积的舍入和缩放误差、冗余的数据转换,以及由碎片化算子引起的调度效率低下问题。它们共同实现了FP8-Flow-MoE的无缝和高吞吐量执行,确保了FP8张量核的优势从孤立的算子扩展到完整的端到端训练流水线。

算子性能评测方法。为了评估这些FP8算子的有效性,我们进行了一系列单元级基准测试,以隔离每个算子的性能和数值行为。测试输入通过模拟生成,其张量形状被选择以反映代表性MoE模型——DeepSeek V2-Lite、DeepSeek V2和DeepSeek V3——在不同专家并行(EP)设置下的真实世界配置。具体来说,这些微基准测试使我们能够直接测量原始算子吞吐量,展示每个算子如何为FP8-Flow-MoE系统的整体效率做出贡献。

3.3.1 融合的置换和填充 (Fused Permute and Padding)

背景与动机。在FP8数据流水线中,引入了一个填充(padding)步骤,以满足FP8 GEMM算子的形状约束——每个专家的输入必须包含16的倍数个条目,以充分利用张量核。在此之前,一个置换(permute)操作会重组分发后的token,使得属于同一专家的样本在内存中是连续的。分别执行这两个轻量级的数据移动算子会产生冗余的HBM访问和算子启动开销。由于这两个操作都是逐元素的且相互独立,它们可以自然地融合成一个单一的算子。

融合实现。基于这些观察,我们设计了一个融合的置换+填充算子,它在一次传递中直接执行专家级的重排序和形状对齐。这种融合是通过一个线程块映射方案实现的,该方案在从全局内存流式传输重排序的输入元素时,动态计算填充后布局中的目标偏移量。类似地,在反向传播阶段,我们应用相同的融合策略将逆置换和逆填充操作组合成一个单一的算子。这种对称的融合确保了前向和反向流之间的数据处理一致性。

性能对比。图4(原文为Figure 4,但描述内容对应Figure 3的标题)和图3(原文为Figure 3,但描述内容对应Figure 4的标题)报告了在BF16和FP8精度设置下,融合与非融合实现的性能比较。在对应典型MoE工作负载的不同张量大小上,融合版本展现出一致的加速效果。对于前向置换+填充,我们观察到高达1.7倍的加速;而对于反向逆置换+逆填充算子,在大规模配置下,性能提升高达6.6倍。

图3:融合与分离实现的置换和填充算子性能比较。
图3:融合与分离实现的置换和填充算子性能比较。
图4:融合与分离实现的逆置换和逆填充算子性能比较。
图4:融合与分离实现的逆置换和逆填充算子性能比较。

3.3.2 融合的SwiGLU和量化 (Fused SwiGLU and Quantization)

性能瓶颈分析。设计一个以FP8为中心的训练流程的动机源于一个关键的性能观察:将FP8算子(如高性能GEMM和通信)简单地插入到BF16计算图中,会引入频繁的量化和反量化操作,这可能极大地抵消FP8加速带来的好处。

通信中的量化开销。为了说明这个问题,我们使用DeepSeek发布的DeepEP库分析了FP8通信。在其默认用法中,量化和反量化紧接在通信之前和之后执行。表1报告了在三个代表性数据规模(对应DeepSeek V2-Lite、V2和V3)和三个专家并行(EP)度(8、16、32)下测得的延迟。结果揭示了两个关键发现:对于较小规模的工作负载,量化和反量化各自消耗的时间与通信算子相似,几乎抵消了任何性能增益;虽然通信成本随工作负载大致线性增长,但(反)量化成本保持不变,然而,这种数据转换导致端到端FP8加速从1.6倍降至1.4倍。

量化开销对性能的影响。值得注意的是,尽管FP8通信理论上将数据传输大小减半,但在实践中,需要同时传输FP8张量及其对应的缩放因子,这使得数据缓冲和同步的数量加倍,将可实现的加速限制在1.6倍左右。然而,仅仅一对量化和反量化算子就将FP8通信的性能增益减少了大约三分之一(从0.6倍的提升降至0.4倍)。考虑到一个典型的MoE前向或反向传播过程涉及大约三对这样的操作(见图2(c)),FP8算子的好处几乎可以被完全抵消。

融合SwiGLU与量化的解决方案。为了缓解这个问题,FP8-Flow-MoE尽可能地将量化直接集成到周围的计算算子中。特别是,我们将SwiGLU激活和量化步骤融合成一个单一的算子。图5展示了我们在几个代表性张量形状上的单元测试结果。融合实现达到了与NVIDIA Megatron框架中原始独立SwiGLU算子几乎相同的延迟,同时无缝地为后续计算生成FP8量化的输出。这表明FP8-Flow-MoE可以在保留混合精度数值效率的同时,消除先前限制端到端性能的量化瓶颈。

表1:通信性能与加速比
表1 通信性能与加速比

图5:融合的SwiGLU和量化算子的性能。
图5:融合的SwiGLU和量化算子的性能。

A4 实验

我们从数值收敛性训练效率两个角度评估FP8-Flow-MoE。所有实验均在一个32节点的NVIDIA Hopper集群上进行,每个节点配备8个GPU(每GPU 80 GB内存)。专家并行(EP)跨节点扩展,以评估在真实大模型工作负载下的训练效率和可扩展性。

4.1 收敛性验证

实验设置: 为了验证FP8-Flow-MoE能保持数值稳定性和收敛行为,我们使用两种精度设置从头开始训练一个DeepSeek-V2-Lite模型(160亿参数):BF16和FP8-Flow-MoE。两个模型都训练了2000亿个token,使用完全相同的优化超参数、学习率调度、数据顺序和并行配置。

实验结果: 如图6所示,两条损失曲线在整个训练轨迹中几乎无法区分——从早期的预热阶段到损失稳定阶段。FP8-Flow-MoE模型没有出现发散、梯度下溢或数值漂移的迹象,表明即使在扩展规模的工作负载下,我们的以FP8为中心的数据流也不会损害收敛保真度。这一观察证实了FP8-Flow-MoE中的感知缩放的转置和融合算子保留了BF16训练的数值动态。

图6:训练2000亿token的损失比较
图6:训练2000亿token的损失比较

4.2 效率评估

实验配置: 我们接下来对多种精度和数据流配置的端到端训练效率进行基准测试。具体来说,我们比较了以下三种代表性配置:
1. BF16: 使用标准FP32-BF16混合精度的基线配置。
2. Blockwise: 采用FP8分组GEMM算子的TE风格分块式方案。
3. FP8-Flow-MoE (ours): 具有融合算子和最小化量化开销的以FP8为中心的数据流。

所有实验都在DeepSeek-V3 (671B)模型上进行,以验证FP8-Flow-MoE在真实大规模训练条件下的可扩展性和鲁棒性。模型使用标准的专家和流水线并行配置(EP/PP = 8/32, 16/16, 32/8)进行训练,并在两种激活检查点(AC)策略下进行评估:AC = full,对所有模块应用完整检查点;以及AC = sel (+MoE expert),选择性地对MoE层进行检查点,但排除专家以评估FP8激活压缩。后者是与Megatron中1F1B重叠调度兼容的最节省内存的配置。为隔离计算效率与通信重叠的影响,所有端到端实验均在没有计算-通信重叠的情况下进行。

吞吐量分析: 表2和表3总结了在三种EP设置下测得的吞吐量(TGS,tokens/GPU/s)和峰值内存(Mem,GB)。在所有激活检查点模式下,FP8-Flow-MoE始终优于BF16基线,实现了+6% (EP8)、+8% (EP16)和+16% (EP32)的吞吐量提升。与Blockwise方案相比,FP8-Flow-MoE进一步将吞吐量提高了3% (EP8)、8% (EP16)和高达21% (EP32)。性能差距在更高的EP级别上扩大,这是因为减少的量化开销和算子融合的累积影响变得更加显著。

内存效率: 在AC = sel (+MoE expert)和EP8的设置下,FP8-Flow-MoE的峰值内存相比BF16减少了约8GB,相比Blockwise减少了16.5GB,这直接得益于FP8检查点压缩机制,该机制将中间激活以FP8而非BF16存储。在最大的EP=32规模下,BF16和Blockwise基线都遇到了内存不足(OOM)错误,而FP8-Flow-MoE保持稳定,清晰地展示了其卓越的可扩展性和内存效率。

表2:AC=full下的吞吐量分析
表2 AC=full下的吞吐量分析

表3:AC=sel (+MoE expert)下的吞吐量分析
表3 AC=sel (+MoE expert)下的吞吐量分析

4.3 讨论与洞见

实验结果共同突显了三个关键发现:

  1. 持久的FP8数据流至关重要。仅仅用FP8替代GEMM或通信算子带来的收益有限,因为在小批量或高通信频率下,量化/反量化成本占据主导。所提出的感知缩放的转置和融合算子设计消除了冗余的精度转换,直接提高了算子启动效率和内存带宽利用率。

  2. 规模化放大了FP8-Flow-MoE的优势。随着EP的增加,基线方法中的通信和量化开销会复合增长,而FP8-Flow-MoE的低精度持久性最小化了这些成本,从而在大专家并行度下实现了高达1.2倍的端到端加速。

  3. 实际训练的稳定性。尽管FP8-Flow-MoE在训练流水线中比基线实现更广泛地采用了较低精度,但它在数百亿token的训练中保持了相同的收敛行为和鲁棒性。这种稳定性源于我们对转置边界处潜在双重量化误差的洞察。通过引入感知缩放的转置,FP8-Flow-MoE有效消除了这种冗余量化,同时简化了数据移动,从而同时提高了数值保真度和效率。

A5 结论

本研究确定了在大规模混合专家(MoE)训练中实现FP8高精度端到端效率的另外两个关键因素。

首先,在整个计算流程中保持低精度数据可以最小化冗余转换并减轻内存带宽压力。一个代表性的例子是我们的FP8转置算子,它直接将行式量化张量转换为列式量化张量,而无需恢复到更高精度——从而在避免双重量化误差的同时实现高数据吞吐量。

其次,FP8精度不可避免地通过引入量化和填充等小算子来碎片化计算图,导致过多的算子启动和内存流量。通过融合轻量级操作和开发定制的高性能实现,FP8-Flow-MoE有效消除了这种结构上的低效。

总之,FP8-Flow-MoE为MoE模型建立了首个以FP8为中心的训练范式,在保持与BF16相当的收敛性的同时,与BF16工作流和朴素的FP8算子替换相比,实现了高达21%的吞吐量提升和每GPU 16.5 GB的峰值内存降低。这些结果证明了完全低精度MoE训练的可行性,并为可扩展、高能效的基础模型开发开辟了新的机遇。