首个开源实现100%可复现的稳定RL训练框架来了！2次结果完全重合

衡宇
2025-09-27
11:30:07

来源：量子位

SGLang团队 投稿

量子位 | 公众号 QbitAI

开源框架实现100%可复现的稳定RL训练！

下图是基于Qwen3-8B进行的重复实验。两次运行，一条曲线，实现了结果的完美重合，为需要高精度复现的实验场景提供了可靠保障。

这就是SGLang团队联合slime团队的最新开源成果。

近期，Thinking Machines Lab (由 OpenAI 前 CTO Mira Murati 创立) 发布了一篇文章——《克服LLM推理中的不确定性》，指出问题的核心在于缺乏批次不变性 (batch invariance)。

自从这篇博客发布后，业界反响热烈，一直期待开源推理引擎能实现稳定可用的确定性推理，或者更进一步，实现完全可复现的RL训练。而现在，SGLang和slime一起给出了答案。

SGLang团队在Thinking Machines Lab发布的批次不变算子基础之上，通过定制一系列注意力算子和采样逻辑，实现了完全确定性推理。该实现同时保持与分块预填充 (chunked prefill)、CUDA Graph、Radix Cache 和非贪婪采样 (non-greedy sampling) 等关键功能的兼容性。

利用CUDA Graph，SGLang可以实现2.8倍的加速。与 Thinking Machines Lab博客中报告的61.5%性能下降相比，SGLang在FlashInfer和FlashAttention 3后端平均仅有34.35%的性能下降。

在此基础上，SGLang团队联合slime团队一起，以极少代码实现了完全可复现的稳定RL训练，具体细节现在全面公开。

问题的本质：为什么LLM推理不确定？

从大语言模型(LLM)推理中获得一致输出的能力越来越重要。例如，推理结果的不确定性可能会隐式地将在线策略强化学习(RL)转换为离线策略RL。然而，即使将温度降到0，由于使用了动态批处理(Dynamic Batching) 和基数缓存(Radix Cache)，采样仍然不是确定性的。

Thinking Machines Lab的研究发现，不确定性的最大来源是变化的批次大小：即使用户重复提交相同的提示，输出也可能在不同的运行中有所变化，因为请求可能与其他用户的请求一起批处理，批次大小的差异导致不确定的推理结果。

更具体地说，不同的批次大小会影响内核的归约（reduction）分割过程。这导致每个归约块的顺序和大小变化，由于浮点运算的非结合性，可能导致不确定的输出。为了解决这个问题，他们实现了批次不变(batch invariant) 的归约算子（RMSNorm、Matmul）。这些算子也作为配套库发布供外部集成。

在Thinking Machines Lab工作的基础上，SGLang提供了强大、高吞吐量的确定性 LLM 推理解决方案，将批次不变(batch invariant)算子、CUDA 图、基数缓存和分块预填充与高效性能相结合。通过全面的测试和强化学习训练实验，其确定性得到了充分验证。

主要增强功能包括：

• 集成Thinking Machines Lab的批次不变(batch invariant)算子。
• 实现固定KV分割大小的批次不变注意力算子。支持多种后端，包括 FlashInfer、FlashAttention 3和Triton。
• 与关键推理性能相关功能完全兼容，例如分块预填充、CUDA图、基数缓存等，当启用确定性推理时，所有这些功能都仍受支持。
• 支持按请求设置采样种子(per-request sampling seed)，即使在temperature>0的非贪婪采样模式下也能实现确定性推理。

实验结果

确定性评估

引入了一个确定性测试，以验证推理结果在不同批处理条件下是否保持一致。该测试包含三个子测试，难度循序渐进：

• 单一（Single）：在不同批次大小下运行相同的提示，检查输出是否保持一致。
• 混合（Mixed）：在同一批次中混合不同类型的提示（短提示和长提示），并验证一致性。
• 前缀（Prefix）：使用来自同一长文本但前缀长度不同的提示，随机进行批处理，并测试结果在不同运行中是否可复现。

以下是50次采样试验的结果。数字表示每个子测试观察到的独特输出数量（数值越低，确定性越高）。

离线推理性能评估:

使用三种常见的RL推理工作负载（256个输入/输出长度不同的请求）来衡量非确定性模式和确定性模式的端到端延迟。

测试结果显示，确定性推理具有良好的可用性，大部分性能下降控制在25%到45%之间，其中FlashInfer和FlashAttention 3后端的平均性能下降为34.35%。大部分开销来自未经优化的基于Triton编写的批次不变算子，这表明性能仍有很大的提升空间。

团队确实观察到确定性推理比正常模式要慢。因此建议主要将其用于调试和复现性。SGLang团队未来的工作之一将专注于加速确定性推理，目标是将性能差距缩小到20%以内，或理想情况下达到与正常模式持平。

使用方法

SGLang环境设置：

# 使用最新的主分支
git clone https://github.com/sgl-project/sglang.git
cd sglang
# 安装 Python 依赖
pip install --upgrade pip
pip install -e "python[all]"

启动SGLang Server：

SGLang支持多种模型的确定性推理。例如，对于Qwen3-8B，您只需在启动服务器时添加—enable-deterministic-inference参数：

python3 -m sglang.launch_server
    --model-path Qwen/Qwen3-8B
    --attention-backend <flashinfer|fa3|triton>
    --enable-deterministic-inference

未来工作

团队未来的努力将集中在以下领域，以增强确定性推理的性能和可用性：

• 更快的批次不变算子：批次不变算子是性能瓶颈，因此将努力优化其配置并可能重写它们以提升性能。这对于提高 RL 推理的速度也至关重要。
• 支持 MoE 模型：目前只支持像 QWen3-8B 或 LLaMa-3.1-8B 这样的Dense模型的确定性推理。未来计划将支持扩展到像Qwen3-30B-A3B或DeepSeek-V3这样的混合专家模型。
• 增强基数缓存功能：将改进基数树，使其与更广泛的注意力算子兼容，超越目前对FlashAttention 3后端的限制。
• 张量并行（Tensor Parallelism）：观察到TP1和TP2是确定性的，可能因为不涉及浮点加法的结合率。将尝试扩展到更大的 TP size， 可能包括修改归约算子以实现确定性。
• FlexAttention集成：除了目前支持的注意力后端，计划未来将确定性推理的支持扩展到FlexAttention。

SGLang的确定性推理和slime的可复现训练功能目前正在积极完善中，团队诚挚欢迎广大用户和开发者积极试用，并提供宝贵的反馈意见。您的使用体验和建议将推动这一重要功能的进一步优化，推动确定性推理技术的发展。

参考链接：

[1]https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/

[2]https://lmsys.org/blog/2025-09-22-sglang-deterministic/

[3]https://github.com/sgl-project/sglang/issues/10278

[4]https://thudm.github.io/slime/_examples_synced/reproducibility/README.html#reproducibility

[5]https://github.com/THUDM/slime/pull/370

Read More