隨著大語言模型從簡單的文本生成轉向複雜推理，強化學習扮演著越來越重要的角色。以 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