随着大语言模型从简单的文本生成转向复杂推理，强化学习扮演着越来越重要的角色。以 GRPO（Group Relative Policy Optimization）为代表的算法推动推理模型通过迭代反馈持续改进。然而，强化学习训练循环分为两个截然不同且强度极高的阶段：生成阶段（具有严格的延迟要求）和训练阶段（需要高吞吐量）。
NVIDIA 在 NeMo RL 框架中实现了端到端的 FP8 精度方案，从生成到训练全过程使用 FP8，在保持模型精度的前提下显著提升吞吐量。相比 BF16 基线方案，FP8 端到端方案在 Llama 3.1 8B 模型的 GRPO 训练中实现了一致的 15-25% 吞吐量提升，且验证精度与 BF16 几乎持平（0.613 vs 0.616）。
FP8 强化学习面临的核心挑战之一是数值不一致性。由于 RL 流水线通常使用独立引擎——vLLM 负责生成、Megatron Core 负责训练——各引擎使用不同的 CUDA 内核，低精度下的量化与反量化逻辑会累积数值差异。研究表明，仅在生成阶段使用 FP8 而训练用 BF16 的方案无法完全弥合精度差距，而端到端 FP8 结合重要性采样技术则能彻底消除这一差距。
在此基础上，NVIDIA 进一步将 FP8 扩展至 KV 缓存和注意力计算。强化学习工作流中，长输出序列下的 KV 缓存增长和注意力计算往往成为瓶颈。NeMo RL 采用动态重标定方案：在每个训练步骤结束时，使用更新后的策略权重重新计算 QKV 缩放因子，并同步到推理引擎。该设计确保生成引擎始终使用源自最新策略状态的最优量化尺度，重标定开销仅占总步骤时间的 2-3%。
启用 KV 缓存和注意力的 FP8 后，生成阶段在 W8A8 线性层基础上再获约 30% 加速，相比 BF16 基线整体加速约 48%。这一增益在长响应长度下尤为显著，因为此时注意力计算占总工作负载的更大比例。
NeMo RL 的 FP8 方案支持稠密模型和混合专家（MoE）模型。在 Qwen3-30B MoE 模型上的实验显示，FP8 端到端方案可达到与 BF16 相同的精度曲线。开发者可通过简单的配置文件启用 FP8 线性层、KV 缓存和注意力，并可自由调节首尾层精度、缩放因子类型等高级参数。
WeChat
Profile