不匹配从哪来,又到哪去?

在 RL 中,我们先用推理引擎采样 rollout,再将采样序列送入训练引擎执行 prefill,根据得到的 logprob 计算梯度并更新模型。 两个引擎的前向可能在许多位置算出不同的 logprob,由此使 rollout 带有 off-policy 偏差;校正这种偏差的方法也不止一种。

训练端评分、推理 prefill 与 decode

推理在 prefill 阶段并行处理 prompt,写入模型的 cache,再通过每次生成一个 token 的 decode 继续生成。训练端则并行为整条完成序列评分。 这些路径上的 kernel 可能采用不同的舍入规则或加法顺序。对于循环层和压缩 attention,prefill 还会写入后续每一步 decode 都需要读取并更新的 state。

++ ++ 1 2−8 2−8 1 2−8 2−8 1 1 + 2−7 (a + b) + c a + (b + c)

舍入取决于加法顺序

BF16 中比 1 大的最小可表示数是 1 + 2−7。若计算 (1 + 2−8),结果会舍入为 1;但若先计算 (2−8 + 2−8),得到的是 2−7。只要涉及 reduction,我们就必须让两个引擎使用相同的加法顺序。

KERNEL A KERNEL B OP 1 BF16 OP 2 OUTPUT A OP 1 FP32 OP 2 OUTPUT B 相同运算 · 不同中间精度

精度

kernel 执行一系列运算时,每个中间值的精度都是计算定义的一部分。若一个引擎把中间结果舍入到 BF16,另一个保持 FP32, 后续运算就会读到不同的值,最终产生不同的输出比特。

XORL TRAINING SGLANG PREFILL SGLANG DECODE COMPLETE SEQUENCE · PARALLEL PROMPT STATE HANDOFF STATE +t STATE +t

缓存状态必须在 prefill 与 decode 之间对齐

Decode 会为每个新 token 读取并更新 cache。对于 Gated DeltaNet 这类带状态的层,prefill 写入的 cache 以及 decode 的每次更新,都必须与训练端的计算结果一致。因此,我们必须确保 prefill 与 decode 之间没有 mismatch。

不匹配如何改变 GRPO 更新

qθ 为训练端评估的策略, Ât 为 GRPO 中 detach 后的 group-relative advantageShao et al.,「DeepSeekMath」(2024)。引入 GRPO 及其 importance ratio。。 对前缀 ht 之后采样的 token at,on-policy 梯度是:

∇J(θ) = Σₜ Eₕₜ∼dqθ, ₐₜ∼qθ(·|hₜ)[Âₜ ∇ log qθ(aₜ | hₜ)]

每个 token 会被评分多次。 推理在采样时会记录该 token 的概率。 若再让推理对完成序列做一次 prefill,记录的概率会略有不同;训练端为这条 rollout 评分时,又会算出另一个值。 记 qt 为训练端选中 token 的概率, sPt 为当前权重下推理 prefill 回放的概率, sDt 为同一权重下带 cache 的 decode 概率,μt 为采样时记下的概率。我们可以用这些不同的前向 路径定义 importance sampling ratio

ρₜ = qₜ / μₜ = [qₜ / sᴾₜ] trainer scoring vs. serving prefill × [sᴾₜ / sᴰₜ] serving prefill vs. serving decode × [sᴰₜ / μₜ] stale state

我们在 RL 中优化的目标函数用 importance sampling ratio 给梯度加权,以校正这些 mismatch。

importance ratio 的细节

在固定前缀 ht 上,用 qt / μt 给样本加权,可以把 μ(· | ht) 上的期望改写为 qθ(· | ht) 上的期望。前缀本身 仍然来自 μ

At fixed hₜ, Eμ[ρₜ f] = Eqθ[f]. Across rollouts, Eμ[ρₜ f] keeps hₜ ∼ dμ, not hₜ ∼ dqθ.

过期权重与缓存状态

Dt 为当前权重下、 并在该权重上重建 cache 后的 decode。两项分别对应权重变化, 以及过期 cache:

sᴰₜ / μₜ = [s̃ᴰₜ / μₜ] stale weights × [sᴰₜ / s̃ᴰₜ] stale cache

我们在更新推理引擎权重后总会清空 cache,因此第二项 是 1。

不截断的 importance sampling

不截断的 importance sampling 将 token 梯度乘以 ρt。构成该比值的 选中 token logprobs 会被 detach,只对 log qt 求导。Tinker 的 importance-sampling loss 使用这一形式。

∇Lᴵˢₜ = −ρₜ Âₜ ∇ log qₜ; Eμ[‖gₜ‖²] = Eμ[ρₜ² Âₜ² ‖∇ log qₜ‖²]

梯度的二阶矩随 importance ratio 的平方缩放。我们在 有限大小的 batch 上训练,因此会出现高方差步:梯度被少数 大比值 token 主导。

截断后的 importance ratio

GRPO 通常用 PPOSchulman et al.,「Proximal Policy Optimization Algorithms」(2017)。引入截断代理目标。 目标来截断这一比值:

Jᴾᴾᴼₜ = min(ρₜ Âₜ, clip(ρₜ, 1 − εlow, 1 + εhigh) Âₜ); ∇Lᴾᴾᴼₜ = −1[not clipped] ρₜ Âₜ ∇ log qₜ

对优势为正的 token,一旦 ρt 超过 1 + εhigh,PPO 就舍弃策略梯度。对 advantage 为负的 token,一旦 ρt 低于 1 − εlow,同样舍弃梯度。

CISPOMiniMax-M1 引入 CISPO。 在乘入梯度之前截断 ρ,从而限制超过阈值的 token 的贡献,但并不舍弃其梯度。它的边界是绝对比值,而不是围绕 1 的偏移;默认 关闭下界,因此比值只会被从上方封顶。

ρ̄ₜ = clip(ρₜ, ρmin, ρmax); ∇Lᶜⁱˢᵖᵒₜ = −sg[ρ̄ₜ] Âₜ ∇ log qₜ; (ρmin, ρmax) = (0, 4)

当每个 token 都有 ρt = 1 时, PPO 不会截断,CISPO 也不会封顶,更新就简化为上面的 on-policy 梯度。

如何让训练端与采样端达成一致

我们希望各次前向的 logprob 对齐,因此从最终的 logprob 出发, 沿训练端与采样端的计算路径反向追踪。对每个 算子,我们记录其输入、舍入点加法顺序。对缓存状态,我们比较 prefill 写出的内容 以及 decode 的每一次更新。对分布式算子,我们记录到达了哪些值,以及它们 被加在一起的顺序。

从选中 token 的 logprob 反向追踪到第一个数值差异 训练端评分、推理 prefill 与推理 decode 从相同的输入和权重出发。比较它们选中 token logprob 的字节,再沿前向反向追踪到第一个不同的算子,并将该算子缩小为一行、一次 reduction 或一次状态更新。对齐其输入、常数、舍入点、完整加法顺序以及状态交接。只有在生产形状下的字节比较再次通过之后,才允许调整调度、布局与传输。

对每个算子,我们向两个引擎提供相同的输入并比较输出字节。一旦出现差异,就把问题 缩小到仍能复现 mismatch 的一行、一次 reduction 或一次状态更新,再对齐这个算子。

对齐一个算子的四种方法

  1. 在训练端前向中调用推理 kernel。我们从 XoRL 调用 SGLang 的 前向 kernel,并为其编写训练用的自定义反向。
  2. 让两个引擎使用同一个 kernel。我们编写一个新 kernel,或扩展已有 kernel,并在 XoRL 与 SGLang 中都调用它。
  3. 让 reduction 对 batch 不变。我们固定决定 reduction 形状的参数, 这样更大的 batch 只是多启动同样的工作,而不会把加法 重新分组。
  4. 编写两个执行相同算术的 kernel。当两个引擎需要不同的 形状或布局时,我们分别编写 kernel,让它们在相同的位置舍入,并按相同顺序相加。
ρₜ = qₜ / μₜ = [qₜ / sᴾₜ] trainer scoring vs. serving prefill × [sᴾₜ / sᴰₜ] serving prefill vs. serving decode × [sᴰₜ / μₜ] stale state

我们希望每一项都等于 1。对于同步 rollout 且 cache 即时重建的情况,最后一项(过期 state)就是 1。我们的工作是让训练 prefill 与推理 prefill 对齐,再让 推理 prefill 与推理 decode 对齐。

所有涉及 cache 的组件都需要额外检查。训练、prefill、eager decode 与 CUDA-graph decode 都会 写状态,四者必须在每个 token 边界写出相同的字节。

训练模型时,我们用 K3,t = ρt − 1 − log ρt 报告 train-infer mismatchSchulman,「Approximating KL Divergence」(2020)。引入 K3 估计量。。它处处 非负,且仅当 ρt = 1 时为零。

我们达到的保证比 K3 = 0 更强:在 1 附近它的行为像 (ρt − 1)2/2,因此在 FP32 里已经变成 0,而 logprob 仍可能有微小差异:差 1e-8 就已经精确为零。 训练端评分、推理 prefill 与推理 decode 的 FP32 选中 token logprob 是 bitwise-identical 的。尽管 SGLang 提供了 batch-invariant kernel 开关,但仅凭 SGLang 仍无法实现 0 train-infer mismatch。我们必须在其基础上继续扩展。

验证阶梯

最终,我们训练完整模型,并在同一权重版本上比较每一个选中 token 的 logprob。测试按难度递增:

  1. 将同一段程序执行两遍,比较输出字节。
  2. 改变 batch 形状与启动配置,比较同一批行。
  3. 向 XoRL 与 SGLang 提供相同的输入。
  4. 在同一权重版本的在线采样中比较它们。
  5. 训练完整模型,并要求 K3 保持为零。

接下来,我们会在一系列模型上测量 train-infer mismatch,以展示这套方法, 从稠密的 Qwen3-8B 开始。

Qwen3-8B:对齐一个稠密 Transformer

在 Qwen3-8B 上,两个引擎可能在常数构造、舍入位置和 reduction 顺序上无法对齐。这里消除 mismatch 相对直接,也为后续处理更复杂的模型提供了可复用的基础组件。 我们会逐个检查模型中的算子,说明 mismatch 从何而来、如何修复,先从 LoRA 开始。

LoRA:前向对齐的两种方式

LoRA 将线性层从 W(x) 改为 W(x) + sB(A(x))。我们支持 两种 serving 模式。合并模式在推理前将 sBA 融入基座权重。分离 模式则保持 AB 分开,并将 adapter 激活加到 W(x) 上。这两条 路径的舍入点不同:合并路径在乘 x 之前先舍入更新后的权值, 分离路径则分别舍入基座投影与 adapter 投影,再相加。

固定 adapter 合并后更快,因为推理无需执行额外的 adapter 投影。在 multi-LoRA 训练中,XoRL 为每个实验保留独立参数与 optimizer state,而 SGLang 则在一次部署中将多个 adapter 的请求合并到同一个 batch。

在合并模式中,我们在 XoRL 与 SGLang 中构造 W′ = W + sBA,并使用相同的权值字节运行模型。 我们在 XoRL 中缓存 W′,直到 AB 改变,并用自定义 autograd 算子将梯度映回 AB。在分离模式中,我们从 XoRL 调用 SGLang 的 LoRA 前向,并实现训练 AB 所需的反向。当量化 expert 需要新的 LoRA 前向时,我们先将 LoRA 运算加入 SGLang,再从 XoRL 调用该前向。

方法 1:在训练端前向中调用推理 kernel

在 XoRL 的前向中调用推理 kernel,并补充一份自定义反向,是消除 train-infer mismatch 的有效方法,但训练效率可能不如专门编写的 kernel。

代码 XoRL 合并模式 LoRA 前向与反向 ↗ XoRL LoRA 反向 ↗ XoRL 对 SGLang LoRA kernel 的调用 ↗

RoPE:对齐表与旋转

CPU 表与 GPU 表会分叉

RoPE 根据与位置相关的正弦和余弦,旋转隐状态中成对的分量。两个引擎可能在 何处构造正余弦表、以及旋转过程中 何时舍入上无法对齐。

对 Qwen, SGLang 在 CPU 上构造逆频率,并在 GPU 上完成位置表, 以便并行计算大外积与正余弦。它只传输较小的 逆频率向量。XoRL 在 CPU 上构造完整表, 因为它在 FSDP 给每个 rank 分配执行 GPU 之前,就在 CPU 或 meta tensor 上构造模块与派生常数。

CPU 与 CUDA 的正弦、余弦实现不同,因此两个引擎可能产出不同的值。当我们用 Qwen3-30B-A3B decode 一条序列时,表的 BF16 值在位置 592 之前都一致。 它们在位置 593 首次不同,下一个 token 上就出现第一次 logprob 不匹配。 这是一条重要经验:必须在足够长的序列上测试 mismatch。

CPU 构造的表与 GPU 构造的表在位置 593 首次无法对齐。从此处起,同一旋转会给出不同输出。

RoPE 首次分叉的位置 CPU 构造与 GPU 构造的表在位置 592 之前给出相同的 BF16 值,并在位置 593 首次不同。再施加同一旋转会得到不同的输出对,下一个 token 上出现第一次不匹配的 logprob。

在 CPU 上只构造一次表

我们让两个引擎以相同方式构造 RoPE 表:在 CPU 上构造完整的 FP32 表, 缓存后再传到 GPU。在任一引擎施加旋转之前,表的字节就已对齐。

表字节对齐之后,两个引擎都可以选用两种旋转之一。Class A 将正弦、余弦 cast 为 BF16,并在整个旋转过程中舍入。Class B 将旋转保持在 FP32,仅在输出处 cast。我们用 Class B,因为它 更快快 1.23×。 也更准确相对 FP64 的 MAE 大约低 40%。

两种 RoPE 实现里的算术与舍入。
Class AClass B
算术 将正弦、余弦 cast 为 BF16,再对中间乘积与和进行舍入 将正弦、余弦与完整旋转保持在 FP32,再将输出舍入到 BF16
舍入 每对旋转执行八次 BF16 舍入:六次中间舍入,两次输出舍入 每个输出只执行一次最终 BF16 舍入

代码 SGLang RoPE 实现 ↗ XoRL RoPE kernel ↗ XoRL 按模型构造 RoPE ↗

RMSNorm:固定 reduction 顺序

RMSNorm 暴露出一个新问题:编译器可以改变 reduction 的执行方式。

为什么会有不匹配?

RMSNorm 将一行中的每个值平方,求和后取均值的倒数平方根,再用其缩放该行。RMSNorm 在计算平方和时需要 reduction,因此两个引擎必须采用相同的 reduction 顺序。

XoRL 在上一层返回前就将 residual 加入,因此下一层 RMSNorm 得到的是一行 BF16。 SGLang 将隐状态与 residual 分开传入,并在 RMSNorm 内部相加。这两种 输入签名会派发到不同的 Triton kernel。

单输入 kernel 载入一行 BF16 并在寄存器中平方。residual kernel 在 FP32 中将两个输入相加, 将结果舍入到 BF16 并存为 residual 输出;然后再存储并重载 FP32 平方。单输入 kernel 计算 1 / sqrt,residual kernel 则用 rsqrt二者在数学上等价,但可能舍入成不同的 FP32 值。

residual kernel 还会为下一层存储 BF16 residual 行。这次存储 并不进入 reduction,但它的存在会把 Triton 赋给 reduction 操作数的 layout 从 sizePerThread = [4] 改成 [8]sizePerThread 是每个线程拥有的连续段,不是总数。32 线程一个 warp、四个 warp,128 个线程覆盖 1,024 个值,因此两种 layout 下每个线程都拥有 8 个值。 无论哪种,每个线程仍然拥有 8 个值。变的是哪 8 个:[4] 下线程取两段长度为 4、相距 512 的段,[8] 下则取一段连续的 8 个。 下图单独展示这一效应:同一行 tl.sum 源码,仅添加必需的 BF16 store,就会改变加法顺序,8,192 个输出中有 2,854 个发生变化。

mean = tl.sum(x, axis=0)

同一行 tl.sum 源码下编译器选出的 layout 单输入路径没有 residual store,每个线程处理 4 个值。residual 路径必须存储 BF16 residual 输出,每个线程处理 8 个值。源码中的 reduction 与指令条数不变,但新 layout 会将值分到不同组,8,192 个输出中有 2,854 个发生变化。 单输入 · 无 residual store · sizePerThread = [4] · thread 0 ← x[0..3] + x[512..515] residual 路径 · 必需的 BF16 store · sizePerThread = [8] · thread 0 ← x[0..7] 相同指令条数 · 不同加法顺序2,854 / 8,192 个输出改变
Triton 发出了什么

两个 kernel 都用 triton.compile 对着固定的 sm90 目标、Triton 3.7.1、四个 warp、关闭 autotuning 来编译。两份 TTIR 模块都包含同一个覆盖 1,024 个 FP32 值的 tt.reduce 运算。layout 分配之后,reduction 操作数在单输入 kernel 里是 sizePerThread = [4],在带 store 的 kernel 里是 [8]。 PTX 每个线程发出的 add.f32 条数相同,两种 layout 都是 7 条,只是操作数不同。

// one-input kernel
"tt.reduce"(%vals) <{axis = 0}> : tensor<1024xf32, #blocked>
#blocked = sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [4]
  thread 0 accumulates  x[0]+x[1]+x[2]+x[3] + x[512]+x[513]+x[514]+x[515]

// same line, with the BF16 residual store present
"tt.reduce"(%vals) <{axis = 0}> : tensor<1024xf32, #blocked>
#blocked = sizePerThread = [8], threadsPerWarp = [32], warpsPerCTA = [4]
  thread 0 accumulates  x[0]+x[1]+x[2]+x[3]+x[4]+x[5]+x[6]+x[7]

按每种 layout 发出的加法顺序回放,就能复现全部 8,192 行的 GPU 输出。

高效的精确一致

方法 2:给两个引擎各加同一个 kernel

我们为两个引擎添加同一个 RMSNorm kernel。kernel 先形成 BF16 行。若隐状态与 residual 分开传入,就在 FP32 中相加并舍入一次到 BF16;若结果已经存在,就直接使用该行 BF16。此后两个引擎都在 FP32 中将该行平方。我们不用 tl.sum,而是采用显式的树形 reduction。然后在 FP32 中施加 rsqrt 与模型的 scale,再将输出 cast 为 BF16。

我们的 kernel 从两种不同输入类型产生同一份 BF16 tensor,再由两个引擎执行同一个融合 kernel。

两种输入形式共用一次显式 RMSNorm reduction 我们的 kernel 从两种不同输入类型产生同一份 BF16 tensor,再由两个引擎执行同一个融合 kernel。

对大多数形状,一次融合启动即可完成计算。当只有少量非常宽的行时,我们将同一次计算拆成三次启动,以便让更多 GPU block 并行执行。

代码 XoRL RMSNorm ↗ SGLang RMSNorm ↗

GEMM:沿 K 执行同一次 reduction

GEMM 沿 K 维累加乘积,得到每个输出。训练会处理数千行 packed token, 而 decode 常常只乘一行。像 cuBLAS 这样的库会为每种形状选不同的 kernel,这些 kernel 可能按 不同顺序加同一组 K 乘积。Split-K matmul 是清楚的例子:若干 thread block 分别累加 K 的片段,再合并部分和。当 M 变化时,库可能选出不同的 split 数,从而改变 reduction 顺序并造成不匹配。

方法 3:让 reduction 对 batch 不变

我们编写一个 Triton kernel:每个 program 计算一块输出 tile,并从左到右累加完整的 K 维。BF16 输入使用 64 个值的 slab。kernel 使用一个 FP32 累加器从第一块遍历到最后一块,只 cast 一次输出。现在 M 变大时,我们只是多启动同样的工作。这就是一个简单的 batch-invariant kernel。

两个 kernel 都从左到右遍历 K 上的 slab,并全程保留同一个 FP32 累加器,直到最后的 BF16 cast。

每种 batch 大小都用同一次从左到右的 K reduction 训练与 decode kernel 都从左到右处理 K 上的 slab,为每个输出带着一个 FP32 累加器直到最后的 BF16 cast。

我们将这份 Triton kernel 作为参考;batch-invariant matmul 如今已有成熟方案,无需再自行将性能优化到极致。对于受支持、能产出与参考 kernel 相同比特且速度更快的形状,我们使用 DeepGEMM;否则使用我们的 Triton kernel,例如 Qwen3.5 MoE 中的 shared-expert gate 投影。

代码 XoRL 固定 K 矩阵乘 ↗ SGLang 固定 K 矩阵乘 ↗

SwiGLU:在乘法之后再舍入

SwiGLU 计算 SiLU(gate) * up,这是 舍入规则的另一个例子。训练与推理可能在中间 SiLU 结果 是否舍入上无法对齐。

我们为两份代码库添加了只舍入一次的 SwiGLU kernel。它将输入转成 FP32,计算 SiLU 与乘法,再将输出一次 cast 为 BF16。去掉中间的 BF16 舍入后,结果更准确MAE 为 4e-4,SGLang 原生实现为 6e-4。。 将两次运算融合为一次启动,并沿行和 hidden dimension 分块,也使其 更快decode 吞吐提升 2.56%,端到端吞吐提升 2.35%,prefill 吞吐提升 0.45%。

代码 XoRL 只舍入一次的 SwiGLU kernel ↗ SGLang 只舍入一次的 SwiGLU kernel ↗

Dense attention:一次流式 reduction

Dense attention 是一种 token 混合算子:它对 value 的低维投影进行加权线性组合,权重来自 key 与 query 的低维投影点积的 softmax。 输入中存在较大值时,直接计算 softmax 可能数值不稳定;由于 softmax 对平移不变,可以先从输入中减去最大值。朴素实现需要扫描整个向量来求最大值。FlashAttention 则采用 流式 softmax:它从每块输入 tile 中提取最大值 m、指数和 l,以及加权 value 和 u。在归约两块 tile 级摘要之前,先将它们缩放到同一个最大值。归约所有 tile 级摘要后,即可得到 attention 输出 u / l

m = max(m₁, m₂); l = exp(m₁ − m)l₁ + exp(m₂ − m)l₂; u = exp(m₁ − m)u₁ + exp(m₂ − m)u₂; Attention = u / l

推理引擎可以通过 多种 attention backend 完成这项计算,例如使用 FlashAttention 或 FlashInfer;训练端则可能使用 FlashAttention、FlexAttention,甚至 SDPA。每个 backend 分块与合并摘要的方式都不同。

训练会处理许多 packed query,而 decode 通常用一条 query 访问很长的 key/value cache。推理 backend 可能将 cache 切分,让多个 block 同时处理这条 query,再通过第二次 reduction 合并各自的 (m, l, u) 摘要。改变 split 数 就会改变加法的先后顺序。不同版本的 FlashAttention 也可能对第一次 reduction 分块得不一样。

对 dense attention,我们在两个引擎中都使用带 num_splits=1 的 FlashAttention 4,因此两边都不再做第二次合并。

LM head:对词表执行一次 reduction

LM head 接收最终隐状态,为词表中的每个 token 计算一个 logit。 对选中的 token y,其 logprob 等于选中 logit 减去整个词表上的 log-sum-exp:

log p(y | h) = zᵧ − log Σⱼ exp(zⱼ)

LM head 将固定 K 的矩阵乘与流式 log-sum-exp 融合在一起。输出取决于 矩阵乘的累加顺序、词表分组边界,以及用来 合并各组摘要的树。

两个引擎都用上固定 K 的矩阵乘之后,它们的 logits 就对齐了,包括选中 token 的 logit zy。剩下的不匹配可能来自词表上的 log-sum-exp。 每个词表组产出一个最大值与一个指数和 (m, l),头再合并这些摘要。训练端可以一次分 256 个 logit,而 decode 分 128 个。分组不同意味着局部最大值不同,缩放后的和就会按不同顺序相加。

我们让两个引擎都以 256 个 logit 为一组,并在组内及合并各组摘要时使用同一棵显式 reduction 树。

方法 4:编写两个执行相同算术的 kernel

SGLang 需要整行 FP32 logits 来采样 token,而 XoRL 只需要被采样 token 的 logprob。因此,XoRL 直接从 GEMM tile 的 FP32 累加器计算选中 logit 和每一份流式 softmax 摘要,无需物化完整的 [M, V] logits tensor,从而节省大量显存与计算。这说明,为了提高效率,两个引擎可以使用不同的 kernel,只要它们执行相同的算术并产生 bitwise-identical 的输出。

什么必须固定,什么可以自由调整

要让两个引擎一致,无需冻结每个启动参数。两边必须固定的是决定 reduction 形状(进而决定顺序)的参数;只影响调度的参数可以自由调整。例如在 LM head 中,我们固定词表维的 tile 宽度和 K slab 宽度,因为它们决定哪些 tile 会被归入同一份摘要,以及沿 K 归约的方式。但每个引擎仍可根据工作负载自行选择行块大小、group size、pipeline stage 与 warp 数。

代码 XoRL LM head ↗ SGLang LM head ↗

采样变换

我们已经得到 bitwise-identical 的 logprobs,但推理引擎通常不会直接从 这个分布采样。我们可能只想从分布的 top-k 下标采样,以防止 reward hacking,或施加 temperature 来调节探索。这些采样运算改变了 我们从中采样的分布,因此需要把 temperature 等参数转给训练端,并 回放同一套变换。

我们回放整套变换(temperature、top-k、top-p、min-p),而不是变换产生的 mask,这样训练端会通过相同的排序与阈值规则得到与采样端一致的 support。

代码 XoRL 采样变换 ↗ SGLang 采样变换 ↗

0 mismatch 与性能

我们在 Wordle 上训练 Qwen3-8B 时,bitwise-aligned 前向使采样吞吐降低了 16.2%。

Qwen3-8B 的采样吞吐:关闭对齐开关、使用 bitwise-aligned 前向,以及从中去掉 batch-invariant dense GEMM。三组都使用相同的 BF16 serving 策略与 CUDA decode graph,任务是 256 条 Wordle 轨迹;每组数据均为单台服务器上去掉 warmup 后三次计时的均值。
配置吞吐相对关闭开关
关闭开关2,116 tok/s1.000×
bitwise-aligned1,774 tok/s0.839×
bitwise-aligned,去掉 batch-invariant 稠密 GEMM1,970 tok/s0.931×

其中大部分来自 batch-invariant 稠密 GEMM。

代码 XoRL 稠密 Qwen 实现 ↗ SGLang 稠密 Qwen 实现 ↗

在不同并行布局下保持同一套算术

用 dense attention 训练稠密模型时,实现 0 mismatch 相对直接。 但我们的目标是在 GLM-5.2 这类模型上实现 0 mismatch。 在讨论各开源权重 MoE 的具体复杂性之前,先明确一点:在这一规模下进行训练和 serving 必须采用 sharding。训练端与采样端可以使用不同的切分方式,因为训练端还要保存 activation 和 optimizer state。这会改变每个值所在的 GPU、reduction 执行的位置,以及 cache state 如何跨设备传递。因此,对每一种并行方式,我们都要确保其输出与未切分计算 bitwise-identical。XoRL 使用 CP、PP、EP,以及用于 attention 和 expert 的 DP。

我们沿四个轴切分:一次训练许多条 rollout,每条 rollout 可以长达数千 token,模型有许多层,每个 MoE 层可以有数百个 expert。

规模的四个轴 一叠 rollout,左侧括号标为许多 rollout,底部括号标为长序列。右侧是从第 0 层到第 99 层的层叠,括号标为许多层;第 k 层展开为注意力块以及从 E1 到 E128 的 expert 网格,标为许多 expert。 许多 rollout ok let me think about how to... import torch; model = GPT2... to train gpt2 we must first... here is my code for training... first we need to tokenize the... the learning rate should follow... 长序列 Layer 0 Layer k Attention E₁ E₂ E₃ E₄ E₅ E₆ E₇ E₈ E₉ E₁₀ E₁₁ E₁₂ E₁₃ E₁₄ E₁₅ E₁₆ E₁₂₈ 许多 expert Layer 99 许多

切分 expert

我们将 128 个被路由的 expert 分成四组,分别对应 4×4 网格的四列:EP₀ 拥有 expert 0–31,EP₃ 拥有 expert 96–127。

沿网格列执行 Expert Parallel 四组 expert,从 0–31 到 96–127,颜色分别对应 4×4 GPU 网格中标为 EP0 到 EP3 的各列。 Experts 0–31 Experts 32–63 Experts 64–95 Experts 96–127 Expert Parallel EP₀ EP₁ EP₂ EP₃ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15

Expert Parallel 中, 每个 rank 计算其拥有的被路由 expert,以及 shared expert 中属于自己的部分,最终每个 token 得到一份 BF16 部分输出。一个 rank 可能拥有某个 token 选中的八个 expert 中的任意数量,因此 expert kernel 按 slot 顺序在 FP32 中累加每个 routing slot; 其他 rank 所拥有 slot 的贡献精确为零,最后只舍入一次到 BF16。若训练端与推理使用不同的 collective,它们可能采用不同的 加法树合并各 rank 的输出。

我们对各 rank 的输出执行一次显式的树形 reduction。首先 交换 BF16 部分输出而不做 reduction,将它们按固定的逻辑顺序排列,再 将相邻项两两相加。每对值 在 FP64 中累加;若贡献者数量为奇数,则将最后一项原样传至下一层。重复这一过程直到只剩一个值, 最后舍入一次到 BF16。交换过程本身无需对齐:训练端 在 expert 组上执行 all-to-all,serving 则在 tensor-parallel 组上使用可由 CUDA graph 回放的 collective 执行 gather。每个引擎只需按相同的逻辑顺序提供相同的部分输出。

这是我们处理 collective 通信的通用策略。默认情况下, 执行 reduction 的 collective 会根据拓扑与消息大小选择 reduction 树,从而造成 mismatch。我们改用只交换值、不执行 reduction 的 collective, 再执行显式的 reduction 树,以控制归约顺序。DeepEP kernel 会在传输途中累加部分和,因此我们无法控制加法顺序。我们只使用 DeepEP kernel 传值,并在 DeepEP 之外执行 reduction。显式 reduction 树在 FP64 中运行,以限制精度损失;如果也对在途执行 reduction 的 collective 使用 FP64,通信量会显著增加,因此通常不会这样做。

我们 gather 八份 BF16 部分输出,按逻辑 rank 顺序排列,再在 FP64 中两两归约,最后舍入一次到 BF16。

对八份 rank 输出做一次 batch-invariant reduction 我们 gather 八份 BF16 部分输出,按逻辑 rank 顺序排列,再在 FP64 中两两归约,最后舍入一次到 BF16。

代码 XoRL expert 输出相加 PR ↗ SGLang expert 输出相加 PR ↗

切分序列

我们将一条 rollout 切成四份,并将每一份分配给网格同一行中的一块 GPU。

沿网格一行执行 Context Parallel 一条 rollout 被拆成十六块 token tile、四份各四块,每份颜色更浅一档,映射到网格同一行的四块 GPU 上,标注 Context Parallel = 4。 长序列 ok let me think about how to train gpt two from scratch using py torch ... 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 Context Parallel = 4

我们在 RL 中的主要目标之一是训练长上下文。对大模型而言, 这要求沿上下文切分。XoRL 用 Ulysses 与 RingAttention 实现 Context Parallel。

Dense attention 相对直接。Ulysses 按 head 将 attention 计算分配到各 rank。 每个 rank 通过 all-to-all 获取其负责的一组 head 上的全部 token。因此,无论训练端是否沿序列切分,输出都相同。

不同的开源权重 MoE 往往采用不同的 attention 机制,其中一些会使 本节讨论的不变量不再成立。我们稍后会具体说明。

切分 batch

四条 rollout 分配到网格的四行。括号跨过这些行,表示参数、梯度与 optimizer state 沿这些行切分。

沿网格行执行 Fully Sharded Data Parallel 四条彩色 rollout 映射到 4×4 GPU 网格的四行,标题为 Fully Sharded Data Parallel,括号跨过全部四行。 ok let me think about... import torch as... to train gpt2 we must... here is my code for... 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 许多 rollout Fully Sharded Data Parallel

Data Parallel 沿 batch 维切分:每个 rank 处理自己的 rollout 并执行完整计算, 因此每条 rollout 的计算过程保持不变。XoRL 沿 context 和 data parallel 维切分 attention,沿 Expert Parallel 和 FSDP 维切分 expert 层。若改变 attention 的切分方式,某个 token 的行会落到不同 rank 上,因此无论如何切分,expert 输出的 reduction 都必须产生相同的比特。

因此,我们不能用 rank id 为expert 输出 reduction 排序。 我们改为给每个贡献者分配一个以 data-parallel 为主、context-parallel 为次的标签,并按 expert shard 计数贡献者。这样,无论 attention 在 CP 还是 DP 上运行,reduction 树的叶子都相同。然后按这些标签执行 reduction,使每个 token 得到与完全不切分时相同的 reduction 树,而不受实际行切分方式影响。 各 rank 拥有的行数也可以不同,因此每一行都携带自己的位置;padding 会被标记为无效,不参与 reduction。

按层切分模型

我们把第 0–49 层放在 stage 0 的十六块 GPU 上,第 50–99 层放在 stage 1。P2P 标记标出激活跨 stage 的位置。

跨两个 stage 的 Pipeline Parallel 模型的两半,第 0–49 层与第 50–99 层,分别映射到拥有 GPU 0–15 的 stage 0 与拥有 GPU 16–31 的 stage 1,由 P2P 标记相连。 Pipeline Parallel Layers 0–49 P2P Layers 50–99 Stage 0 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 P2P Stage 1 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31

Pipeline Parallel 按层将模型切分到不同设备。在每个 decoder 层中,我们 返回已经加上 residual 的 BF16 输出。这样只需将已存储的 tensor 传过 pipeline 边界,无需增加计算或 cast。无论训练端是否使用 Pipeline Parallel,输出都相同。

代码 XoRL 头切分与 pipeline stage 对齐 PR ↗ XoRL context-parallel GDN 分支 ↗ XoRL 逻辑行归属 ↗

每次交换后,接收 rank 得到按 head 划分的 BF16 值,或跨 stage 传递的 BF16 行——与未切分时计算出的值完全相同。

Qwen3.5/3.6:Gated DeltaNet 与 mixture-of-experts 层

Qwen3.5/3.6 系列带来了新的挑战:Gated DeltaNet 这样的循环层会将一份 FP32 state 从一个 token 传到下一个,state 中的任何差异都会沿序列传播。对 Qwen3-8B,我们只需对齐 XoRL 与 SGLang;现在还要对齐 SGLang prefill 与 SGLang decode。这也是我们遇到的第一个 MoE,因此还会讨论 MoE 对齐的特殊难点。

Gated DeltaNet:对齐循环状态

Gated DeltaNetYang, Kautz, and Hatamizadeh, “Gated Delta Networks” (2024)。引入本架构使用的循环线性注意力层。 维护一个持续更新的 FP32 矩阵 H,汇总当前 token 之前的全部信息,并在每个 token 上更新一次。对于单个 head 和单个 token,其循环形式为:

H̄ₜ = exp(gₜ)Hₜ₋₁, δₜ = βₜ(vₜ − H̄ₜᵀk̂ₜ), Hₜ = H̄ₜ + k̂ₜδₜᵀ, oₜ = Hₜᵀq̂ₜ/√dₖ

这里 g 是衰减门,β 是更新强度,v 是当前 token 的 value, 是 L2 归一化后的 query 与 key。更新 Ht = (Iβttt)t + βttvt 就是 delta rule。

训练与 decode 对循环的分组不同

训练、prefill 与 decode 对循环采用不同的分组方式。XoRL 每次计算 64 个 token。 SGLang 在 prefill 时也用 64-token 的 chunk,但 decode 一次更新一个 token 的 H。 64-token 计算与逐 token 计算 对 FP32 乘法与加法的分组不同,因此可能在 chunk 边界产生不同的 state。下一个 token 读取该 state 后,差异会沿序列继续传播。

64-token block 内部的计算也不同。SGLang 的 kernel 将 state 存为 [V, K],采用对应的转置乘法顺序,再通过单独的小块 pass 和一次 merge 求逆 64 行三角因子。XoRL 将 state 存为 [K, V],并使用一次融合 merge。两边计算的是同一组和,但 FP32 分组不同,因此在长上下文中会有少量元素无法对齐。

在 decode 时回放 64 行计算

两个引擎使用相同的 reduction 顺序进行 q/k 的 L2 normalization,将 gβ 保持在 FP32,并在对衰减取指数之前 mask 掉未来位置。Serving 也使用 XoRL 的 [K, V] state kernel 及其融合三角 merge,因此两边 在构造 chunk 边界的那些阶段执行同一套 kernel 体。对 XoRL 的前向, 我们调用 SGLang 的因果卷积 kernel 并实现其反向。

XoRL 训练与 SGLang prefill 使用一次 64 行状态更新,prefill 在每个完整 chunk 之后存储 FP32 状态。要 decode token p,我们从该边界出发,按训练顺序重算第 0 到 p 行。到达下一个 chunk 边界之前,最多重算 64 行。

不匹配 XoRL 训练每次计算 64 行,而 SGLang decode 每次更新一个 token 的 H。SGLang prefill 采用相同的 64 行分组,但按 serving 方向传递 state。

不同的 chunk 边界状态会改变之后的每一个 chunk XoRL 训练每次计算 64 行,而 SGLang decode 每次更新一个 token 的循环 state。SGLang prefill 使用 64 行分组,但按转置后的 serving 方向传递运行 state,并用单独的小块 pass 求逆三角因子。不同的浮点程序可能在 chunk 边界产生不同的 FP32 state。

已对齐 要 decode token p,SGLang 从与训练、prefill 相同的 FP32 chunk 边界出发,按相同顺序计算第 0 到 p 行。

在 decode 时回放训练用的 chunk 要 decode token p,SGLang 从与训练、prefill 相同的已完成 FP32 state 边界出发,在取指数之前 mask 掉未来位置,并按相同顺序计算第 0 到 p 行。

循环具有因果性,因此较早的行不能依赖后续 token。这意味着我们可以在 SGLang 中缓存中间值,只重算新增行涉及的部分。prefill 之后,一个 batched initializer 会为 prompt 最后一个不完整的 chunk 初始化 cache。每次 graph replay 新增一行。

完成一次 eager decode 后,我们重新执行缓存行涉及的阶段,使 CUDA-graph decode 能从预期的 cache 继续运行。一个 chunk 完成时,我们在 SGLang 中保存下一个 FP32 边界,并开始新的 64 行 slab。

CUDA graph 要求固定的 tensor 地址,但请求会在调度 slot 之间移动,每个活跃 chunk 每步都会增加一行。我们为每个 slot 预留 64 行,并将当前 slot ID 与 chunk 长度写入 GPU tensor。我们确保这些 SGLang 改动仍然支持 CUDA Graph。

每个调度 slot 预留 64 行。GPU tensor 选择 slot 与当前行。Graph replay 复用缓存行,更新第 p 行,并执行与完整重扫相同的 64 行输出计算。

在 CUDA graph 中更新一行 GDN 每个调度 slot 在固定地址预留 64 行。GPU tensor 选出 slot 与当前行。Graph replay 更新该行,产出与完整重扫相同的输出。

循环层不能按 head 切分,因为 state 沿序列维传递,而不是跨 head 传递。若在 chunk 中间切分序列,就会重新组织循环、改变 reduction 顺序,进而改变 state,并将差异传播到后续所有 token。因此,我们只在每份文档的 64-token chunk 边界切分序列。collator 会在文档之间添加 padding,使每个切分点都落在边界上。每个 rank 在自己的 shard 上执行未经修改的 64 行计算,并将 FP32 边界 state 传给下一个 rank。因果卷积需要读取切分点之前的三个 token,因此前一个 rank 会发送这些投影后的输入,供接收 rank 继续计算。

代码 XoRL GDN chunk 前向与反向 ↗ SGLang 增量 GDN decode ↗ SGLang 整 chunk GDN decode ↗ SGLang Qwen GDN 选择 ↗

Mixture-of-experts 层

Router 选出八个 expert 并赋权。每个 rank 计算本地被路由的 expert 以及 shared expert 的一份 shard。我们收集八份 BF16 部分输出,并按固定顺序归约。

Router:对齐 expert id 与权重

Router 使用固定 K 的矩阵乘,为每个 token 计算对应各 expert 的分数,再执行 FP32 softmax, 选出最高的八个概率并重新进行 normalization。expert 选择是离散的:gate 投影中最后一个比特的差异就可能使第八和第九个 expert 互换。即便 expert id 一致,以不同顺序累加这八个选中概率,也会改变 BF16 routing 权重。

我们在两个引擎中都执行固定 K 的投影。然后从左到右累加这八个被选的 FP32 概率,将每个概率除以总和,再将 routing 权重 cast 为 BF16。两个引擎都按概率顺序返回被选 expert,因此后续沿八个 slot 执行的每次求和都会以相同顺序访问它们。

代码 XoRL top-8 路由 ↗ XoRL router 矩阵乘 ↗

Expert:将 routing 权重保留在累加器中

训练端可能在乘 routing 权重之前,先将 down-projection 结果舍入到 BF16;推理端则在 FP32 down-projection 累加器中乘以该权重,最后只舍入一次。

对 XoRL 的前向,我们调用 SGLang 已有的 Triton expert kernel,并实现训练用的自定义反向算子。XoRL 的权重以转置 view 传入,routing 权重保留在 FP32 累加器中,结果只在最后舍入一次。

Shared expert 使用稠密模型中的固定 K 投影SwiGLU。每个 rank 先将自己的 shared-expert shard 加到 routed-expert 输出上,再由八个 rank 合并结果。

代码 XoRL 被路由与 shared expert ↗

此时,两个引擎对每个 token 都持有相同的八份 BF16 部分输出(每个 rank 一份),并使用固定的 expert 输出 reduction 将其相加;八个贡献者需要三层 reduction。

0 mismatch 与性能

Wordle 采样吞吐与墙钟时间:关闭开关的 baseline、除 GDN 外已对齐的前向,以及 bitwise-aligned 前向。 每种配置的数据均为 256 条 Wordle 轨迹上去掉 warmup 后三次计时的均值。
配置吞吐采样墙钟相对关闭开关
关闭开关7,863 tok/s104.0 s1.000×
除 GDN 外已对齐5,778 tok/s134.4 s0.735×
bitwise-aligned4,767 tok/s164.2 s0.606×

Bitwise-aligned 前向使采样吞吐降低 39.4%,并使采样墙钟时间增加 57.9%。去掉精确 GDN decode 可以收回约一半的新增墙钟时间,即 49.5%。因此,GDN 是可以明确归因的最大单项开销,另一半则分布在已对齐前向的其余部分。

墙钟时间不能直接比较,因为不同策略学到的行为不同。Bitwise-aligned 运行更早开始解开游戏,因此生成的 token 数也不同。采样只是 RL 单步的一部分,因此这并不代表完整训练的实际开销;我们会在结语中进行端到端测量

GLM-5.2:稀疏注意力与 FP8 expert

GLM-5.2 又引入两项复杂因素:稀疏注意力,以及使用 FP8 原生权值的 expert。要正确训练 GLM-5.2,我们需要组合多种并行方式:Context Parallel、Expert Parallel、Tensor Parallel、Fully Sharded Data Parallel 与 Pipeline Parallel。XoRL 支持对 GLM-5.2 进行全参数、LoRAQLoRA 训练。

稀疏注意力

一些 GLM block 会选择较早的 token 位置参与 attention,后续 block 会复用这些位置。若 selector 的选择不同,mismatch 会沿序列级联。我们先对齐 selector,再对齐 attention。

选出稀疏注意力位置

Selector 返回后续稀疏 attention block 将读取的位置。XoRL 与 SGLang 构造输入的方式不同:XoRL 分别投影 key 和 head gate,再为 BF16 query、key 评分;SGLang 联合投影 key 和 head gate,并在 RoPE 之后将 query、key cast 为 FP8,各自使用独立的 scale。这些差异可能改变最终选中的位置。

我们让 XoRL 采用推理端的 selector 计算:在 BF16 中联合投影 key 和 head gate,并在 RoPE 之后进行 BF16 舍入;然后分别按各自的 scale 将 query 和 key cast 为 FP8。两个引擎都对完整历史评分,分数相同时选择较小的位置,并按升序将位置返回给后续稀疏 attention block。

Prefill kernel 写入 prompt 的 key,decode 则追加每个生成 token 的 key。Prefill 与 decode 可能在不同位置执行 normalization、RoPE 和 BF16 舍入。为了在 XoRL 中复现 serving cache,我们标记 prefill 边界,并使用对应的 SGLang prefill 或 decode 计算重建每一行。

两个引擎可以采用不同的历史存储方式,只要每个逻辑位置上的值相同。XoRL 将历史存储在连续 tensor 中,SGLang 则存储在 64-token page 中。

我们在 XoRL 中使用对应的 serving 计算重建每一行 prompt 和生成行。因此,selector 在 XoRL 的连续历史与 SGLang 的 paged 历史中读取相同的值,从而选择相同的位置。

两种内存布局上的同一个稀疏注意力 selector 我们从 XoRL 的连续历史与 SGLang 的 64-token page 构造相同的 FP8 query、key 以及相同的 FP32 head-gate。两个 selector 返回相同位置。

我们从 XoRL 调用 SGLang 的 FlashMLA 前向,并在 TileLang 中实现其反向。选中的行直接从 XoRL 的连续 tensor 传入,绕过 SGLang 的 page table。

代码 XoRL selector 输入 ↗ XoRL 被选位置 ↗ XoRL FlashMLA 前向与反向 ↗ SGLang 稀疏 selector ↗ SGLang paged 稀疏 selector ↗

Mixture-of-experts 层

在 GLM 中,router 会添加校正 bias,expert 使用 FP8 权值运行,且多个 rank(最低配置为 EP=16) 会共同产生每个 MoE 输出。Adapter 训练与全权值训练共用同一套 FP8 expert 前向。 两种模式只在反向之后如何更新 expert 权值上不同。

选出 GLM 的八个 expert

GLM 只在选择 expert 时加入校正 bias;expert 权重仍来自加入 bias 之前的 sigmoid 分数。训练端可能将 bias 加载为 BF16,而推理端保持 FP32,从而改变被选中的 expert。训练端也可能在调用 expert 之前,将经过 normalization 的 routing 权重舍入到 BF16;推理端则将其保留为 FP32,并在 expert kernel 中乘以 2.5。

为消除这些差异,我们在两个引擎中运行相同的固定 K gate 投影,并从 XoRL 直接调用 SGLang 的 router。router 在选出 top-8 expert 之前加入 FP32 校正 bias,再 gather 加入 bias 之前对应的 sigmoid 分数,在 FP32 中进行 normalization,最后传给 expert kernel;expert kernel 再乘以 2.5。

代码 XoRL GLM router ↗ SGLang 分组 top-k ↗

LoRA 路径:带 adapter 的 FP8 expert

训练端可能将 expert 权值转成 BF16 并单独计算 LoRA,而推理端读取 checkpoint 中的 FP8 权值与块 scale,并在 expert kernel 中施加 LoRA。两条前向会在不同位置舍入。

我们扩展 SGLang 的 FP8 expert kernel,在 gate/up 投影和 down 投影之后各施加一次 LoRA。基座权值保持 FP8,routing 权重乘法与最终 BF16 舍入仍保留在推理端原有的位置。我们从 XoRL 的前向调用该 kernel,并为 FP32 LoRA 矩阵实现自定义反向。

代码 XoRL FP8 expert 与反向 ↗ SGLang LoRA MoE ↗ SGLang FP8 expert ↗

全权值路径:训练 FP8 expert 权值

要训练 FP8 expert,训练端会保留 FP32 master 权值,并在前向中量化。若推理端再次量化,两个引擎就必须在量化器的每个舍入点上对齐。

我们改为在每次 optimizer step 之后,按 SGLang 的块 layout 对 XoRL 的 FP32 master 权值量化一次。在下一次 optimizer step 之前,所有前向都读取缓存的 FP8 code 与块 scale。我们将它们与上次发送的版本比较,只发送发生变化的部分。单步内每个块的 scale 变化很小,因此大约只有 1% 的条目会改变。

LoRA 或全权值前向完成后,我们将十六份部分输出按逻辑 rank 顺序排列,再施加固定的 expert 输出 reduction;十六个贡献者需要四层。

Distributed LM head

推理可以将 GLM 的 154,880-token 词表分片到十六个 rank,而训练端一次投影完整词表。形状不同的矩阵乘可能按不同 顺序累加乘积,从而产生不同的 logits。

我们将 XoRL 的 LM head 按相同的十六份词表 shard 分片,并将每一行隐状态与目标 token id 发送到持有对应 shard 的 rank。每个 rank 施加 固定 K 投影,再按 rank 顺序 gather logits,最后施加同一套流式 softmax

代码 XoRL distributed LM head ↗ SGLang distributed LM head ↗ SGLang logprob 计算 ↗

0 mismatch 与性能

使用 bitwise-aligned 前向采样时,吞吐从 147.6 降至 107.2 tok/s,降低 27.4%。

GLM-5.2 采样吞吐:未对齐 baseline、bitwise-aligned 前向,以及从中去掉一项对齐后的前向。每种配置的数据均为 TP16/EP16/CP16 上十六个样本、六十四个输出 token、三次计时的均值。
配置吞吐相对关闭开关
关闭开关147.6 tok/s1.000×
bitwise-aligned107.2 tok/s0.726×
bitwise-aligned,去掉 expert 路径对齐125.9 tok/s0.853×
bitwise-aligned,去掉稀疏注意力 selector110.2 tok/s0.747×
bitwise-aligned,去掉 batch-invariant 稠密 GEMM105.8 tok/s0.717×

额外开销主要来自 expert 路径。若关闭已对齐的 router、FP8 expert 前向,以及按 rank 排序的 expert 输出 reduction,吞吐会恢复到 125.9 tok/s;这意味着近一半的对齐开销来自 expert 路径。精确稀疏 attention selector 与 top-k 大约占两个百分点。batch-invariant 的 addmmbmmlog_softmaxmeanmm 的开销处于测量噪声范围内,因为它们不改变 expert 投影,只改变 dense matmul 与 attention 投影,而后者只占 FLOPS 的一小部分。

DeepSeek-V4:四路残差与压缩历史

DeepSeek-V4 引入多路残差流、压缩注意力历史,以及 MXFP4 expert 权值;我们使用 LoRA 训练这些 expert,且无需解包权值。

四路残差混合

融合会改变残差计算

每个 block 包含四路 BF16 residual stream。模型的 mHC mixer 在每个 attention / expert 层之前将这些 stream 合成为一行 BF16,再将该层输出加回重新混合后的 stream。这个 mixer 会为每个 token 计算 FP32 权重和一个 4×4 FP32 混合矩阵。XoRL 与 SGLang 可能在 mixer 的融合方式,以及内部舍入和相加的位置上无法对齐,从而产生 mismatch。

共用 serving 前向

在 XoRL 的前向中,我们调用 SGLang 的 pre-mix 与 post-mix kernel。自定义反向会将梯度传回全部四路 residual stream。训练时对 mixer 系数执行 stop-grad,因此没有梯度流入 mixer 权值。

pre-mix kernel 将四路 residual stream 合成为该层的输入行;post-mix kernel 再将该层输出加回重新混合后的 stream。

四路残差流如何穿过一层 融合的 pre-mix 与 RMSNorm kernel 把四路 BF16 残差流合成一行。注意力或 expert 层之后,融合的 post-mix kernel 再把结果分回四路。

代码 XoRL mHC 前向与反向 ↗ SGLang 融合 mHC kernel ↗

压缩稀疏注意力

每一层都使用 128-token 的滑动窗口注意力。C0 层看这个窗口。 C4 层还会看重叠的摘要。这些摘要覆盖八个 token,每四个 token 建一次。C128 层对每个不重叠的 128-token tile 有一份摘要。每个 query 在一次 softmax 里同时看自己的摘要与滑动窗口。

压缩前缀

要构造一行压缩结果,模型先在 BF16 中投影隐状态,再使用 FP32 权重执行 pooling 和 normalization, 然后施加 RoPE 并存为 FP8。投影采用与稠密层相同的 batch-invariant matmul 和同一份 RoPE kernel。 训练端看到完整序列,而推理在 prefill 时构造这份状态、在 decode 时更新它。因此两边可能在不同的 token 边界上创建压缩行。我们只在对应的 token 组完成时才创建 C4 与 C128 行,以便对齐 serving 边界,确保后续 query 读到正确的压缩行。

在 XoRL 中,我们使用 SGLang 的摘要计算与 FP8 cast。由于 SGLang 的稀疏 prefill kernel 不是 bitwise-stable,我们逐个因果 query 行调用 decode kernel 来回放 prompt。随后使用同一个 kernel,每次推进 paged cache 中的一个 token。在 token 之间,我们传递最近 token 的 cache、FP32 compressor state 和 C4/C128 行,并在与 serving 相同的边界添加一行压缩结果。

对压缩行做注意力

我们在 XoRL 中根据 SGLang paged cache 暴露的逻辑位置构造 FlashMLA 输入,并将压缩前缀放在最近 128 个 token 之前。我们调用 SGLang 的 FlashMLA 前向,并用同一 attention 的 FP32 参考实现反向。回放单个 token 时,直接在 paged cache 上调用 serving 的 decode 入口。

我们用 decode kernel 回放 prompt,以初始化 paged cache。随后每一步 decode 都将最近 token cache、FP32 compressor state 与压缩 cache 传到下一步。

把压缩注意力状态从 prefill 带到 decode 对 C4/C128 层,用 decode kernel 回放 prompt 来初始化 paged cache。每一步单 token decode 都会推进原始 128-token 窗口,将新 token pool 到 FP32 compressor state,并在与 serving 相同的边界把新的 FP8 摘要写入 C4/C128 cache。

代码 XoRL DeepSeek-V4 compressor ↗ XoRL C4 indexer ↗ XoRL FP8 注意力前向与反向 ↗ SGLang DeepSeek-V4 compressor ↗ SGLang C4 indexer 与 page 映射 ↗

哈希路由与学习到的 expert 路由

训练端可能使用自己的矩阵乘生成 router logits,并在 PyTorch 中计算两套路由方案;推理端则可能用 persistent kernel 生成 BF16 gate 投影,再扩展为 FP32,交给融合的哈希或学习路由 kernel。路由是离散的,因此在学习层中,gate logit 上 1 ULP 的变化就可能改变 sqrt-softplus 权重、被选中的 slot 及其顺序。顺序不同自然会导致 mismatch。

为对齐这两条路径,我们在 XoRL 中使用 batch-invariant matmul 生成 gate logits。前三个 block 使用一张冻结表,将每个 token id 映射到六个 expert,因此 logits 只决定权重;后续 block 使用带校正 bias 的学习 router。在整个 expert 前向中,我们保留返回的权重、expert id 与 slot 顺序。

代码 XoRL DeepSeek-V4 路由 ↗ XoRL 哈希与校正 bias 选择 ↗ SGLang token 到 expert 的路由 ↗ SGLang 学习路由 ↗

带 LoRA 的 MXFP4 expert

MXFP4 还有一项 FP8 不具备的约束:Marlin 在不同行数下可能返回不同的比特。因此,两个引擎每次都只运行一行 packed GEMM,与 serving 在 decode 时采用的无 padding 启动方式一致。训练端的反向将每个 expert 权值解包为 BF16 并单独计算 LoRA,而推理端则在处理 packed 权值的 Marlin runner 内部加入 LoRA。

与 FP8 expert 类似,我们扩展了 SGLang 的 runner,在 packed gate/up 和 down 投影之后施加 BF16 LoRA,全程不解包基座权值,并从 XoRL 的前向调用该 runner。我们沿用 SGLang 的行分解、clamp、SwiGLU 计算和最终 BF16 舍入,并在相加之前,用 routing 权重分别乘以 packed 基座输出和 LoRA delta。

routing 权重分别乘以 packed 基座输出和 LoRA delta;六个 slot 随后在 FP32 中累加,并舍入一次到 BF16。两个引擎使用相同的行数。

Packed MXFP4 expert 前向里的两处 LoRA 插入点 Marlin expert runner 在 packed gate/up 投影之后加入 BF16 LoRA,MXFP4 基座权值保持 packed。在 down 投影中,它在相加之前用 routing 权重分别乘以 packed 基座输出和 LoRA delta;六个被路由的 slot 随后在 FP32 中累加,并舍入一次到 BF16。两个引擎使用相同的行数与 clamp,并按返回顺序对 slot 求和。

代码 XoRL packed MXFP4 权值与前向 ↗ SGLang DeepSeek-V4 LoRA 映射 ↗ SGLang Marlin LoRA runner ↗

每个 rank 收到的行数不同,交换已经按 rank 顺序交出全部八份 rank 输出,因此我们只需用固定的 expert 输出 reduction 把它们相加。

代码 XoRL 可变行交换与 reduction ↗ SGLang expert 输出 reduction ↗

最终混合与 LM head

最后一个 decoder block 之后,mixer 产生最后一行 BF16 隐状态,再经过最后的 RMSNorm。 LM head 用词表权值投影该行,按 rank 顺序 gather TP shard,再执行 log-softmax。这三步中的任何一步,只要两个引擎采用不同的计算方式,都会导致 mismatch。

我们复用前面 LM head 中的固定 K 投影与有序 gather。SGLang 的 LoRA kernel 在普通基座 matmul 之上加入 rank-1 delta。我们按 rank 顺序 gather BF16 logits,从 XoRL 调用 SGLang 的 BF16 log-softmax,并用 FP32 log-softmax 实现 XoRL 的反向。

代码 XoRL DeepSeek-V4 LM head ↗ SGLang DeepSeek-V4 模型 ↗ SGLang BF16 log-softmax ↗

0 mismatch 与性能

我们在同一 Wordle 任务上使用 LoRA 训练 DeepSeek-V4。采样时在单个节点上使用 TP8/DP8/EP8,bitwise-aligned 前向使 decode 吞吐降低 12.7%。

DeepSeek-V4 采样吞吐:关闭开关、使用 bitwise-aligned 前向,以及从中去掉已对齐的 expert 几何与 combine;配置为 TP8/DP8/EP8,基座权值冻结。每组数据均为八个请求、六十四个输出 token、三次计时的均值。
配置吞吐相对关闭开关
关闭开关46.0 tok/s1.000×
bitwise-aligned40.2 tok/s0.874×
bitwise-aligned,去掉 expert 路径对齐42.5 tok/s0.925×

额外开销仍主要来自 expert 路径:去掉已对齐的 expert 几何与 combine,可以收回 43.8% 的新增墙钟时间。在三个 MoE 中,DeepSeek-V4 实现 bitwise alignment 的成本最低。

代码 XoRL DeepSeek-V4 PR ↗ SGLang DeepSeek-V4 PR ↗

代价与收益

我们已经能以 0 train-infer mismatch 训练 Qwen 3.x、GLM-5.y 与 DeepSeek-V4 这类 MoE,也有一套易于部署并可扩展到其他模型的方法。但为什么值得这样做?请注意:从这里开始,我们将从可直接核验的 kernel 转向更依赖经验结果的 RL 实验。 代价是多少?Async RL 会受到什么影响?更换目标函数是否足以缓解 mismatch?是否存在更好的替代方案?接下来,我们会在 Qwen3.6-35B-A3B 的 Wordle 任务上进行一组 ablation,尝试回答这些问题。我们还会与托管 RL API TinkerRiver 对比,评估 baseline 的 mismatch 水平和最终模型表现是否与闭源 RL 引擎相当。当然,我们看不到这些引擎的内部细节,因此不能对它们的表现下定论。

为什么要消除不匹配?

要理解为什么应该消除 mismatch,先看 Tinker 上的运行。图 1 展示的是三次运行中表现最好的一次;另外两次都在第 62 步发散。

显示
  • Tinker
    Run 1
  • Tinker
    Run 2
  • Tinker
    Run 3
从同一个第 56 步 checkpoint 出发的三次 Tinker 运行:solve rate、K3 与训练端平均 logprob Run 1 与 Run 3 的 K3 在第 62 步升至 0.2466 与 0.2520,同时训练端平均 logprob 降至 −0.461 与 −0.477。该步采样端平均 logprob 在三次运行中都维持在约 −0.20。Run 2 的 K3 保持在 0.0178,训练端平均 logprob 为 −0.215。随后 Run 1 与 Run 3 发散;Run 2 完成训练,并在 850 局 held-out 评测中解出 418 局。图例中的每一项都可以开关对应曲线。

从第 56 步同一 checkpoint 出发的三次 Tinker 运行中,有两次在第 62 步发生灾难性发散。

第 62 步本身并无特殊之处。就我们所见,Tinker 的基础设施没有系统错误。 三次运行中,采样端选中 token 的平均 logprob 都维持在约 −0.20,采样得到的 rollout 看起来也合理。但在发散的两次运行中,训练端 logprob 在第 62 步明显更低,从而导致灾难性的 mismatch。

由于训练采用 importance sampling,这些比值对梯度的影响没有上界。在 Run 1 中,仅一个 importance ratio 为 19,276.5 的 </ token 就占平方系数质量的 65.7%。每个 token 的更新系数是 c = ρÂ,即 importance ratio 乘以 GRPO advantage。平方系数质量是该 token 的 c2 占整个 batch 的比例。 Run 3 没有这么极端,但前四个 token 仍占 72% 的质量。Run 2 则不然:第 62 步最大的单个 token 只占 0.39% 质量。

Run 1 失去了格式约束,大多数回合开始输出非法动作。Run 3 仍能执行合法的 Wordle 步骤,但猜测质量很差。Run 2 没有发散,不过最终表现也不理想。

当然,无上界的 importance sampling 更像是刻意放宽的 baseline;我们可以通过调整目标函数来修正这个问题。

什么目标函数能缓解不匹配?

显示
  • 0 mismatch
  • Unclipped IS(不截断 IS)
  • PPO clipping(PPO 截断)
  • CISPO
  • Unclipped IS
    (Tinker)
  • CISPO
    (Tinker)
不截断 importance sampling 与 PPO clipping 的 Wordle 运行:solve rate、train-infer mismatch 与 held-out 评测 精确 0-mismatch 运行在第 128 步达到 70.8% 训练 solve rate 和 77.4% held-out solve rate,全程 K3 精确为零。关闭开关时,不截断 importance-sampling 运行在第 128 步达到 58.1% 训练 solve rate 和 63.9% held-out solve rate,后期 K3 约 0.017。PPO clipping 运行达到 21.9% 训练 solve rate 和 20.2% held-out solve rate,后期 K3 约 0.00068。CISPO 运行在第 128 步以 67.4% 训练 solve rate 结束,低于第 58 步 80.5% 的峰值,后期 K3 约 0.012,held-out solve rate 为 72.5%。在 Tinker 上,不截断 importance-sampling 达到 51.4% 训练 solve rate 和 49.2% held-out solve rate,而 CISPO 达到 67.1% 训练 solve rate 和 75.5% held-out solve rate,后期 K3 约 0.0083。图例中的每一项都可以开关对应曲线;0-mismatch 曲线默认隐藏。

文章开头简要讨论了不同目标函数如何处理 importance sampling ratio。展开下方内容可以回顾 IS、PPO 与 CISPO 的定义。

Importance sampling、PPO 截断与 CISPO

ρt = qt / μt 为训练端对所采样 token 的概率,与推理在采样时记下的概率之比, Ât 为 detach 后的 group-relative advantage。三个目标函数的差别在于如何处理乘入梯度之前的 ρt。它们都不对 ρt 求导;我们求的是 log qt 的导数。

不截断的 importance sampling

不截断的 importance sampling 把 token 梯度乘以 ρt。我们把构成该比值的 选中 token logprobs detach 掉,只对 log qt 求导。

∇Lᴵˢₜ = −ρₜ Âₜ ∇ log qₜ; Eμ[‖gₜ‖²] = Eμ[ρₜ² Âₜ² ‖∇ log qₜ‖²]

梯度的二阶矩随 importance ratio 的平方缩放。我们在 有限大小的 batch 上训练,因此会出现高方差步:梯度被少数 大比值 token 主导。

截断后的 importance ratio

GRPO 通常用 PPO 目标来截断这一比值:

Jᴾᴾᴼₜ = min(ρₜ Âₜ, clip(ρₜ, 1 − εlow, 1 + εhigh) Âₜ); ∇Lᴾᴾᴼₜ = −1[not clipped] ρₜ Âₜ ∇ log qₜ

对优势为正的 token,一旦 ρt 超过 1 + εhigh,PPO 就舍弃策略梯度。对 advantage 为负的 token,一旦 ρt 低于 1 − εlow,同样舍弃梯度。

CISPO 在乘入梯度之前截断 ρ,从而限制超过阈值的 token 的贡献,但并不舍弃其梯度。它的边界是绝对比值,而不是围绕 1 的偏移;默认 关闭下界,因此比值只会被从上方封顶。

ρ̄ₜ = clip(ρₜ, ρmin, ρmax); ∇Lᶜⁱˢᵖᵒₜ = −sg[ρ̄ₜ] Âₜ ∇ log qₜ; (ρmin, ρmax) = (0, 4)

当每个 token 都有 ρt = 1 时, PPO 不会截断,CISPO 也不会封顶,更新就退化成上面的 on-policy 梯度。

XoRL 上的 IS 比 Tinker 上的 IS 表现更好,但仍落后于 bitwise-aligned 运行。 我们使用与 Tinker 相同的形式。 训练后期,IS 每个回合比 bitwise-aligned 运行多生成 44% 的 token。 这些 token 大多不是真正的推理,而是模型在下一次猜测前反复检查格式、词数和准备状态。 如前所述,大比值引发的大梯度会积累在 optimizer state 中,并影响后续许多步。 我们发现绝大部分质量来自 IS 比值大于 1 的 token(也就是训练端给出的似然高于采样端)。

PPO 运行的表现要差得多。我们使用与 Tinker 相同的 clipping threshold。 它丢弃了太多重要 token 的梯度,几乎无法学到有效策略。 测得的 mismatch 虽然远低于 IS,但主要原因是策略只会产生退化的 rollout。

CISPO 正是为了解决这些问题:它不允许 IS ratio 无限增大,也不丢弃梯度,而是限制其贡献。我们使用默认值。实验表明,CISPO 能有效缓解 mismatch。XoRL 上的 CISPO 运行达到 72.5% held-out solve rate,明显高于不截断 IS 的 63.9%,但仍低于 bitwise-aligned 运行的 77.4%。CISPO 在 Tinker 上的效果也很显著,将表现提升了 25%,几乎达到 0-mismatch 运行的水平。按报告的 K3 比较,Tinker 上的 CISPO 也远低于 IS baseline。

在我们尝试的目标函数中,没有一个能完全达到 0-mismatch 运行的水平,尽管 CISPO 已经相当接近。不过,经过仔细调参的目标函数仍有可能近乎完全校正 mismatch。

Total Router Recall

显示
  • 0 mismatch
    (XoRL)
  • Total Router Recall
    (XoRL)
  • Replay IDs
    (River)
  • Baseline
    (River)
  • Baseline
    (XoRL)
Bitwise-aligned、Router Replay 与 baseline Wordle 运行的训练 solve rate、train-infer mismatch 与 held-out 评测 0 mismatch 运行在第 128 步达到 70.8% 训练 solve rate 和 77.4% held-out solve rate,全程 K3 为零。Router Replay 运行达到 72.1% 训练 solve rate 和 77.2% held-out solve rate,后期 K3 约 0.00027。只回放捕获到的 expert id 的 River 运行达到 67.4% 训练 solve rate 和 74.8% held-out solve rate,后期 K3 约 0.0087。不做 Router Replay 的 River baseline 达到 68.4% 训练 solve rate 和 76.4% held-out solve rate,后期 K3 约 0.017。XoRL baseline 达到 58.1% 训练 solve rate 和 63.9% held-out solve rate,后期 K3 约 0.017。图例中的每一项都可以开关对应曲线;0 mismatch 曲线默认隐藏。

我们的大多数实验都在 MoE 上进行,因为 train-infer mismatch 对 MoE 的影响可能比对稠密模型更严重。 在网络的大部分位置,mismatch 往往平滑累积:这里相差几个比特,那里再相差几个比特。 但 MoE 层包含离散选择。一旦 router 的 top-k 在训练端与采样端选出不同的 expert,mismatch 就会显著放大。 我们推测,这可能正是导致部分 Tinker 运行失败的模式。

先前工作提出了一种方法:Router Replay。 它在推理时记录被选 expert 的下标,再将其传给训练引擎。Group Sequence Policy Optimization 描述了 Qwen 更早的 Routing Replay 策略,Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers 提出 Rollout Routing Replay。 这一做法见于 GSPOR3。但在上面的两次 River 运行中,我们没有观察到只回放 router 下标带来的收益。 Router Replay 的确能略微降低 mismatch,但如果用它挽救 mismatch 很高的 rollout,问题就会显现。假设推理引擎给 expert 0 较高的分数,而训练端给出的分数较低;即使回放相同的 expert 0,仍会出现严重 mismatch:expert 下标一致,routing 权重却不同,因为训练端会根据自己的 logits 重新计算权重。因此,只回放下标在 mismatch 很小时几乎不起作用,在极端 mismatch 时甚至可能有害,只有在中等区间才可能有帮助;总体效果预计会相互抵消。

Cursor 在训练 Composer 2 时改进了 Router Replay:过滤掉训练端 router 给出异常低 gating score 的回放 expert,并用训练端 top-k 中的候选替换。根据其报告,这能降低基础回放残留的 p99 数值 mismatch。这是一种合理的策略;下面再看第三种方案。

我们将 Router Replay 扩展为同时记录并传输 routing 权重。R3 作者特意只传输下标,以保持 router 可训练;这是一个值得肯定的目标。 SGLang 为每个请求返回两份 payload:被选 expert id(整数),以及经过 normalization 的 FP32 routing 权重;每个 token、每个 MoE 层、每个被选 expert 各有一份。XoRL 直接使用传来的权重,不再重新计算。这份 payload 随 token × MoE 层 × 被选 expert 增长,因此我们通过 handle,经 Mooncake 或文件系统传输,而不将其内嵌到 rollout 记录中。

我们将扩展后的 Router Replay 称为 Total Router Recall。它的效果很好:held-out accuracy 几乎与 bitwise-aligned 运行相同,且 K3 全程维持在 3e-4,而没有像 baseline 那样持续上升。

但 Total Router Recall 也有代价:我们无法再按常规方式训练 router。 技术上仍可用 stop-gradient 技巧保留梯度。 若前向传入 w − sg[w] + w*,前向值仍是采样端权重 w*,反向则是 ∂w/∂θ。但后续层使用的是 w*,因此 ∂L/∂w 在回放的 router expert 上求值,而 ∂w/∂θ 在训练端 logits 上求值。两个因子来自不同的 router 选择,因此实际优化的是采样端权重附近 loss 的线性化,梯度误差项相对于两个 router 的差距是一阶的。 训练端现在也与采样端的内部实现耦合:它需要获取采样端的 expert id 和经过 normalization 的权重,并按照采样端的 layout 覆盖每个 token、每个 MoE 层的数据。 每个请求从采样端传到训练端的数据量都很大。Mooncake 能缓解一部分传输压力,但捕获并跨设备发送这些 metadata 仍会带来不可忽视的开销。

我们通过一套与 Tinker 兼容的接口在 XoRL 中运行实验:先启动训练与 采样服务器,再由 轻量客户端 向它们发送采样和训练请求。 我们不希望每个请求都将大量 router metadata 从采样端经客户端再送回训练端。 但这也意味着,如果客户端修改了 token / logprobs,例如删除第一个答案之后的 token,就必须对 routing payload 做出相应修改。

到这里可以看出,这仍是一种权宜方案。Tinker 甚至不支持基础的 Router Replay。 尽管如此,Total Router Recall 仍是我们目前找到的、最接近 bitwise-aligned 系统的替代方案。 如果需要一种容易接入、能够缓解 mismatch,且不要求训练 router 的方法,它是一个不错的选择。 任何能够替代 0-K3 的方案都有吸引力,因为它可能避免 0-mismatch 方案的额外开销。

0-mismatch RL 要付什么代价?

前文已经讨论过在采样中强制 0 mismatch 的成本。从完整 RL step 来看,0 mismatch 带来的额外成本约为 20%。

Qwen3.6-35B-A3B 上关闭开关的 CISPO 运行与 zero-K3 运行的 Wordle 生产步时间,对 100 步取平均。右列两边都去掉了权值同步。
运行每步更新秒数去掉权值同步
关闭开关,CISPO374.3 s334.8 s
Zero K3447.0 s409.9 s

20% 是相当可观的额外开销,但需要注意,这项开销在 Qwen 上最为明显,主要来自用精确 GDN decode 对齐 SGLang 内部的 prefill 与 decode。

未来仍有许多优化可以降低单步开销;我们预计,0-mismatch RL 的成本会很快下降。我们认为这项投入值得,因为它能排除 mismatch 导致训练平台化的可能性。不过,强制 0 mismatch 还有另一项代价。

每个采样条表示一条轨迹。所有资源都在工作,但权值更新时仍在生成的轨迹会继续读取过期 cache。

Async RL 日程、其权值版本与过期 cache 采样与训练两条车道持续忙碌。轨迹在比训练端正在产出的版本更旧的权值下采样;跨越一次权值同步的轨迹会继续读在上一版权值下写出的 cache。

尽管我们投入了很多工作来实现 0 train-infer mismatch,Async RL 对生产规模的运行仍然很有吸引力。但它会因过期权重与过期 cache而引入显著的 mismatch。

显示
  • 0 mismatch
  • Async on 0-mismatch base + CISPO
  • Async on 0-mismatch base
    + Router Recall + CISPO
同步与流水线化的 0-mismatch Wordle 运行:训练 solve rate、train-infer mismatch 与 held-out 评测 同步 0 mismatch 运行在第 128 步达到 70.8% 训练 solve rate 和 77.4% held-out solve rate,全程 K3 为零。使用同一套 bitwise-aligned kernel、但带过期权值与 cache 的 PipelineRL 与 Streaming 运行达到 50.4% 训练 solve rate 和 59.4% held-out solve rate,来自过期 state 的 K3 峰值约 0.108。加入 Total Router Recall 后,该流水线运行恢复到 74.7% 训练 solve rate 和 76.6% held-out solve rate,后期 K3 约 0.0002。图例中的每一项都可以开关对应曲线。

即使在 0-mismatch 栈上使用 CISPO——这是我们找到的缓解 mismatch 效果最好的目标函数——异步仍会严重降低表现。过期权值与过期 cache 会使 mismatch 迅速升高,模型很早就进入平台期。加入 Total Router Recall 后,可以恢复大部分表现,同时保留 Async RL 的加速收益。在 TRR + CISPO + Async 运行中,报告的 K3 仍会在早期升高,但随后持续下降。

Streaming / Async RL 的目标是让所有训练 GPU 与推理 GPU 始终保持满负载。Streaming RL 可以在不引入任何 mismatch的情况下,让训练与推理部分重叠。

SGLang 每完成一条轨迹就立即返回,客户端组装出完整 group 后再交给 XoRL。XoRL 对每个完整 group 执行一次前向和反向——因为只需在 group 内对 advantage 进行 normalization——并在最后统一执行一次 optimizer step。因此,得到的梯度更新与将全部数据组成一个大 batch 一次处理完全相同。由于等待时间服从几何分布,训练端与采样端可以实现较好的重叠,而无需采用异步更新。

每条轨迹都在同一权值版本下运行。一个 group 在最后一条轨迹完成时立即交接,因此空闲只出现在第一个 group 到达前和最后一个 group 完成后。

单一权值版本下的 Streaming RL 日程 每条轨迹都在同一权值版本下采样。已完成的组在后续组仍在生成时就交给训练端,因此训练端只在第一组到达前空闲,采样端只在最后一组期间空闲。最后一次优化器步进与权值同步复现「一个大 batch」的梯度,节省 46% 的串行墙钟。
Streaming RL,以及与 PipelineRL 组合后的 Streaming RL,在热循环上的墙钟时间;取我们 Wordle 运行第 2、3 步的均值,每步 1,024 条 rollout。每秒 rollout 数是 1,024 除以循环时间。
模式热循环Rollouts/s相对 Streaming RL
Streaming RL338.5 s3.03基线
PipelineRL 与 Streaming RL244.0 s4.201.39×

在 Wordle 实验中,Streaming RL 将串行墙钟时间缩短了 46%。但随着更多轨迹完成,采样端利用率会逐渐下降,因此加速效果仍不及 Async RL。

Async RL 用过期策略的新请求填补异构序列长度造成的空隙,但代价是引入 mismatch。另一种做法是通过 multi-LoRA 训练,用不同策略的请求填补这些空隙。当一个任务只剩少量耗时较长的轨迹时,另一个任务可以用新的 batch 填充采样队列。每个请求相对于自己的 adapter 仍是 on-policy,因此这种重叠不会引入 mismatch;前提是这些任务共享同一个基座模型。这不会缩短单个任务的训练时间,但如果一套 XoRL 训练服务和一套 SGLang 采样服务需要同时运行多个实验,multi-LoRA 可以有效提高 GPU 利用率。

结语

我们相信,在开源权重 MoE 上实现 0 mismatch,是任何 RL 栈都能够达到的目标,而且收益明确。这些技术目前都已用于 XoRL,也可以移植到其他栈。虽然 XoRL 目前使用 SGLang 进行推理,但同样的方法也适用于其他推理引擎。

From: Blog of Ashwinee Panda
with help: Ashwinee Panda, Qingyang Wu, Zhongzhu Zhou.

致谢:感谢 Core Automation、Recursive、Ricursive 与 River AI 的朋友审阅初稿并惠赐意见。TogetherAI 提供的 GPU 机时与 token,使本研究得以完成。

附录 A:算子对照表

显示每一个算子
我们如何对齐每一个算子,以及什么决定输出比特。
算子方法决定输出比特的因素
共享机制
RoPE构造同一张表并对齐旋转同一张位置表、同一套 cast 与旋转算术
RMSNorm执行同一套 batch-invariant RMSNorm 计算residual 的相加位置、固定的 FP32 reduction 顺序、rsqrt、FP32 scale 与最终 cast
GEMM使用我们的 Triton 参考实现,或输出相同的库 kernel沿 K 按同一固定顺序累加乘积,不用 split-K
SwiGLU在两个引擎中执行同一个、只舍入一次的 SwiGLU kernelFP32 SiLU 与乘法,再执行一次 BF16 输出 cast
Attention两个引擎都调用 FlashAttention 4相同的 backend、构建和参数,以及对 key/value block 执行一次 reduction
LM head对齐投影与流式 reduction同一 GEMM 累加顺序、256-logit 分组与摘要合并树
LoRA(合并)在两个引擎中构造相同的合并权值字节投影前相同的 W + sBA 字节
LoRA(分离)从 XoRL 调用 SGLang 的 LoRA 前向并实现反向传给 SGLang 前向的同一对 BF16 A、B 矩阵
采样变换将采样参数传给训练端并回放变换采样前按同一顺序施加相同的 temperature、top-k、min-p 及其他变换
Qwen3.5/3.6 组合
Gated DeltaNet对齐训练、prefill、eager decode 与 CUDA graph 上的 state同一套 normalization 与卷积、state 方向与三角 merge、64 行计算、FP32 边界 state 与行 cache
Qwen router对齐投影与被选权重的 normalization累加 gate 投影并对被选权重重新进行 normalization 时所使用的顺序
Qwen expert 前向从 XoRL 调用 SGLang 的 Triton expert 前向并实现反向FP32 down-projection 累加器中的 routing 权重,以及一次最终 BF16 舍入
Qwen expert 输出相加对各 rank 输出执行 batch-invariant reduction按逻辑 rank 顺序排列的同一八份 BF16 部分输出、同一套括号结构,以及在 FP64 中累加后舍入一次到 BF16
GLM-5.2 组合
GLM 稀疏 selector用对应的 SGLang prefill 或 decode 计算重建每一行 cache相同的编码 query、key 与 scale、合法历史、平局规则与被选位置
GLM 稀疏注意力调用 FlashMLA 并实现反向按升序逻辑顺序的同一批被选行,以及同一份 FlashMLA 前向
GLM router从 XoRL 调用 SGLang 的 router选择八个 expert 时使用的同一固定 K gate 投影与校正 bias,以及同一套 FP32 权重 normalization
GLM expert 前向从 XoRL 调用 SGLang 扩展后的 FP8 expert 前向并实现反向相同的 FP8 权值与 scale、BF16 LoRA 矩阵以及 routing 权重的放置位置
GLM expert 输出相加对十六个 rank 执行同一套 batch-invariant reduction按逻辑 rank 顺序排列的同一十六份 BF16 部分输出、同一套括号结构,以及在 FP64 中累加后舍入一次到 BF16
GLM LM head投影同一十六份 shard 并按 rank 顺序 gather相同的输入行、十六份词表 shard、投影 kernel、按 rank 排序的 gather 与最终概率 reduction
DeepSeek-V4 组合
DeepSeek-V4 残差混合从 XoRL 调用 SGLang 的 pre/post-mix kernel 并实现反向从四路合成一路时使用的同一套 pre-mix、RMSNorm、post-mix 与最终 mix
DeepSeek-V4 压缩在 prefill 与 decode 中执行 SGLang 的 compressor相同的 C4 与 C128 行、FP32 compressor state、block 边界与 FP8 cache 存储
DeepSeek-V4 压缩注意力在 prefill 与 decode 中调用 FlashMLA 并实现反向相同的 prefill 与 decode 入口,以及按相同顺序排列的压缩行与最近行
DeepSeek-V4 路由从 XoRL 调用 SGLang 的 gate 投影与 selector相同的 token 查找或固定 K gate 投影、sqrt-softplus、校正 bias、top-k slot 顺序与权重 normalization
DeepSeek-V4 expert 前向从 XoRL 调用 SGLang 扩展后的 Marlin expert runner 并实现反向相同的 packed MXFP4 权值、LoRA 插入点、行数、clamp、routing 权重放置与 BF16 输出求和
DeepSeek-V4 expert 输出相加在可变行交换后执行同一套 batch-invariant reduction相加前恢复到逻辑 rank 顺序的同一八份源 rank 输出
DeepSeek-V4 LM headGather 词表 shard,并从 XoRL 调用 SGLang 的 log-softmax同一套最终 mix、按 rank 排序的 BF16 词表 gather,以及 batch-invariant BF16 log-softmax
并行
Ulysses attention通过 all-to-all 原样传输这些值通过 all-to-all 传输的同一批 BF16 值,以及任意 head 数下同一套按 head 算术
Context-parallel Gated DeltaNet在 64-token chunk 网格上切分并传递 FP32 边界 state每个切分点上的相同 chunk 网格、边界 state 与卷积输入
Pipeline stage在 decoder 层之间切分传给下一 stage 的 BF16 行,与下一层原本会读取的行完全相同