End-to-End MXFP8 Reinforcement Learning on Blackwell: What Matching Buys, When It Pays Off, and How It Fails
作者: Wengee Zhang
分类: cs.LG
发布日期: 2026-10-26
备注: 36 pages, 8 figures, 13 tables
💡 一句话要点
提出MXFP8强化学习以优化大语言模型训练效率
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 强化学习 低精度算术 大语言模型 训练效率 模型一致性 MoE模型 KL散度 权重同步
📋 核心要点
- 现有的强化学习方法在大语言模型训练中面临高成本和一致性问题,尤其是在低精度算术的应用上。
- 本文提出了一种端到端的MXFP8强化学习框架,通过优化采样器与学习者之间的匹配来提高训练效率。
- 实验结果表明,MXFP8在MoE模型中显著降低了训练和生成时间,同时在特定条件下提高了模型的一致性和效率。
📝 摘要(中文)
低精度算术可以降低大语言模型的强化学习(RL)成本,但要求采样器与学习者之间达成一致。本文实现了端到端的MXFP8 RL,使用B200 GPU对Qwen3-8B和Qwen3-30B-A3B进行测量,分析了该设计的优势、速度提升及其失败情况。研究发现,MXFP8学习者的采样器在学习者的网格上更为一致,训练回放的KL散度显著降低。MXFP8训练在密集设置中仅在每个微批次超过4.5k个token时才有经济效益,且在MoE模型中,MXFP8减少了10.7%的步骤时间和24%的生成时间。最后,论文还提出了七个权重同步必须保持的条件,并通过实验验证了这些条件的有效性。
🔬 方法详解
问题定义:本文旨在解决大语言模型强化学习中的高成本和低一致性问题,尤其是在低精度算术应用时的挑战。现有方法在采样器与学习者之间缺乏有效匹配,导致训练效率低下。
核心思路:论文提出的MXFP8强化学习框架通过在学习者的网格上优化采样器,确保采样器与学习者之间的一致性,从而提升训练效率和模型性能。
技术框架:整体架构包括MXFP8学习者和采样器,采用Megatron/Transformer Engine进行训练,结合vLLM和SGLang进行回放,支持密集模型和MoE模型的训练。
关键创新:最重要的技术创新在于通过使用MXFP8格式的采样器与学习者的网格匹配,显著降低了训练过程中的KL散度,并提升了训练的一致性。与传统的高精度采样器相比,MXFP8在特定条件下表现出更好的性能。
关键设计:论文中设定了七个权重同步必须保持的条件,并通过实验验证了这些条件的有效性。关键参数设置包括微批次大小和token数量,损失函数设计确保了训练过程的稳定性。
📊 实验亮点
实验结果显示,MXFP8在MoE模型中将步骤时间减少了10.7%,生成时间减少了24%。在特定条件下,MXFP8的训练效率显著高于传统的BF16方法,尤其是在每个微批次超过4.5k个token时,表现出更优的经济效益。
🎯 应用场景
该研究的潜在应用领域包括大语言模型的训练和优化,尤其是在资源受限的环境中。通过降低训练成本和提高效率,MXFP8强化学习框架可广泛应用于自然语言处理、对话系统和智能助手等领域,未来可能推动更多低精度计算技术的应用。
📄 摘要(原文)
Low-precision arithmetic can make reinforcement learning (RL) for large language models cheaper, but RL requires the sampler and the learner to agree. Several frameworks now run FP8 on both sides. We implement end-to-end MXFP8 RL in verl (Megatron/Transformer Engine training; vLLM and SGLang rollout; dense and MoE models) and measure on B200 GPUs, with Qwen3-8B and Qwen3-30B-A3B, what this design buys, when it is faster, and how it fails. (1) Matching. For an MXFP8 learner, a sampler on the learner's grid is more consistent than a higher-precision BF16 sampler: in one engine at fixed weights, train-rollout KL is 11% lower with half of the sampler's linear layers on the grid and 16% lower with all of them, and the synchronized weights we checked equal the learner's FP8 forward weights byte for byte. On the MoE model, an FP8 pipeline whose sampler uses another grid has twice the shared-grid pipeline's KL from the first step. Without a quantized KV cache, BF16 end to end remains 2.5-10x more consistent. (2) Speed. MXFP8 training adds a fixed cost per micro-batch and saves a per-token cost, so in our dense setting it pays off only above about 4.5k tokens per micro-batch; a micro-batch sweep on identical data confirms the per-micro-batch cost, and a trace shows the GPU waiting for the host for most of the added time. End to end, MXFP8 cuts MoE step time by 10.7% and generation time by 24%, and slows a short-sequence dense workload by 5%. (3) Correctness. We state seven conditions every weight synchronization must preserve and checks that test them against the live engine. Each of the seven failures we met on hardware broke one condition, five silently; the checks caught two at their first opportunity, each missed by the other checks, and CPU fault-injection tests reproduce five. All runs are single-seed and at most 60 steps long; we make no claim about final model quality.