THUDM/slime项目使用指南：从参数配置到高级功能详解

项目概述

THUDM/slime是一个基于强化学习的大语言模型训练框架，它巧妙地将Megatron-LM的训练能力与SGLang的高效推理能力相结合，为研究人员和开发者提供了一个强大的工具链。本文将深入解析slime的使用方法，从基础参数配置到高级功能定制，帮助用户充分发挥这一框架的潜力。

核心参数配置

1. 集群资源分配

slime采用训练与推理分离的架构设计，通过以下参数实现资源精细分配：

• 训练资源：

  • --actor-num-nodes：指定训练任务所需的节点数量
  • --actor-num-gpus-per-node：每个节点上分配的GPU数量

• 推理资源：

  • --rollout-num-gpus：推理部分所需的总GPU数量
  • --rollout-num-gpus-per-engine：每个推理引擎使用的GPU数量（类似SGLang的tp_size）

特殊场景：当需要训练与推理共享资源时，可启用--colocate参数，此时系统将忽略--rollout-num-gpus，使训练和推理使用相同数量的GPU资源。

2. Megatron模型加载

Megatron作为slime的训练核心，其配置需要特别注意：

模型参数配置

以qwen3 4B模型为例，典型配置包括：

--num-layers 36
--hidden-size 2560
--num-attention-heads 32
--group-query-attention
--normalization "RMSNorm"

注意：

• 项目已提供常见模型的预设配置（位于scripts/models目录）
• 无需配置--seq-length，slime采用数据打包(varlen/thd)训练方式

并行策略优化

Megatron提供多种并行策略，可根据硬件条件灵活组合：

• 基础并行：TP(--tensor-model-parallel-size)、PP(--pipeline-model-parallel-size)
• 高级优化：SP(--sequence-parallel)、CP(--context-parallel-size)
• MoE支持：EP(--expert-model-parallel-size)、ETP(--expert-tensor-parallel-size)

检查点管理

支持两种主流格式：

• torch格式：需严格匹配训练并行策略
• torch_dist格式（推荐）：支持自动并行分片

关键参数：

• --ref-load：参考模型检查点
• --load：训练模型检查点（未设置时从ref-load初始化）
• --save：模型保存路径

3. SGLang推理配置

SGLang作为推理引擎，配置相对简单：

• --hf-checkpoint：Hugging Face格式的模型检查点
• --sglang-context-length：覆盖模型原始上下文长度限制
• --sglang-mem-fraction-static：显存分配比例（共置训练时需调整）

数据格式规范

slime目前仅支持.jsonl格式，每条数据示例：

{
  "prompt": [{"content": "数学问题...", "role": "user"}],
  "label": "34"
}

对应配置参数：

• --input-key prompt：指定输入字段
• --label-key label：指定标签字段
• --apply-chat-template：应用对话模板

高级功能定制

1. 自定义Rollout函数

slime支持不同粒度的数据生成定制：

完全自定义

通过--rollout-function-path指定自定义函数，需满足接口：

def generate_rollout(args, rollout_id, data_buffer, evaluation=False) -> list[Sample]

部分定制

仅替换生成逻辑时使用--custom-generate-function-path，基础实现：

async def generate(args, sample: Sample, sampling_params) -> Sample

2. SGLang深度集成

slime通过HttpServerEngineAdapter桥接SGLang：

• 参数传递：所有SGLang参数需添加--sglang-前缀

  • 例如：--mem-fraction-static → --sglang-mem-fraction-static

• 路由管理：

  • 默认启动内置router
  • 可通过--sglang-router-ip/port指定外部router

3. Megatron扩展支持

针对定制化Megatron实现，提供多种钩子：

• --custom-megatron-init-path：初始化扩展
• --custom-megatron-before-log-prob-hook-path：概率计算前钩子
• --custom-megatron-before-train-step-hook-path：训练步骤前钩子

最佳实践建议

1. 资源分配：

   • 大规模训练优先采用分离式部署
   • 小规模实验可使用--colocate节省资源

2. 模型配置：

   • 推荐使用torch_dist格式检查点
   • 复杂模型优先启用SP+CP组合

3. 性能调优：

   • 共置训练时适当降低--sglang-mem-fraction-static
   • 长上下文推理显式设置--sglang-context-length

4. 开发调试：

   • 从小规模--rollout-num-gpus-per-engine开始验证
   • 利用预设模型配置快速启动

通过本文的详细解析，开发者可以全面掌握slime框架的使用方法，从基础训练到高级定制，充分发挥这一强化学习训练框架的强大能力。实际应用中，建议结合具体任务需求，灵活组合文中介绍的各项功能和配置策略。