THUDM/slime项目使用指南：从参数配置到自定义训练流程

项目概述

THUDM/slime是一个基于强化学习的大模型训练框架，它整合了Megatron-LM的训练能力和SGLang的推理能力，实现了高效的训推一体化流程。本文将详细介绍如何使用slime进行大模型训练，包括参数配置、资源分配、模型加载以及自定义训练流程等核心内容。

集群资源配置

基础资源配置

在slime中，GPU资源分配主要涉及以下关键参数：

1. 训练资源配置

   • --actor-num-nodes：指定用于RL训练的actor节点数量
   • --actor-num-gpus-per-node：每个actor节点使用的GPU数量

2. 推理资源配置

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

训推一体模式

通过设置--colocate参数可以启用训推一体模式，此时：

• 训练和推理将共享相同的GPU资源
• --rollout-num-gpus参数将被忽略
• 需要特别注意显存管理，可通过调整--sglang-mem-fraction-static控制显存占用

模型加载与配置

Megatron模型加载

slime深度集成了Megatron-LM训练框架，加载Megatron模型需要三个关键配置：

1. 模型架构参数

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

--num-layers 36
--hidden-size 2560
--ffn-hidden-size 9728
--num-attention-heads 32
--group-query-attention
--num-query-groups 8
--normalization "RMSNorm"
--norm-epsilon 1e-6

2. 并行策略配置

Megatron提供多种并行优化策略：

• 基础并行：

  • --tensor-model-parallel-size：张量并行大小
  • --sequence-parallel：建议与张量并行同时启用
  • --pipeline-model-parallel-size：流水线并行大小

• 高级并行：

  • --context-parallel-size：序列并行（类似ring attention）
  • --expert-model-parallel-size：MoE专家并行
  • --expert-tensor-parallel-size：MoE专家张量并行

3. 检查点管理

slime支持两种Megatron检查点格式：

• Torch格式：传统格式，需保持并行策略一致
• Torch_dist格式（推荐）：支持自动并行切分

关键参数：

• --ref-load：参考模型检查点路径
• --load：训练模型检查点路径
• --save：模型保存路径

SGLang推理配置

SGLang的配置相对简单：

• --hf-checkpoint：指定HuggingFace格式的模型检查点
• --sglang-context-length：可覆盖模型原始上下文长度限制
• 其他SGLang参数需添加--sglang-前缀，如--sglang-mem-fraction-static

数据格式与处理

输入数据格式

slime目前仅支持.jsonl格式，每行包含一个JSON对象，示例：

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

关键配置参数：

• --input-key：指定输入字段（如"prompt"）
• --label-key：指定标签字段（如"label"）
• --apply-chat-template：是否应用聊天模板

自定义训练流程

1. 自定义Rollout函数

slime支持不同级别的数据生成定制：

基础定制：替换默认的generate_rollout函数

• 通过--rollout-function-path指定自定义实现
• 函数需返回包含tokens、response_length、reward等字段的Sample对象

高级定制：仅替换生成逻辑

• 通过--custom-generate-function-path指定
• 实现异步生成接口，与SGLang router交互

2. 自定义Reward模型

通过--custom-rm-path参数可以指定自定义的reward模型实现，用于替代默认的奖励计算逻辑。

技术深度解析

SGLang集成原理

slime通过HttpServerEngineAdapter桥接层实现了与SGLang的深度集成：

1. 使用sglang-router进行负载均衡
2. 所有SGLang server通过/add_worker注册到router
3. 推理请求通过HTTP发送到router进行分发

Megatron集成原理

slime通过复用Megatron的核心组件实现兼容：

1. 直接使用Megatron的parse_args解析参数
2. 复用checkpoint保存/加载逻辑
3. 支持通过PYTHONPATH加载自定义Megatron版本

关键扩展点：

• --custom-megatron-init-path：自定义初始化逻辑
• --custom-megatron-before-train-step-hook：训练步前回调

最佳实践建议

1. 资源分配：

   • 训推分离时确保推理资源充足
   • 训推一体时合理设置--sglang-mem-fraction-static

2. 模型配置：

   • 优先使用torch_dist格式检查点
   • 确保并行策略与硬件匹配

3. 性能优化：

   • 合理配置重计算参数
   • 根据任务特点调整数据生成策略

4. 定制开发：

   • 从小范围hook开始扩展功能
   • 充分利用现有Sample数据结构

通过本文介绍，您应该已经掌握了使用THUDM/slime进行大模型训练的核心方法。无论是基础训练还是深度定制，slime都提供了灵活的配置选项和扩展接口，能够满足各种复杂场景下的训练需求。