4 min read#speculative-decoding#mtp#eagle#rl#slime#draft-model

Online Training of EAGLE MTP for RL Speculative Decoding

Apr 7, 2026

投机解码(Speculative Decoding) 用轻量级 draft model 先生成多个 token，再由 target model 批量验证。验证通过的 token 直接接受，验证失败的用 target model 重新采样。由于验证是并行的，理想情况下一次前向就能接受 $k$ 个 token。

在 RL rollout 中直接应用 speculative decoding 有一个核心矛盾：RL 训练会持续改变 target model 的策略，导致 target model 和 draft model 之间的分布差异不断加大，accept rate 逐步衰减。slime 的解决方案是在训练过程中同步在线更新 draft model——让 MTP 层随主模型一起学习，始终保持较高的 accept rate。

EAGLE MTP 的预测目标

主流的 speculative decoding draft model 采用 EAGLE MTP(Multi-Token Prediction) 架构。与标准自回归模型用 $t$ 时刻输入预测 $t+1$ 时刻 token 不同，MTP 的预测目标是 $t+2$ 时刻的 token：

$\text{Input}(t) + \text{Input}(t+1) \to \text{Output}(t+2)$

具体来说，MTP 层接收两个输入：

target model 在 $t$ 时刻的隐状态(hidden state) $h(t)$ ：代表主模型在 $t$ 时刻的”思考”
$t+1$ 时刻的 token embedding $e(t+1)$ ：代表下一个位置的输入信息

MTP 层需要根据这两个输入，预测 $t+2$ 时刻的 token。

以 target model 生成序列 [a, b, c, d, e] 为例：

位置	hidden state	token embedding	MTP 预测目标
0	$h(a)$	$e(b)$	$c$
1	$h(b)$	$e(c)$	$d$
2	$h(c)$	$e(d)$	$e$
3	$h(d)$	$e(e)$	(无效)
4	$h(e)$	(无效)	(无效)

这里 token embedding 的构造通过 roll_tensor 实现——把 target model 的 token 序列整体左移一位后过 embedding layer。

在线训练流程

前向计算

MTP 层的在线训练与主模型的 GRPO 训练在同一个 backward() 中完成，但有两条独立的 loss flow：

主模型：正常做 forward → 拿到 log probs → 在 Megatron 外部计算 GRPO loss
MTP 层：在 Megatron 内部拿到主模型的 hidden states 和 token embeddings → MTP forward → 通过共享的 lm head 得到 logits → 计算 cross-entropy loss

# MTP 的前向流
draft_hidden_state = mtp(concat([token_embedding, hidden_state]))
mtp_logits = shared_lm_head(draft_hidden_state)

labels = roll_tensor(roll_tensor(tokens, shift=-1), shift=-1)  # 左移两位
mtp_loss = cross_entropy(mtp_logits, labels)

调用 backward() 时同时触发 GRPO loss 和 MTP CE loss 的反向传播。

梯度隔离

关键设计：MTP 的 CE loss 只更新 MTP 层自身的参数，不能反传到主模型的 hidden states、embedding 层或 lm head。否则辅助 loss 会干扰主模型的 GRPO 训练。

具体做法有三处切断：

1. 切断 hidden state 和 embedding

decoder_input = embedding(input_ids).detach()       # embedding 输出 detach
hidden_states = make_viewless_tensor(hidden_states,
    requires_grad=True, keep_graph=False)            # hidden state detach

对 decoder_input 做 detach() 后，前向数值不变（仍用主模型的 embedding 权重），但 MTP loss 的梯度不会传回 embedding 权重。hidden states 同理。

2. 切断 lm head

使用 torch.func.functional_call 把 lm head 当作纯函数调用，传入 detach 过的参数副本：

output_layer_params = {k: v.detach() for k, v in self.output_layer.named_parameters()}
mtp_logits, _ = torch.func.functional_call(
    self.output_layer,
    {**output_layer_params, **output_layer_buffers},
    (mtp_hidden_states,),
    {"weight": output_weight.detach(), ...}
)

functional_call 的语义是：把 nn.Module 当成 $f(\text{params}, \text{inputs})$ 的纯函数，这次调用用一份 detach 过的参数视图，不修改模块本身的 requires_grad 状态。所以主 LM loss 仍可正常更新同一个 lm head，两条 loss flow 互不干扰。

这一设计的整体效果是：

参数	GRPO loss 更新	MTP CE loss 更新
主模型参数	是	否（hidden state detached）
Embedding	是	否（detached）
LM Head	是	否（functional_call + detach）
MTP 层参数	否	是

Loss Mask 处理

RL 训练中，只有模型自己生成的 token 才参与 loss 计算（见 multi-turn-rl）。MTP 的 loss mask 需要额外处理，因为 MTP 在位置 $t$ 的预测依赖 $t+1$ （中间 token）和 $t+2$ （目标 token）两个位置的有效性。

以 loss mask [1, 0, 1, 1, 0] 为例：

位置 1 的 mask 为 0，说明 $b$ 不是模型生成的 token
MTP 在位置 0 要根据 $h(a)$ 和 $e(b)$ 预测 $c$ ——但 $b$ 无效，这个预测没有意义
因此需要同时检查 $t+1$ 和 $t+2$ 位置的 mask

loss_mask_1 = roll_tensor(loss_mask, shift=-1)   # [0, 1, 1, 0, ❌]  (t+1 有效性)
loss_mask_2 = roll_tensor(loss_mask_1, shift=-1)  # [1, 1, 0, ❌, ❌]  (t+2 有效性)
mtp_loss_mask = loss_mask_1 * loss_mask_2          # [0, 1, 0, ❌, ❌]  (取交集)

只有 $t+1$ 和 $t+2$ 都有效的位置才计算 MTP loss。注意位置 $t$ 本身是否有效不影响——因为 MTP 用的是 $h(t)$ 而不是 $t$ 位置的 token。

Sequence Packing 适配

为了最大化训练吞吐，slime 把一个 batch 的多条序列用 sequence packing 拼成 bs=1 的连续张量，消除 padding 浪费。这要求 roll_tensor 操作对每条序列独立执行——不能跨序列边界做平移。

在开启 上下文并行(Context Parallelism, CP) 的场景下，CP 切分先于 sequence packing 进行：先按 CP 规则把每条序列补齐、切分到各 rank，再在 rank 内做 packing。因此每个 rank 内的 roll_tensor 逻辑和非 packing 场景几乎一致：对每条序列分别平移即可。

实验效果

在 H200 集群、Mimo-7B-RL 模型、DAPO-Math-17k 数据集（max response length = 24k）上的测试结果：

配置	rollout 吞吐 (tokens/s)	长尾效率 (tokens/s)	训练耗时 (s)
Online SFT spec	1667	65	148
Frozen spec	1464	57	143
No spec	1219	47	141

关键观察：

Online SFT vs No spec：rollout 吞吐提升约 35%，训练耗时仅增加约 5%（MTP 层的额外训练开销很小）
Online SFT vs Frozen spec：整体提升约 14%，训练后期差距扩大到约 25%——因为 frozen draft model 的 accept rate 随主模型策略漂移而持续下降，而 online SFT 的 accept rate 稳步上升
训练效果：引入 speculative decoding 不影响主模型的 reward 曲线——这符合理论预期，因为 speculative decoding 在概率上和 target model 直接采样完全一致

一个实用理解

冻结 draft model 在 RL 中不可行——主模型策略在不断变化，accept rate 会快速衰减
在线 MTP 训练的核心技巧是梯度隔离：让两条 loss flow 共享前向计算（hidden states、embedding、lm head），但在反向传播时互不干扰
这套方案的额外开销很小（训练时间 +5%），但吞吐收益显著（+35%），因此在 rollout 是瓶颈的 RL 系统中值得默认开启