大数跨境
0
0

AMD x slime: ROCm Day-0支持SGLang 原生强化学习框架,助力推理性能跃迁

AMD x slime: ROCm Day-0支持SGLang 原生强化学习框架,助力推理性能跃迁 AMD开发者中心
2025-10-15
1




原文作者: Yusheng Su, Yuzhen Zhou, Jin Pan, Gowtham Ramesh, Xiaodong Yu, Jialian Wu, Ze Wang, Ximeng Sun, Jiang Liu, Hao Chen, Zicheng Liu, Emad Barsoum.


AMD 很高兴为 SGLang 原生的强化学习(RL)框架 slime 提供 Day-0 支持在本文中,我们将提供有关我们支持与优化的更多细节,以及 slime 在大规模 RL 训练中的优势。


首先,我们将介绍 slime 背后的工程工作包括对代码库的修改、面向 ROCm 软件的内核级内存管理,以及对第三方依赖(Megatron-LM、SGLang 和 torch_memory_saver)的修改——并提供可在 AMD GPU 上实现高效执行的 Docker 镜像。在架构层面,slime 支持两种训练模式:同步与异步。围绕这两种模式,我们还给出系统级优化及其对应的应用场景。具体而言,在同步设置下,我们的 rollout 优化相较于未使用该优化的方案,在 AMD GPU 上带来 40% 的吞吐量提升在异步设置下,我们开发了一个多轮 RL  agent框架,用于训练内核生成模型。读者也可以在MLsys – SGLang 官方博客中阅读更多关于该支持的信息:https://lmsys.org/blog/2025-07-09-slime


核心要点

  • slime 框架及其优势

  • 代码库级与内核级支持及相应发布的Docker 镜像

  • 系统级优化:APRIL -- 主动式部分rollout(Active Partial Rollouts)

  • 多轮内核代理(kernel-agent)RL 训练框架


slime 框架及其优势

要开发智能化的大规模基础模型,RL 训练与预训练同等关键。自 2023 年底 ChatGPT 发布以来,RLHF 在增强大规模预训练语言模型(LLM)方面的有效性日益凸显。近期,多款采用 RLHF 训练的 O1/R1 系列模型的推出再次证明了其能力,尤其体现在提升推理能力方面。此外,训练阶段 RL 的投入也在不断增加。因此,开源社区中涌现出越来越多的 RL 训练框架,包括 OpenRLHF、verl、Areal 以及 slime。在这些框架中,slime 是目前唯一已被验证可用于大规模混合专家(MoE)模型 RL 训练的框架,模型规模可达 3550 亿参数。



1,Slime整体框架


slime 框架不仅支持监督式微调(SFT),还作为一个 RL 框架,通过可定制的 rollout 接口和多样的训练配置(无论是同机还是解耦、同步还是异步)来强调通用性。如图 1 所示,它通过 SGLang 执行推理、并以原生方式利用 Megatron-LM 进行训练,从而实现高性能,同时保持轻量且易维护的代码库,平滑衔接从 Megatron 预训练到 SGLang 部署的过程。此外,slime 通过重新思考 RL 数据采样以更好地满足用户需求,从而简化开发:它通过 sgl-router 集中管理所有 SGLang 服务器,暴露单一 HTTP 端点,使环境能够通过兼容 OpenAI 的 API 进行通信,同时确保训练与部署的一致性。同时,它借助基于 Ray 的资源管理简化实验搭建,使用户能够在同机与解耦的 GPU 配置之间无缝切换且全部在统一框架内完成。


体而言,slime 提供:

  • 可定制的 rollout 接口与多训练设置(同机/解耦/异步)。

  • 原生 SGLang 体完整 MoE 支持、配置透调试模式。

  • 量代码库,最少抽象,易展、易维护

  • 已在大模代智能体 RL 训练中得到验证,适用于企业场景。


码库级与内核支持

为确保 slime 在 AMD GPU 上高效运行,我们引入了多项关键增强。首先,我们更新了 slime 代码库以实现与 ROCm 的全面兼容,从而在 AMD 硬件上实现稳定高效的执行。此外,我们通过改进虚拟内存管理并新增基于 HIP 的内存卸载与上传实现,解决了与 ROCm 相关的问题,从而提升 RL 训练效率。最后,我们提供 Dockerfile 及对应的 Docker 镜像(https://github.com/THUDM/slime/blob/main/docker/Dockerfile.rocm,以简化 slime 训练环境的搭建。以上可以参考slime仓库发布(https://github.com/THUDM/slime/pulls?q=is%3Apr+Hardware+is%3Aclosed需求列表。


快速

你可以尝试以下脚本在你的系统上运行 slime。

首先,启动 Docker 镜像。接着,运行示例脚本。


接着,运行示例脚本:


如前所述,slime 为 RL 训练提供两种模式:同步(GPU 同机模式)与异步(GPU 解耦模式)。在同步模式下,推理与训练引擎共享同一组 GPU,系统在 rollout 生成(推理)与梯度更新(训练)阶段之间交替进行。这与异步架构形成对比,即在异步架构中,为每个引擎分别配置独立的 GPU 集群。同步模式的主要优势在于降低硬件门槛,使训练能够在相对有限的算力资源上运行。然而,它需要更复杂的系统管理,因为不同的引擎必须在 GPU 上卸载与重新加载。相比之下,异步模式会消耗更多 GPU,但避免了管理 GPU 交替带来的开销。在接下来的两个章节中,我们将讨论各模式最适用的场景,并重点说明我们所解决的瓶颈。


统级优

 GPU 资源有限且 rollout 生成长度相对均匀(例如数学推理任务)时,我们采用同步模式。然而,rollout 生成是 RL 训练中的主要开销,通常占据总训练时间的 90% 以上。此外,训练任务中往往包含少量极长的 rollout 生成,如图 2(A) 所示。


这些“长尾”会导致 GPU 在等待最长序列完成时处于空闲,从而降低利用率并拖慢整个训练周期,如图 2(B) 所示。



2. (A) 不同任务中 rollout 长度的分布情况。(B) 标准 rollout 在强化学习训练过程中会出现大量“气泡”现象。(C) APRIL 机制在强化学习训练中有效缓解了这些“气泡”问题。


为缓解这一问题,我们提出Active Partial Rollout(APRIL)策略(https://github.com/RLsys-Foundation/APRIL),并已合并至上游 slime 仓库(https://github.com/THUDM/slime/pull/2),用于解决长尾瓶颈:

  • :启超出需求数量的 rollout 求(例如目批次 32, 64)。

  • 提前:一旦达到所需批次大小,立即止其余求。

  • 收集与复用:被中止、未完成的迹存入冲区,在下一次迭代中接上次继续生成。


机制原生集成于 SGLang 的路由器,可通过 --partial-rollout 与 --over-sampling-batch-size 简单启用,能够显著减少训练过程中的“气泡”,且在各类 RL 算法与 LLM 上表现良好。在下文中,我们给出一个对比实验,使用常见算法 GRPO,并结合 Qwen3-4B 模型与 DeepMath 推理数据集:

  • 吞吐量:在最大响应长 16k 的定下,端到端吞吐量从 9.3k 提升至 13k tokens/秒(提升 40%)。

  • 定性与准确性:与原始方法相比,APRIL 更定,且准确率略有提升,尽管其中包含一些 off-policy 的 rollout(即完成上一个时间步的未完成 rollout)。 on-policy  off-policy 混合的 rollout 生成方式可能提升数据多性,从而来更健的训练



 3 [左] APRIL(蓝色)相比原始基线(橙色)吞吐量提升了40%;[右] 与原始方法(橙色)相比,APRIL(蓝色)表现更稳定,且准确率略有提升。


内核代理(Multi-turn Kernel-agent)RL 训练框架

在诸如代码生成等代理型任务中,同步模式的 RL 框架会遇到瓶颈。由于该系统的同步性,每一步都必须包含推理(rollout)与训练(模型权重更新)。然而,代理型任务的 rollout 分布往往呈现出比前述更为严重的“长尾”问题。


为此,slime 还提供了完全异步、解耦的设计:在 rollout 引擎与训练引擎之间对 GPU 进行分区。rollout 引擎承载 LLM 服务器,持续流式生成轨迹;训练引擎消费这些轨迹,计算损失并更新 actor。随后,新权重会广播至 rollout 服务器;进行中的请求会被安全抢占,并在更新后的模型上继续生成。


其优势如下:

  • 无批次栅栏迹随到随理,避免首阻塞。

  • 更高利用率:rollout 与训练并行运行,最大化 GPU 吞吐量并墙钟时间

  • 内存定性:不再繁加/卸,降低碎片化在大景尤

  • 清晰的横向展:件相互独立,便于按需 rollout 或训练能力。


 4 展示了该资源分配与工作流,说明 rollout 引擎与训练引擎如何在时间轴上独立运行。


4. SLIME 异步框架概览


基于该架构,我们设计了一个多轮代理训练框架TritonForge(https://github.com/RLsys-Foundation/TritonForge一个kernel 生成沙盒项目,使 RL 引擎 slime 能与之交互。更具体地说,该沙盒作为一个基于 kernel-bench 的环境,用于 Torch 到 Triton 的任务。在 RL 训练过程中,输入 Torch 代码,由推理引擎处理以生成 Triton 代码。生成的 Triton 代码随后在沙盒中执行,以获得由编译成功、正确性与加速比定义的奖励。随后,这些奖励被用于优化策略(policy)。此外,该框架同时支持单轮与多轮两种 RL 训练模式。在我们的实验中,多轮模式在代码生成性能上始终表现更佳。我们在 Qwen3-8B 上进行了内核代理单轮训练,对应的训练曲线见图 5

5. 单轮 Qwen3-8B kernel agent强化学习训练曲线,显示reward随时间稳步提升


总结

随着 RL 训练成为推动基础模型超越预训练的关键,slime 凭借其原生集成 SGLang、支持高达 3550 亿参数的大规模 MoE 模型,成为唯一经验证的框架而脱颖而出本文展示了 AMD 在 GPU 上对 slime 的 Day-0 ROCm 支持,涵盖我们的内核级优化、Docker 部署,以及在同步模式下实现 40% 吞吐增益的创新 APRIL 策略。我们同时给出了同步与异步训练架构,并通过面向 Torch-to-Triton 代码生成的多轮内核代理框架展示其实际应用。凭借轻量代码库、可定制的 rollout 接口以及无缝的基于 Ray 的资源管理,该框架使 AMD Instinct 用户能够充分发挥 slime 在可扩展 RL 训练中的全部潜力,最终加速高级推理与代码智能体的研发。


引用
如果您觉得这篇博客对您的工作有帮助,请随意引用。我们也鼓励您查看我们的论文APRIL: Active Partial Rollouts in Reinforcement Learning to tame long-tail generation,并在相关工作中引用它。


项目贡献者

核心贡献者:Yusheng Su, Yuzhen Zhou, Jin Pan, Gowtham Ramesh, Zicheng Liu

更多贡献者:Xiaodong Yu, Jialian Wu, Ze Wang, Ximeng Sun, Jiang Liu, Hao Chen, Emad Barsoum


鸣谢

感谢 AMD AIG Bingqing Guo 和 Guruprasad MP 赞助AMD Developer Cloud进行实验。感谢微软亚洲研究院的 Yang Wang 协助我们完成对 ROCm-6.3.4 的支持。感谢 slime 开源社区的 Zilin Zhu 在 SLIME 的高层理解与实现方面所做的贡献。感谢 MLSys – SGLang 开源社区的 Jiajun Li、Xiang Long 和 Chenyang Zhao 在系统级优化和内核代理框架方面给予的帮助。


AMD × slime:让大规模 RL 训练更高效、更智能、更可靠。即刻体验,成就更强 AI!


了解更多ROCmAI开请访问https://www.amd.com/en/developer/resources/rocm-hub.html






【声明】内容源于网络
0
0
AMD开发者中心
AMD开发者中心为开发者提供工具、技术和资源,助力构建AI解决方案。ROCm、Ryzen AI软件和ZenDNN,帮助您实现模型加速与部署。开发者可通过文档、SDK及教程快速上手。立即关注AMD开发者中心,开启智能未来!
内容 65
粉丝 0
AMD开发者中心 AMD开发者中心为开发者提供工具、技术和资源,助力构建AI解决方案。ROCm、Ryzen AI软件和ZenDNN,帮助您实现模型加速与部署。开发者可通过文档、SDK及教程快速上手。立即关注AMD开发者中心,开启智能未来!
总阅读127
粉丝0
内容65