论文
大规模、长上下文 RL 中的 推测解码 在线草案联合训练 后训练
Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training
摘要
推测解码 加速了部署生成，这在强化学习 (RL) 后训练 的成本中占主导地位。在线协同训练可以进一步提高草稿的准确性，从而实现更大的加速。然而，将这种方法扩展到具有长上下文的大型模型上的协同训练会带来两个障碍：（1）标准因果上下文并行（CP）实现不支持分支注意力，（2）目标特征跨越管道并行（PP）阶段。我们通过大规模在线选秀协同训练的端到端系统来解决这两个问题。对于 CP，我们通过将排名局部分支注意力与因果主序列注意力合并来扩展打包、负载平衡的锯齿环注意力。对于 PP，TapChannel 通过单独的路径跨阶段传输中间目标特征，使管道调度不受影响。实验表明，共同训练的草案密切跟踪策略基线，同时在高达 122B 的模型规模上实现大幅部署和端到端加速。我们的 CP 设计实现了 256K 词元的强大扩展，与之前的工作相比显着节省了内存，并且我们的 PP 传输产生了适度的开销。代码可以在 https://github.com/NVIDIA-NeMo/RL/issues/3698 找到。