FP8-Flow-MoE: Efficient FP8 Training for MoE Models

• FP8-Flow-MoE is an FP8-centric pipeline that eliminates redundant quantize-dequantize casts to enhance efficiency and numerical stability.
• It introduces scaling-aware transpose and fused FP8 operators, reducing kernel launches and memory traffic while boosting throughput by up to 21%.
• Designed for both TransformerEngine and Megatron-LM, it supports hundred-billion parameter models in supervised and reinforcement learning settings.

FP8-Flow-MoE is an FP8-centric training and deployment pipeline for large Mixture-of-Experts (MoE) models, specifically designed to eliminate the computational and memory overheads associated with traditional mixed-precision approaches. It achieves near-theoretical FP8 acceleration by removing redundant quantize-dequantize (Q/DQ) casts and introducing quantization-consistent dataflows, a scaling-aware transpose operator, and kernel-fused FP8 operators. FP8-Flow-MoE is implemented for both supervised and reinforcement learning (RL) settings at the hundred-billion parameter scale and is compatible with TransformerEngine and Megatron-LM (Wang et al., 4 Nov 2025, Zeng et al., 8 Dec 2025).

1. Motivation and Problems with Prior Art

Modern MoE training pipelines employ mixed-precision, typically based on BF16 or FP16 for computations, and rely on frequent Q/DQ between operator and communication boundaries. Substituting GEMMs with FP8 arithmetic does not yield maximum throughput because each operator boundary still incurs a three-step conversion: dequantize FP8→BF16, run the BF16 op, and quantize BF16→FP8. In an MoE forward-backward pass, this pattern leads to about 12 Q/DQ steps per example, each incurring a kernel launch, device memory traffic, and dynamic scaling computation. These overheads negate much of FP8's 2× theoretical arithmetic gain.

Attempting to eliminate all intermediate casts by running fully in FP8 leads to double quantization error. Since per-block scaling can differ along orthogonal tensor dimensions (e.g., row and column tiles), sequential quantization steps introduce inconsistency, yielding unbounded amplification of rounding error and unstable gradients. The error for input $X$ is

$$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$

where $s_{row}$ and $s_{col}$ are generally not equal, so roundings do not commute.

2. FP8-Flow-MoE Dataflow and Core Operators

FP8-Flow-MoE resolves both redundant casting and double quantization by prescribing an FP8-centric dataflow with only two Q/DQ steps: at MoE entry (input quantize) and MoE exit (output dequantize). All intermediate data (activations, dispatches, intermediate expert results) remains in FP8. Only two BF16 boundaries are retained: before the first GEMM accumulation (to prevent exponent overflow) and after the output of the last GEMM (for final combination/all-to-all).

Scaling-Aware Transpose

To preserve scale consistency across row-wise and column-wise quantized tiles, a direct exponent-only scaling-aware transpose is introduced. For a block $X_{row}$ in FP8 with per-row scales $S^{row}$, the matrix is transposed so that the per-column scales of the output $X_{col}$ are set to $S_{max} = \max_{i} S^{row}_i$. Elements' exponents are adjusted by

$$e' = e - \Delta_{j,i}, \quad \Delta_{j,i} = \log_2 \frac{S^{col}_j}{S^{row}_i},$$

ensuring scale consistency without extra rounding or loss of precision within a power-of-two rescaling.

Fused FP8 Operators

FP8-Flow-MoE fuses lightweight activation and data movement operations with quantization, reducing kernel launch counts and memory movements:

• Fused Permute + Padding / Unpermute + Unpadding: Arranges FP8 dispatch arrays into GEMM-aligned blocks in memory.
• Fused Swiglu + Quantization: Computes $C = A \odot \mathrm{GELU}(B)$ in BF16, then quantizes $$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$0 to FP8 (per 128-row tile) in the same kernel, producing FP8 activations and per-row scales.
• Fused GEMM + Accumulate (+Quantization): Multiplies FP8 row-scaled input $$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$1 and col-scaled $$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$2 (with appropriate scales), accumulates in BF16, and optionally quantizes the result to FP8.

3. Quantization Scheme and Numerical Properties

FP8-Flow-MoE employs symmetric, per-tile FP8_E4M3 quantization with a standard tile size of 128:

• Quantization: For $$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$3, $$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$4 (where $$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$5 for E4M3), and $$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$6.
• Dequantization: Either on-the-fly inside GEMM accumulation or by $$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$7.

By aligning all scales within each transpose block to the same power-of-two $$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$8, one ensures row- and column-quantized encodings are identical after a simple exponent shift. This removes double quantization error, bounding per-element error to $$E_i = Q_{col}(D(Q_{row}(x_i))) - Q_{row}(x_i)s_{row},$$9, matching the single-pass quantization error of FP8.

Empirically, this approach shows convergence curves that are indistinguishable from BF16 mixed-precision training on large-scale models across 200B tokens.

4. System Integration and Adoption

FP8-Flow-MoE is implemented as a plug-and-play extension for both TransformerEngine and Megatron-LM:

• TransformerEngine: Provides a custom operator library (fp8_flow_moe.so), adds a new DataType.FP8_FLOW enum, and configures grouped linear layers with row/column tensor shapes and per-tile metadata. Usage involves specifying the FP8 config and substituting standard grouped linear layers with FP8-Flow-aware variants.
• Megatron-LM: Replaces linear operator calls with FP8-Flow-MoE-aware versions, inserts scaling-aware transpose in expert dispatch, and adjusts configuration YAML to specify "precision: fp8_flow", with custom block size and scaling settings.

Adoption requires replacing grouped linear layers, inserting scaling-aware transposes, enabling fused permutations, and ensuring only the two precise BF16 boundaries are preserved at the expert stack entrance/exit. Sanity-checks for NaN/Inf and small problem settings are recommended before scaling up.

5. Experimental Evaluation

FP8-Flow-MoE was benchmarked on a 671B-parameter DeepSeek-V3 MoE with 32 nodes × 8 H100 GPUs. Parallelism degrees varied (EP=8/16/32, PP=32/16/8), using AdamW and large batch sizes. Results demonstrate consistent improvements:

Method	EP=8 (TGS/Mem)	EP=16 (TGS/Mem)	EP=32 (TGS/Mem)
BF16	1,109 / 39	939 / 36	671 / 43
TransformerEngine	1,146 / 37	938 / 41	644 / 51
FP8-Flow-MoE	1,176 / 37	1,012 / 39	779 / 49

For activation checkpointing over MoE experts only (AC=sel), FP8-Flow-MoE achieves up to 8 GB/GPU memory savings and enables EP32 training instances previously out-of-memory (OOM) under BF16/blockwise. Speedup over TransformerEngine blockwise implementations reaches up to 21% at maximal parallelism. Training loss curves for FP8-Flow-MoE and BF16 are nearly identical.

6. FP8-Flow-MoE in Reinforcement Learning Pipelines

Within RL systems for hundred-billion-scale MoE models, FP8-Flow-MoE is utilized as the quantization backend for both rollout and value estimation (Zeng et al., 8 Dec 2025). The core design includes:

• Length-aware Scheduling: Prompts are sorted into batches with balanced predicted sequence lengths, minimizing idle time.
• FP8 Rollout Kernels: All linear and MoE blocks use FP8-quantized weights (E4M3 or E5M2 as per layer sensitivity), with activations dynamically quantized. Throughput increases by ~1.6× relative to FP16 rollout.
• Overlapped Reward Computation: Post-generation, trajectories are detokenized and scored asynchronously, hiding up to 85% of reward-stage latency.
• MoE Router-Replay: During RL actor updates, recorded routing decisions are replayed to ensure train/infer consistency and reduce discrepancy in output probabilities.
• Performance: End-to-end RL rollout iteration wall time is reduced by ~40%, with the stage time dropping from 61% to ~42.5% of the overall RL loop.

Gradients remain in FP16, and the gating MLP is retained in FP16 for router stability. FP8-Flow-MoE's quantization noise does not measurably affect downstream RL policy performance when FP16 log-prob recomputation is used for the updates.

7. Reproducibility and Best Practices

Reproduction of FP8-Flow-MoE requires:

• FP8_E4M3 format for activations/weights, UE8M0 only for shared scale factors.
• Block size 128 with dynamic, per-tile scaling.
• AdamW optimizer with standard hyperparameters and batch sizes.
• Only two BF16 boundaries (entry, exit for expert stack).
• Autopipelining on NVIDIA H100 for optimal throughput.
• Monitoring overflow/underflow counters in FP8 kernels.

A checklist (see original work) ensures proper integration, kernel activation, and numerically stable convergence. Scaling up parallelism should be accompanied by throughput and memory monitoring and comparison of loss curves versus BF16 for validation.

---

References:

FP8-Flow-MoE: "FP8-Flow-MoE: A Casting-Free FP8 Recipe without Double Quantization Error" (Wang et al., 4 Nov 2025) Utilization in RL: "Each Prompt Matters: Scaling Reinforcement Learning Without Wasting Rollouts on Hundred-Billion-Scale MoE" (Zeng et al., 8 Dec 2025)