mirror of
https://github.com/radixark/miles.git
synced 2026-10-02 07:14:53 +08:00
5.7 KiB
5.7 KiB
slime
slime 是为 RL scaling 设计的 LLM post‑training 框架,提供两大核心能力:
- 高性能训练:通过连接 Megatron 与 SGLang,支持各种模式的高效训练;
- 灵活的数据生成:通过自定义数据生成接口以及 server based engine,实现任意的数据训练数据生成流程。
目录
架构总览
模块说明:
- training (Megatron):负责主训练流程,从 Data Buffer 读取数据,训练完后将参数同步至 rollout 模块;
- rollout (SGLang + router):生成新数据(含 reward/verifier),存储至 Data Buffer;
- data buffer:桥梁模块,管理 prompt 初始化、自定义数据与 rollout 生成方法。
快速开始
环境准备
基于镜像 zhuzilin/slime:latest(已预装 SGLang 0.4.7 和 Megatron):
docker run --rm --gpus all --ipc=host --shm-size=16g \
--ulimit memlock=-1 --ulimit stack=67108864 \
-it zhuzilin/slime:latest /bin/bash
git clone https://github.com/THUDM/slime.git
cd slime
pip install -e .
- 我们正在与 SGLang 社区合作,将 sglang.patch 中的更新合入主分支;
- 对于不方便使用 docker 的场景,请参考 从零搭建环境。
示例:GLM-4-9B 与 Qwen3-4B
我们提供了 GLM-4-9B 和 Qwen3-4B 的使用样例,请查看:
我们也提供了 MoE 模型的样例,请查看:
Checkpoint 格式转换
由于 slime 使用 megatron,而 megatron 不支持加载 huggingface checkpoint,我们需要将模型转换至 megatron 可以支持的 torch_dist 格式。
HF → Megatron torch_dist ckpt
使用 mbridge 转换:
cd slime/
PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
--hf-checkpoint /root/GLM-Z1-9B-0414 \
--save /root/GLM-Z1-9B-0414_torch_dist
⚠️ 如果出现找不到 slime 的问题,请在 slime 目录下 pip install -e .。
Megatron torch_dist → HF ckpt
将训练过程中的存储的 torch_dist ckpt 转为 hf ckpt:
cd slime/
PYTHONPATH=/root/Megatron-LM python tools/convert_torch_dist_to_hf.py \
--input-dir /path/to/torch_dist_ckpt/iter_xxx/ \
--output-dir /root/GLM-Z1-9B-0414-iter_xxx \
--origin-hf-dir /root/GLM-Z1-9B-0414
⚠️ 由于 mbridge 转换的 torch_dist ckpt 目前不保存 args,不能基于上一步的 torch_dist ckpt 反转回 HF。
任意 Megatron ckpt → HF
适用于自定义保存格式(如 --ckpt-format torch)。
转化方式的原理是直接复用训练中,从 megatron 向 sglang 更新参数的函数,也就是直接复用一下训练脚本,将原先的:
ray job submit --address="http://127.0.0.1:8265" \
--runtime-env-json='{
"env_vars": { ...}
}' \
-- python3 train.py \
... # 其他训练 args
改成:
torchrun --nproc_per_node ${NUM_GPU} tools/convert_to_hf.py \
--load /your/saved/megatron_ckpt \
--output-dir /your/converted/hf_ckpt \
... # 其他训练 args
即,保持所有的参数不变,将:
- 任务启动从 ray 变成 torchrun,把 gpu 数量保存为 megatron 并行的不带 dp 的最小 gpu 数,例如如果是 tp4,就设成 4;
- 确认把
--load改成了需要 load 的路径; - 增加
--output-dir对应要保存的 hf_ckpt。
启动训练流程
整个程序需要使用 ray 进行启动,首先需要启动一个 ray 集群,即在 node 0 运行:
# Node0(HEAD)
ray start --head --node-ip-address ${MASTER_ADDR} \
--num-gpus 8 --disable-usage-stats
# 其他 Node
ray start --address=${MASTER_ADDR}:6379 --num-gpus 8
在 ray 集群启动后,可以在 node 0 提交任务,例如:
ray job submit --address="http://127.0.0.1:8265" \
--runtime-env-json='{
"env_vars": {
"PYTHONPATH": "/root/Megatron-LM/",
... # e.g. no_proxy、接口变量等
}
}' \
-- python3 train.py \
--...(其他 Megatron/SGLang/slime 参数)
参数说明
参数分为三类:
- megatron 参数:slime 会读取
PYTHONPATH中的 megatron 里设置的所有参数,可以通过传入如--tensor-model-parallel-size 2的方式配置 megatron; - sglang 参数:支持环境中安装的 sglang 的所有参数,这些参数需要以
--sglang起始,例如--mem-fraction-static需要通过--sglang-mem-fraction-static传入。 - slime 自身的参数:请见:slime/utils/arguments.py
完整使用说明请查阅 使用文档。
开发指南
-
欢迎贡献! 若有功能建议、性能调优或使用体验反馈,欢迎提交 Issue / PR 😊
-
使用 pre-commit 保证提交代码风格:
apt install pre-commit -y pre-commit install -
调试技巧请参考 debug 指南
常见 Q&A 与致谢
- 常见问题请见 Q&A
- 特别感谢以下项目 & 社区:SGLang、Megatron‑LM、mbridge、OpenRLHF、veRL 等。
