Jet-RL：通过统一训练与推理精度流实现 On-Policy FP8 强化学习

统一训练和推理的 FP8 精度流，解决 RL 滚动加速中的训练不稳定问题

前置知识

FP8 量化 (E4M3 格式)

FP8 是一种 8 位浮点数格式，其中 E4M3 表示 4 位指数 + 3 位尾数，最大可表示值 $\Delta_{max}=448$。与 BF16（16 位）相比，FP8 将数据宽度压缩一半，理论上可实现 2× 的计算吞吐提升。量化过程为 $\hat{X} = \lceil X / S_X \rceil$，其中缩放因子 $S_X = \max(|X|) / \Delta_{max}$。FP8 在 NVIDIA H100 的 Tensor Core 上有原生硬件支持，已被广泛用于 LLM 推理加速。

本文的核心就是用 FP8 加速 RL 训练中的滚动阶段，理解 FP8 的精度特性和硬件实现是理解全文的基础。

Reinforcement Learning from Human Feedback (RLHF) 与 PPO

RLHF 是用人类偏好信号训练语言模型的范式。PPO（Proximal Policy Optimization）是其中最常用的算法，包含四个模型：Actor（待训练的策略模型）、Reference（用于计算 KL 散度的冻结模型）、Reward（打分模型）和 Critic（价值估计模型）。每步训练分为 Rollout（Actor 自回归生成回复）、Evaluation（其他模型前向计算）和 Update（Actor 更新权重）三个阶段。

Jet-RL 的优化目标正是 PPO 训练流程中的 Rollout 瓶颈，理解 PPO 的多模型架构和三阶段流程是理解优化空间的前提。

On-policy 与 Off-policy RL

On-policy 要求训练数据来自当前策略的实时采样——即 Actor 生成回复时的模型权重，必须和更新梯度时的模型权重一致。Off-policy 则允许使用历史数据或其他策略生成的数据训练。在 LLM RL 训练中，如果推理引擎和训练框架使用不同精度的模型，就相当于引入了 off-policy 偏差，因为推理时的输出分布和训练时的梯度计算不一致。

本文的核心论点是：BF16 训练 + FP8 推理的混合精度方案本质上破坏了 on-policy 假设，这是训练不稳定的根源。理解这一概念是理解 Jet-RL 设计动机的关键。

GEMM 量化布局（Row-wise 与 Column-wise）

在量化线性层时，不同的矩阵乘法对输入矩阵的存储布局有不同要求。NVIDIA FP8 Tensor Core 硬件要求第一个操作数按行存储（Row-wise, R），第二个操作数按列存储（Column-wise, C）。对于线性层的三个 GEMM 操作：FProp（前向传播 $Y = X \times W^\top$）需要激活 Row、权重 Row；WGrad（权重梯度 $\nabla W = \nabla Y^\top \times X$）需要激活 Col、梯度 Col；DGrad（激活梯度 $\nabla X = \nabla Y \times W$）需要梯度 Col、权重 Col。

Jet-RL 的技术实现核心之一就是精心设计每个 GEMM 的量化粒度和存储布局，确保训练和推理共享一致的量化行为。

Per-group / Per-block 量化粒度

Per-tensor 量化对整个张量使用单一缩放因子，粒度最粗但容易受离群值影响导致精度损失。Per-group 量化（如 $1 \times 128$，即每 128 个连续元素一组）和 Per-block 量化（如 $128 \times 128$）提供更细粒度的缩放因子，能更好地保留局部数值范围的信息。DeepSeek-V3 等工作已验证细粒度量化对 LLM 训练稳定性的重要性。Jet-RL 对激活和梯度使用 $1 \times 128$ per-group 量化，对权重使用 $128 \times 128$ per-block 量化。

量化粒度的选择直接影响训练精度和计算效率的平衡，是 Jet-RL 实现既高效又稳定的关键技术决策。

研究动机

在大语言模型的强化学习（RL）训练中，Actor 模型的自回归滚动生成（Rollout）是最大的性能瓶颈。论文在 Qwen3-8B-Base 模型上进行了详细 profiling：当滚动长度超过 8K token 时，仅 Rollout 阶段就占总训练时间的 75% 以上；在 16K 滚动长度下，这一比例高达 81.2%。这种极端的时间不平衡严重拖累了端到端训练效率。一个自然的想法是用 FP8 量化来加速 Rollout 阶段，同时保留 BF16 精度用于训练阶段以确保稳定性。这种「BF16 训练 + FP8 推理」的混合精度策略已被 VeRL、SLIME、NeMo-RL、OpenRLHF 等主流 RL 框架广泛采用。然而，论文发现这个看似合理的方案存在两个严重的失效场景：一是当滚动长度较长（>8K token）时，精度不匹配累积导致训练崩溃；二是当目标任务较难或基础模型较弱时，训练曲线迅速发散。在 Llama3.1-8B 上，该方案的平均分数相比 BF16 基线下降了 10.2%；在 Qwen2.5-7B 上甚至完全无法收敛。

本文的目标是论文的目标是建立一个可靠、稳定的 FP8 RL 训练框架，在不牺牲模型精度的前提下，实现 Rollout 阶段和端到端训练的显著加速。具体而言，论文希望实现：（1）在所有模型和任务设置下都能稳定收敛；（2）与 BF16 基线的性能差距控制在 1% 左右；（3）Rollout 阶段达到 1.3× 以上的加速比；（4）端到端训练步时间实现 1.16× 以上的加速。这些目标需要在保持 on-policy 训练一致性的前提下实现。

与已有工作不同的是，已有工作（如 SLIME、NeMo-RL）声称直接将 BF16 权重 cast 到 FP8 不会影响精度，但论文发现这一结论极为脆弱——它仅在模型已有较强任务先验、且滚动长度较短时成立。论文的独特切入角度是：将 BF16-train + FP8-rollout 的失效原因定位为 on-policy 假设的破坏。具体来说，训练阶段使用 BF16 计算梯度，而推理阶段使用 FP8 生成回复，这两个精度流之间的数值不匹配在长序列生成中逐步累积，使得 Actor 实际学习的数据来自一个「不同策略」的分布，导致 RL 训练本质上变成了 off-policy 训练。这个视角将问题从「量化精度损失」重新框架化为「策略一致性破坏」，为解决方案的设计提供了清晰的方向。

核心方法

Jet-RL 的核心直觉可以用一个简单的类比理解：如果你用一个精度较低的「副驾驶」（FP8 推理引擎）来生成训练数据，却用一个精度较高的「主驾驶」（BF16 训练框架）来更新模型，那么主驾驶学到的经验就不是副驾驶实际执行时的经验。长此以往，两者之间的偏差会越来越大。Jet-RL 的解决方案是：让主驾驶和副驾驶使用完全相同的精度——统一使用 FP8 进行前向计算。具体技术路线是：在训练框架的前向传播中，也采用与推理引擎相同的 FP8 量化方案来计算激活值，确保训练时的前向计算图 $\mathcal{G}_{train}^{fwd}$ 和推理计算图 $\mathcal{G}_{infer}$ 具有完全一致的量化行为。权重的主副本仍保持 BF16 精度以稳定训练，但所有 GEMM 操作（FProp、WGrad、DGrad）都使用 FP8 执行，激活的保存也使用 FP8 精度。

Jet-RL 最本质的创新在于：不是在推理侧做更多校准来逼近训练精度，而是反其道而行之——把训练侧的前向计算也降到和推理相同的 FP8 精度。已有方法试图用 PTQ（Post-Training Quantization）校准来缩小训练-推理差距，但校准过程极其耗时（对 8B 模型就需要数十分钟），且必须在每步权重同步时重复执行，完全抵消了加速收益。Jet-RL 通过强制 $\mathcal{G}_{infer}$ 成为 $\mathcal{G}_{train}^{fwd}$ 的子图，从根本上消除了精度不匹配。这意味着训练和推理看到的模型是「同一个模型」，RL 训练恢复了 on-policy 属性。这一设计不需要任何 inter-step 校准，开销极低。同时，论文精心设计了每个 GEMM 的量化粒度：前向 FProp 使用 $(1 \times 128) \times (128 \times 128)$ 的 FP8 矩阵乘法，DGrad 复用相同配置，WGrad 使用 $(1 \times 128) \times (128 \times 1)$ 的配置，并且巧妙地融合了正反向传播中的量化过程以减少额外开销。

方法步骤详情

Jet-RL 的完整工作流程分为以下几个步骤。第一步，量化策略设计：对权重采用 $128 \times 128$ per-block 量化，对激活和梯度采用 $1 \times 128$ per-group 量化。所有量化使用 E4M3 FP8 格式，$\Delta_{max}=448$。第二步，前向传播统一：在训练的前向传播中，输入激活 $X$ 通过 $1 \times 128$ per-group 量化转为 FP8，权重 $W$ 通过 $128 \times 128$ per-block 量化转为 FP8，两者执行 FP8 GEMM 计算 $\hat{Y} = \hat{X} \times \hat{W}^\top$。这一过程与推理引擎（vLLM）中的量化行为完全一致。第三步，激活保存：前向传播中的激活以 FP8 精度保存用于反向传播，避免精度不一致。第四步，反向传播设计：梯度 $\nabla Y$ 保持 BF16 精度以稳定训练。DGrad GEMM 中，对激活（已保存为 FP8）需要从 $1 \times 128$ 重新量化为 $128 \times 1$ 布局；权重的量化方案沿对称性可在正反向间复用，仅需转置。WGrad GEMM 使用 $(1 \times 128) \times (128 \times 1)$ 的配置。第五步，硬件加速：使用 DeepGEMM 库的 FP8 Tensor Core kernel 执行所有 GEMM 操作，使用 Triton 实现自定义的量化、转置和融合 RMSNorm kernel。

技术新颖性

Jet-RL 的技术新颖性体现在多个层面。首先，在方法论层面，已有工作将 FP8 RL 加速视为「推理加速」问题，试图通过更好的 PTQ 校准或直接 cast 来解决精度损失。Jet-RL 则将其重新定义为「训练-推理一致性」问题，提出了统一精度流的范式转换。其次，在量化方案设计上，论文没有简单复用预训练或 SFT 阶段的 FP8 训练方案，而是针对 RL 训练的特殊性（Actor 权重频繁更新、需要在训练和推理框架间同步）进行了专门设计。特别是对反向传播中激活量化的重新量化策略——正向保存 $1 \times 128$，反向需要 $128 \times 1$——这一设计既满足了硬件 kernel 的布局要求，又通过量化噪声的随机性反而有益于训练稳定性。最后，在系统集成层面，Jet-RL 基于 vLLM + VeRL 构建，使用 DeepGEMM kernel，展示了如何在不修改现有 RL 框架架构的前提下，通过统一精度流实现端到端加速。

实验结果

论文在多个模型（Llama3.1-8B、Qwen2.5-7B、Qwen3-8B-Base）、多个数据集（GSM8K+MATH、DeepMATH）和多个滚动长度（8K、16K）设置下进行了全面评估。核心发现如下。在 8K 滚动长度的 GSM8K+MATH 设置中，BF16-train-FP8-rollout 方法在 Llama3.1-8B 上平均分数下降 10.2%（23.2% → 13.0%），在 Qwen2.5-7B 上完全无法收敛，在 Qwen3-8B-Base 上下降 2.9%。Jet-RL 在所有三个模型上都成功收敛，性能差距缩小到 1-2%：Llama3.1-8B 上甚至超出 BF16 基线 2.0%（25.2% vs 23.2%），Qwen2.5-7B 上仅下降 1.0%（55.9% vs 56.9%），Qwen3-8B-Base 上下降 1.1%（62.7% vs 63.8%）。在更具挑战性的 16K 滚动长度设置中，BF16-train-FP8-rollout 在 Qwen3-8B-Base 上完全崩溃无法收敛，Jet-RL 仍然稳定收敛并仅下降 2.7%。在 DeepMATH 数据集上（103K 高难度数学问题，16 次生成），BF16-train-FP8-rollout 的 MATH500 分数暴跌 25.6%（83.4% → 57.8%），Jet-RL 仅下降 3.2%（83.4% → 80.2%）。在效率方面，FP8 在 32B 模型（TP=2）上实现了 1.33× 的 Rollout 加速，在 8B 模型上实现了 1.41× 的训练阶段加速（主要来自 Actor Update 1.54× 和 Reference 推理 1.80× 的加速），端到端步时间加速 1.16×。

**Table 2**: Rollout 长度为 8K 时 Jet-RL 与基线的性能对比

**Table 3**: Rollout 长度为 16K 或 DeepMATH 上的性能对比

查看结构化数据

任务	指标	本文	基线	提升
Llama3.1-8B (8K, GSM8K+MATH)	Average	25.2%	23.2% (BF16)	+2.0% 超出 BF16 基线
Qwen2.5-7B (8K, GSM8K+MATH)	Average	55.9%	56.9% (BF16)	-1.0% vs BF16，而 BF16-train-FP8-rollout 完全不收敛
Qwen3-8B-Base (8K, GSM8K+MATH)	Average	62.7%	63.8% (BF16)	-1.1% vs BF16，而 BF16-train-FP8-rollout 下降 2.9%
Qwen3-8B-Base (16K, DeepMATH)	MATH500	80.2%	83.4% (BF16)	-3.2% vs BF16，而 BF16-train-FP8-rollout 下降 25.6%
Rollout 加速 (32B, TP=2)	吞吐加速比	1.33×	BF16 基线	33% 加速
训练阶段加速 (8B)	吞吐加速比	1.41×	BF16 基线	41% 加速
端到端加速 (8B, 8K rollout)	步时间加速比	1.16×	BF16 基线	16% 加速

局限与改进

论文的局限性体现在多个方面。首先，效率评估主要集中在 8B 模型的端到端训练上，14B 和更大模型的端到端加速数据未完整展示，作者也承认「A full scaling study on 14–32B models is left to future work given resource constraints」。其次，评估的下游任务集中在数学推理领域（GSM8K、MATH、GPQA、SuperGPQA），未涉及代码生成、指令遵循等其他 RL 训练的常见场景，Jet-RL 在这些场景下的表现尚不明确。第三，论文使用 GRPO 算法进行评估，未验证在 PPO、DAPO 等其他 RL 算法上的兼容性和效果。第四，Jet-RL 的反向传播中仍保留 BF16 精度的梯度，这意味着通信开销（特别是在张量并行场景下）并未得到优化。此外，论文在 DeepMATH 实验中因格式问题未报告 GSM8K 结果，实验完整性有一定折扣。

独立分析的弱点

从独立分析的角度，Jet-RL 存在以下几个值得关注的弱点。第一，统一精度流方案要求训练框架也执行 FP8 前向计算，这意味着如果未来 FP8 推理引擎本身的精度问题被更好的校准方法解决（如 SmoothQuant、AWQ 等进阶 PTQ 方法），Jet-RL 的训练侧也需要同步更新量化方案，缺乏独立优化的灵活性。改进方向：可以设计一种自适应精度流机制，根据训练阶段和任务难度动态调整量化粒度。第二，论文的加速收益在小模型（8B）上主要来自训练阶段而非 Rollout，这是因为 8B 模型的 Rollout 内存访问开销占比大，FP8 的计算加速效果有限。对于实际用户而言，如果使用 14B 或更大的模型，Jet-RL 的收益才更显著，但论文在这方面的数据不足。改进方向：在 14B、32B 甚至 70B 模型上进行完整的端到端基准测试。第三，Jet-RL 对 Rollout 引擎（vLLM）和训练框架（VeRL）的耦合较深，移植到其他框架（如 SGLang + OpenRLHF）可能需要较多工程工作。改进方向：抽象出通用的 FP8 精度流接口，实现框架无关的集成。

未来方向

论文和作者提出了一些未来方向，但基于 Jet-RL 的成果可以延伸更多。作者指出大模型（14-32B）的端到端加速研究受限于资源约束，是直接的下一步。此外，以下几个方向值得探索：（1）将 Jet-RL 的统一精度流思想扩展到 FP4 或更低精度，随着硬件对 FP4 的支持日益成熟（如 NVIDIA Blackwell 架构），在更激进的量化下维持训练稳定性将更具挑战也更有价值；（2）结合异步 RL 框架（如 AReaL）使用 Jet-RL，异步 RL 已通过解耦 Rollout 和训练提升吞吐，如果 Rollout 本身也能加速，两者可能产生乘积效应；（3）探索 Jet-RL 在非数学推理任务上的效果，如代码生成（CodeRL）、工具使用（Tool-use RL）等，这些任务的 Rollout 长度分布和难度特征与数学推理不同；（4）将 Jet-RL 的思路与 LoRA 结合，如 QeRL 框架所示，低秩适配 + 量化可能在单 GPU 上实现更大的模型训练。

复现评估

论文承诺将在去匿名化后开源代码和预训练模型，这对复现至关重要。从技术角度看，Jet-RL 的实现依赖三个关键组件：vLLM（推理引擎）、VeRL（RL 训练框架）和 DeepGEMM（FP8 kernel 库），三者都是开源项目。量化部分使用 Triton 实现，这也是开源的。复现所需的主要硬件是 NVIDIA H100 GPU，论文所有实验均在 H100 上完成。数据集方面，GSM8K（8.5K 问题）、MATH（12.5K 问题）和 DeepMATH（103K 问题）都是公开数据集。复现难度中等——虽然组件都是开源的，但将 FP8 精度流正确集成到 RL 训练管线中需要对量化 kernel 和分布式训练有较深理解。如果作者按承诺开源代码，复现难度将大幅降低。预计复现成本：8B 模型在 8×H100 上约需数小时的训练时间。

Jet-RL：通过统一训练与推理精度流实现 On-Policy FP8 强化学习 Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow