Slime 框架深度解析：面向大规模RL的训推一体化实践

最近 RL 领域非常热闹，算法上涌现了各种神奇的操作：从少量数据、单条数据，乃至模型自产数据进行强化学习，甚至连标签错误的 RL 数据也能学到知识。这些进展都很有趣，但个人感觉尚未触及特别本质的突破。因此，我将更多精力转向了基础设施（Infra）层面。本文旨在分析 Zilin 的 Slime 框架。无论算法如何演进，RL 终究是一个算法与系统协同的问题，其规模化（scale）的价值是明确的。

整体设计

Slime 采用 Ray 作为单控制器（Single Controller）进行调度， Megatron-LM 作为训练后端， SGLang 作为推理后端，构建了一个简洁且支持大规模 RL 的框架。它兼容 资源训推一体化 和 分离 两种部署模式。正如 Zilin 所述，其设计目标是减少训练与推理之间数据传递的开销，并使所有环节尽可能对齐真实生产环境中的组件，从而实现灵活的大规模 RL 训练。从实践来看，这套系统成功地实现了其核心设计目标。

本文将沿着 RL Infra 的脉络，学习和巩固相关知识，为 RL 算法迭代和 RL 框架的协同设计提供参考。

从数据流的角度看，以 Buffer 为视角最易于理解。 Buffer 是一个 Ray Actor，它不仅维护了一个 rollout 数据集，还能加载自定义的 generate_rollout 方法，以实现灵活的 rollout 策略。通过 Ray 的远程调用，训练和推理模块可以与 Buffer 进行数据交互。 Rollout (推理/样本生成) : 使用 SGLang 的 Router 进行数据并行（DP），router 依据缓存感知（cache-aware）原则进行负载均衡，再将请求分配给相应的 SGLang Engine。 rollout_generator 直接持有 Buffer 的句柄，双向管控数据的生产和取回。 Train (训练) : 尽量复用 Megatron 的 trainer。5D 并行（TP/PP/DP/CP/EP）在兼容 GPTModel 的基础上，更多地交由 Megatron 的原生 trainer 处理。

整体控制流程

一张图可以总结整体的控制流程，细节将在后续章节中分析。

Placement Group 组创建

Slime 使用 Ray 作为其资源细粒度调度的框架，以支持分离（separated）或统一（co-located）的资源放置策略。这是如何实现的呢？

首先，系统根据当前模式（ train 或 debug ）设定一个具名的 Placement Group (PG)。 - Co-located 模式 : rollout 和 actor (训练) 指向同一个 PG，实现资源共享。 - Debug 模式 : rollout Ray 的 Placement Group Scheduling 机制是实现这一点的关键。它通过分配逻辑上的 虚拟资源 来管理 GPU。例如，可以分配给 train group 中的每个 actor 0.8 个 GPU ，然后 rollout group 中的每个 engine 再要求 0.2 个 GPU 。Ray 本身不限制 GPU 的实际使用，而是像一个“账本师傅”，按需进行逻辑分配。分配逻辑会先满足 0.8 的 GPU 需求，若第一张卡剩余 0.2，不足以再分配一个 0.8，则会转向下一张卡，直到满足所有训练 actor 的需求，每张卡上都留下 0.2 的“空闲”逻辑资源。当 rollout 引擎请求 0.2 个 GPU 时，它会恰好填满这些剩余的逻辑资源，从而实现 co-location。# slime/ray/placement_group.pydefallocate_train_group(num_nodes,num_gpus_per_node,pg,debug_rollout_only):returnRayTrainGroup(num_nodes=num_nodes,num_gpus_per_node=num_gpus_per_node,pg=pg,num_gpus_per_actor=0.8,# 训练 Actor 请求 0.8 GPUdebug_rollout_only=debug_rollout_only,)# slime/ray/rollout.pyforiinrange(num_engines):num_gpus=0.2# 推理引擎请求 0.2 GPUnum_CPUs=num_gpusscheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg,placement_group_capture_child_tasks=True,placement_group_bundle_index=reordered_bundle_indices[i*num_gpu_per_engine],)rollout_engines.append(RolloutRayActor.options(num_cpus=num_cpus,num_gpus=num_gpus,scheduling_strategy=scheduling_strategy,).remote(args,rank=i,data_buffer=data_buffer,))

对于非 co-located 的情况，由于初始创建的 PG 就是分离的，actor 和 rollout 会自然地被调度到各自指定的 PG 上。

训练初始化

Ray Train Group / Actor 初始化

理解这个逻辑可以分为三层：

顶层 ( RayTrainGroup ) : 负责 TrainRayActor 实例的创建、初始化和资源分配。它接收顶层指令（如 train.py ），并将其广播到其管理的所有 TrainRayActor 实例。 中层 ( TrainRayActor ) : 作为分布式训练的基本执行单元，每个 TrainRayActor 负责一部分计算任务，并与 RolloutRayActor 建立连接以进行权重同步。 底层 (Megatron-LM) : 通过 megatron_utils 封装的核心训练逻辑。 RayTrainGroup def 定义。这是因为在 Ray 中，调用 .remote 方法会立即同步返回一个句柄对象（类似 Future），而真正的异步任务执行发生在后台。只有当显式调用 ray.get 时，才会阻塞并获取结果。 权重管理 : 初始化时使用了与 vLLM 相同的 CuMemallocator 。这个分配器比 PyTorch 的默认管理更底层，支持通过标签（tag）定向地卸载（offload）或加载（onload）参数，这对于实现 sleep_mode 至关重要。
allocator=CuMemAllocator.get_instanceifself.args.offloadelseNonewithallocator.use_memory_pool(tag="model")ifallocatorelsenullcontext:(self.model,self.optimizer,self.opt_param_scheduler,loaded_rollout_id)=(megatron_utils.initialize_model_and_optimizer(args,with_optimizer=True))start_rollout_id=loaded_rollout_id+1 模型 On/Offload : 接口与 vLLM 的 sleep 和 wake_up 对齐，可以灵活地将模型参数在 GPU 和 CPU 之间移动。 defsleep(self,tags):assertself.args.offloadassert"model"intagswithtimer("sleep"):ifisinstance(tags,str):tags=(tags,)clear_memoryprint_memory(f"before offload model")self.update_cpu_params_dictallocator=CuMemAllocator.get_instanceallocator.sleep(offload_tags=tags)clear_memoryprint_memory(f"after offload model")defwake_up(self,tags):assertself.args.offloadwithtimer("wake_up"):clear_memoryprint_memory("before wake_up model")ifisinstance(tags,str):tags=(tags,)allocator=CuMemAllocator.get_instanceallocator.wake_up(tags)clear_memoryprint_memory("after wake_up model")

Rollout Group / SGLang Engine 初始化

同样可以分为三层来理解：

顶层 ( RolloutGroup ) : 负责宏观的生命周期管理，包括启动 SGLang 负载均衡器、创建数据缓冲池以及实例化一组工作 Actor。 中层 ( RolloutRayActor ) : 作为面向 Ray 分布式环境的接口，封装了核心的 SGLang 推理引擎。 底层 ( SglangEngine ) : 包含了 SGLang 的具体实现，负责实际的文本生成任务。 启动 Router ( start_router ) : 通过 multiprocessing.Process 启动一个独立的 sglang_router 进程。这个路由器作为负载均衡器，接收推理请求并分发给后端的 SGLang 服务，实现了请求提交与实际推理的解耦。 创建 Buffer Actor : 创建一个中心化的 Buffer Actor，负责管理整个数据流的生命周期：生成 prompts、接收生成结果、存储数据供训练 Actor 获取。 创建 Rollout Engines ( create_rollout_engines ) : 这是初始化的核心。它会创建一个 RolloutRayActor 集群，并为每个引擎精心分配网络端口（HTTP、NCCL 等），这对于跨节点分布式推理至关重要。 创建 lock Actor : 创建一个简单的分布式锁 Actor，用于在权重更新时防止数据竞争，确保 rollout 引擎不会使用不一致的模型参数。

模型与参数同步

模型同步初始化

TrainRayActor 的方法最终调用 megatron_utils 来完成分布式环境和模型的初始化。 分布式环境初始化 : megatron_utils.init(args) 会初始化 torch.distributed ，并设置 Megatron-LM 的模型并行/流水线并行组（MPU）。 分词器串行加载 : 为了避免多个 rank 同时读写 Hugging Face 缓存目录导致文件损坏，配置和分词器采用串行加载（rank by rank）的方式。 模型和优化器初始化 : 使用 CuMemAllocator 上下文管理器，为后续的显存卸载（sleep）功能打下基础。 模型初始化遵循 Megatron-PAI 的实践，直接提供 GPTModel 的 model_provider`，减少盲目魔改，增强生态兼容性。 Reference Model 初始化 : 加载与主模型相同的初始权重，但在训练中保持不变，作为计算 KL 散度的基准。同样支持卸载到 CPU 以节省显存。

参数同步连接

权重同步相关的代码（转换、通信协议、连接建立）在大规模框架中通常占很大比重。Slime 的处理方式非常值得学习。

顶层调用 ray.get(actor_model.async_init_weight_update_connections(rollout_generator)) 显式地将训练和推理模块连接起来。# 文件: slime/ray/ppo_actor.pyclassRayTrainGroup:defasync_init_weight_update_connections(self,rollout):self.rollout=rollout# 1. 设置数据缓冲区ray.get([actor.set_data_buffer.remote(rollout.data_buffer)foractorinself._actor_handlers])# 2. 获取并同步并行化配置actor_parallel_configs=ray.get([actor.get_parallel_config.remoteforactorinself._actor_handlers])# ... (聚合配置信息) ...ray.get(rollout.async_set_parallel_config(parallel_config))# 3. 触发真正的连接建立return[actor.connect_rollout_engines.remote(rollout.rollout_engines,rollout.rollout_engine_lock,)foractorinself._actor_handlers]在执行层，根据 actor 和 rollout 是否 co-located，采用两种不同策略建立通信组： 策略1: Co-located (CUDA IPC) : 当训练和推理 Actor 在同一批 GPU 上时，为每个 SGLang 引擎创建一个只包含同 GPU 训练 Actor 的 torch.distributed 通信组。后端使用 nccl ，它会自动选择最高效的 CUDA IPC 进行 GPU 显存间的直接数据拷贝，避免了网络开销。 策略2: Distributed (TCP/NCCL) : 当 Actor 分布在不同机器时，由 TrainRayActor 的 rank 0 作为主节点，广播其 IP 和端口。所有 TrainRayActor 和 RolloutRayActor 加入这个新的通信组，最终建立一个跨节点的、基于 TCP 的 torch.distributed 通信组。

由于所有 Megatron 权重最终都转换为 HF 格式供 rollout 侧消费，这里不涉及复杂的权重 resharding 操作。

训练主流程

所有初始化和配置工作都是为了训练能够高效稳定地运行。主流程核心包括： 样本生成 -> 模型训练 -> 权重更新 。

训练流程 — 生成样本

Buffer Actor。# 文件: slime/ray/rollout.pyclassRolloutGroup:defasync_generate(self,rollout_id,evaluation=False):returnself.data_buffer.generate.remote(rollout_id,evaluation=evaluation) Buffer Actor 负责编排整个流程。它会动态加载用户指定的函数，这使得数据生成逻辑完全可插拔，为实现复杂的 Agentic 逻辑提供了极大的灵活性。# 文件: slime/rollout/sglang_example.pyasyncdefgenerate_rollout_async(args,rollout_id,data_buffer):# 1. 从 Buffer 获取 Promptssamples=awaitdata_buffer.get_samples(sampling_batch_size*args.n_samples_per_prompt)# 2. 异步提交生成和 RM 评估任务forsampleinsamples:state.pendings.add(asyncio.create_task(generate_and_rm(...)))# 3. 异步收集结果whilelen(data)

这种以模型为中心、高度灵活的设计，非常适合 Agent 时代的需求，避免了厚重框架对模型自主探索的限制。一个巧妙之处在于，我们可以只启动 rollout 部分来独立调试生成策略，极大地提升了开发和迁移效率。

Rollout 过程产生的经验（ Sample dataclass）被转换为标准的 Python 字典格式（ train_data ），存储在 Buffer 中，通过 Ray 对象存储高效地传输给训练 Actor。

训练流程 — 模型训练

获取数据 : TrainRayActor 的 rank 0 从 Buffer 获取数据，然后通过 dist.broadcast_object_list 广播给所有其他 rank。 数据分片 : 每个 TrainRayActor 根据其数据并行（DP）排名对接收到的数据进行切分。 存入本地存储 : 切分后的数据存入一个名为 LOCAL_STORAGE 的进程内全局字典中，后续的 log probs, values, rewards 等中间结果也存储于此。 计算 Reference Log Probs : 使用参考模型（ self.ref ）进行一次前向传播，计算并存储 ref_log_probs 到 LOCAL_STORAGE 。 计算 Actor Log Probs : 使用主策略模型（ self.model ）进行一次前向传播，计算并存储 log_probs 。 计算优势函数 (Advantage) : 从 LOCAL_STORAGE 中取出 log probs 和 rewards，计算 KL 散度，并根据 GRPO（或 PPO）公式计算 advantages 和 returns ，再存回 LOCAL_STORAGE 。 执行训练步骤 : 使用计算出的 advantages ，调用 megatron_utils.train 来执行一个或多个优化步骤（前向传播计算 loss -> 反向传播 -> 梯度裁剪 -> 参数更新）。

核心 Loss 计算细节 :

Policy Loss : 实现了标准的 PPO clip loss，并支持 DAPO 的 clip_higher 和 Dual-Clip 变体。 Entropy Loss : 由于张量并行（TP）会将词表切分，这里借助 Megatron Core 的 fused_vocab_parallel_cross_entropy 来高效、正确地计算交叉熵。 Context Parallelism (CP) Loss : 这是实现上最复杂的部分。当一个序列被切分到多个 GPU 上时（CP），计算序列的平均损失需要全局信息（如序列总长度），代码中通过精细的 loss_mask chunk 计算来处理这一问题，确保 loss 被正确归一化。

训练流程 — 权重同步

训练完成后，需要将更新后的权重同步到 SGLang 推理引擎。

路径一：分布式 (TCP/NCCL)

暂停推理 准备参数 : 训练 Actor 的 rank 0 通过 Ray 调用，暂停所有 SGLang 引擎的生成任务。 重组权重 转换为 HF 格式 : 使用 all_gather 收集参数分片，得到完整的参数。 移除词表 padding。 在指定 rank 上将 Megatron 格式的权重转换为 SGLang 能识别的 Hugging Face 格式。 广播权重 : 在之前建立的 _model_update_groups 通信组内，由 rank 0 使用 dist.broadcast 将转换好的权重数据块广播给所有 SGLang 引擎。

路径二：Co-located (CUDA IPC)

加载并对齐参数 : 将 CPU 上的权重加载到其源 rank 的 GPU 上，并通过内部广播（如在 PP group 内）确保所有训练 Actor 都有完整的参数。 重组权重 转换 : 流程与分布式类似，但操作对象是已在 GPU 上的权重。 发送 IPC 句柄 : 将转换好的权重序列化，获取一个指向 GPU 显存的轻量级 IPC 句柄。 在 IPC 通信组内收集所有句柄。 主进程通过 Ray Actor 调用，将 IPC 句柄列表发送给 SGLang 引擎。SGLang 接收后直接从对应的显存地址读取数据，完成更新，效率极高。 总结

Slime 不仅仅是一个 RLHF 的实现，它更是一个将大规模分布式训练与高性能推理服务深度融合、可扩展的 “决策大模型”在线学习框架 。它的核心价值在于，成功地解决了在真实世界中部署和持续优化一个大型 AI Agent 所面临的最棘手的工程挑战。

核心优势总结

工业级的“训推一体”架构 : 通过 Ray Actor 模型，将专为大规模设计的离线训练引擎 (Megatron-LM) 和专为低延迟、高吞吐量设计的在线推理引擎 (SGLang) 无缝地“粘合”在一起。 高效的“知识”流动闭环 : 框架实现了 RL 的核心飞轮： 生成 -> 训练 -> 同步 。 数据流 : 通过异步的 Buffer Actor 和高效的 NCCL Broadcast，实现了从“探索”到“学习”的高效数据回传。 权重流 : 通过自适应的权重同步机制（分布式 NCCL Broadcast 和同地部署 CUDA IPC），实现了从“学习”到“行动”的低延迟知识更新。 高度的灵活性与可扩展性 : 可插拔的探索与奖励 : 通过配置文件指定 rollout_function_path 和 custom_rm_path ，用户可以轻松自定义 Agent 的探索行为和价值判断标准，无需修改框架核心。 模型无关性 : 通过 slime_plugins 机制，可以轻松适配和扩展支持新的模型结构。

未来挑战与展望

尽管 Slime 已经非常强大，但在 Agentic 场景下，当前的主流 RL 框架仍面临一些共同的挑战：

长序列与延迟奖励 : 如何支持更灵活的、涉及超长序列的决策问题，其中奖励信号（RM feedback）可能在多轮 rollout 后才能获得。 彻底的异步与解耦 : 将生成、训练、评估彻底分解为独立的异步集群，使用分布式数据库或对象存储构建持久化的大规模经验池（Experience Pool）。 复杂的 Agent 模式 : 如何更好地支持多智能体（Multi-agent）协作、多模态（Multi-modal）输入等，以扩展模型的能力边界。