此前智谱开源新一代旗舰模型 GLM-4.5 以及轻量版 GLM-4.5-Air。8 月,GLM-4.5 的技术报告正式发布,在披露预训练、后训练细节之外,还介绍了智谱专为强化学习扩展自研的 post-training 框架 —— slime。摘要:此前智谱开源新一代旗舰模型 GLM-4.5 以及轻量版 GLM-4.5-Air。8 月,GLM-4.5 的技术报告正式发布,在披露预训练、后训练细节之外,还介绍了智谱专为强化学习扩展自研的 post-training 框架 —— slime。
#开发者亲自答 :
知友@朱小霖 作为 slime 开发者 ,他进一步聊了聊在做大规模 MoE 架构的强化学习训练时,slime v0.1.0 针对速度慢、耗卡多、容易爆显存等痛点上做一次性系统级急救的设计思路。智谱研发
9 月 2 日发布于知乎
在社区的帮助下,我们终于在开源 2 个月后给 slime 打上了第一个 tag:
一句话来概括这个版本:
slime v0.1.0 提供了大规模 MoE RL 训练所需的所有基本性能优化。
具体来说, 性能上:
提供了 MoE 模型的高效推理,特别是 fp8 rollout + deepep + mtp;
设计了通用的训练框架显存 offload 方案,节省出更多 kv cache 空间,提升推理并发度;
更快的参数更新;
用 cpu adam 实现用更少的 GPU 训练进行更多训练;
支持了 Megatron 的全部并行策略以及 deepep;
功能上:
针对 MoE 模型训练支持了 GSPO;
针对 fp8 rollout 支持了 TIS。
正确性上:
加入了 dense 与 MoE 模型 CI,严格检查 kl 等指标。
我们希望能通过 slime v0.1.0 展示我们对高性能 RL 训练框架的理解,并有机会成为未来性能对比的 baseline~
下面我展开介绍一下上述功能背后的设计思路。
性能优化
提升上限:优化单条数据的推理速度
对于 pretrain 和 sft 以及更传统的深度学习训练任务,有一个万能的提速方案,那就是加卡:通过降低每张卡上的计算数据,总是可以明显地降低训练任务的端到端 latency。
而这个办法到 RL 这里行不通了,因为推理的 latency 是无法通过增加 GPU 而降低的:即使我们有无尽的 GPU,我们也只能等待回复最长的那条 sample decode 结束。虽然我们可以通过增大吞吐来提升每个 rollout 的训练数据量,但是过大的推理 batch size 带来的 off policy 问题目前来看还是有一些局限性。
我认为这是当前 RL 范式下对 Infra 的最大挑战,即:
我们希望 scale inference compute,但是我们无法 scale inference latency。
单条数据的 decode 速度决定了 RL 训练速度的上限。对于较大的 MoE 模型,目前主要有以下的 3 种常规优化方案来提升这一上限,我们也在每个方向上都做了一定的尝试:
通过量化来降低访存:考虑到 RL 训练中不能进行长时间的 calibration,slime 选择做 fp8 量化;
使用 deepep low latency 模式来降低跨机 all2all 的时延:为了配合 deepep,slime 推荐使用 fp8 的 blockwise 量化来开启 sglang 的相关配置;
开启投机采样:slime 允许推理部分加载任意的 draft model(目前还不支持训练中更新 draft model。)
使用上述 3 种优化,我们可以将 GLM4.5 355B-A32B 这样的模型,从单条数据小于 10 token/s,提升至 60~70 token/s,从而极大提升 RL 训练速度上限。slime 也会在监控推理的吞吐之外,监控
perf/longest_sample_tokens_per_sec从而更好地掌握推理部分的性能优化空间。
用更少的卡做更多的实验:充分 offload Megatron
在对上限进行优化后,我们注意到 RL 训练的另外一个特性:只要 kv cache 不溢出,推理 batch size 的提升并不会明显影响训练的延时。
这里 kv cache 溢出是指推理过程中,当数据的回复长度都很长时,kv cache 空间不够,就需要把某些生成到一半的数据先踢出队列,等其他数据推理完,腾出 kv cache 的空间,再重新进行 prefill 和后续的推理。如果一条回复长度为 64k 的数据在推理过程中等待了其他数据 decode 32k token,相当于他的总时长对应了 decode 96k token,这对 RL 训练速度有很大影响。
因此,一个比较合适的训练配置是根据推理部分的 batch size、数据的平均回复长度和单个 server 能预留的 kv cache 空间,计算一个在 kv cache 不溢出情况下的最少 GPU 数量,以这些 GPU 为一组进行训练。例如我们有 512 张卡,然后计算出来 256 卡的 kv cache 就已经充足,那么就应该并行运行 2 个实验,而不是用 512 卡一起起实验。
基于这样的考量,我们注意到了 2 个优化点。
第一个是最佳的卡数可能不足以支持加载训练部分,因为推理只需要上个 fp8 参数就够了,而训练一般是要 18 倍参数量以上的显存(bf16 param、fp32 grad、fp32 master param、fp32 m 和 v)。为了解决这个问题,slime 选择开启 megatron 自带的 CPU Adam 来节省训练部分的显存。我们也是基于这样的策略提供了 8 node 训练 GLM 4.5 355B-A32B 以及 16 node 训练 DeepSeek R1 的方案。
第二个则是要提升每个 sglang server 能预留的 kv cache,也就是开大
--mem-fraction
对于现在更为常见的训推一体的训练任务,限制 mem_fraction 的主要是将训练部分 offload 至 CPU 后的残留显存,所以我们需要想个办法,正确地把 Megatron 部分占用的显存 offload 到 GPU。
如何通用地 offload GPU tensor
首先先看 Megatron 里的各种 GPU tensor。一个比较粗暴的方式是找到 Megatron 分配的所有 GPU tensor,然后把他们全部
.to("cpu")
这种做法有 3 个难点:
很难捕获 Megatron 分配的全部 GPU tensor;
由于 Megatron 的 distributed optimizer 会把所有的参数重新整理到一些连续的 GPU buffer 里,然后再通过各种 slice 划分出去,很难处理好所有的引用从而正确释放 GPU tensor;
Megatron 每次版本更新都要重新查一遍,不太好维护。
能不能有一个更通用的方案呢?
注意到 sglang 中的 torch_memory_saver 和 vllm 中的 cumem_allocator 提供了一种较为通用的 offload 方案,我们是不是能学习他们的思路呢?
他们大致的原理是,CUDA 10.2 提供了一系列 virtual memory management API,类似操作系统的虚拟地址与物理地址(va 和 pa)一样,在分配显存的时候会返回一个显存映射的 handle,而不是实际的物理地址。因此我们在 offload 的时候只需要偷偷释放这个映射对应的显存,然后在需要这段显存的时候重新分配上就可以了,上层的应用无需感知这样的变化。
CUDA 10.2 提供的 VMM api 是通用 offload 的根本
所以我们只需要将所有 pytorch 的显存分配中的
cudaMalloc
都替换成
cuMemMap
系列 API。在 sglang 与 vllm 使用的工具中,会使用
torch.cuda.memory.MemPool
来分隔出一个单独的显存池,然后使用
torch.cuda.memory.cudaPluggableAllocator来将这个新的显存池中的显存分配操作替换成
cuMemMap
系列 API,也就是:
sglang 与 vllm 的 offload 逻辑
一个自然的想法就是用这个方式接管 RL 中训练部分的全流程。这里出现了一个问题,那就是上述流程没法复用 pytorch 的 CUDACachingAllocator 了,没有 cache 极致会让显存碎片更加明显,训练过程很容易 OOM。
为了能继续复用原生的带 cache 的 allocator,我们不能使用 CUDAPluggableAllocator了。再次注意到 slime 的架构中训练和推理是在不同进程的,所以我们只需要通过 LD_PRELOAD 直接替换训练进程中 CUDACachingAllocator 使用的 cudaMalloc 和 cudaFree 为 VMM API。这样我们就可以完整且通用地 offload pytorch 分配的所有 GPU tensor 了。
同时我们还要注意一个细节,就是在 VMM API 与 cudaIPC API(例如 cudaIpcGetMemHandle)是不匹配的,所以对于训推一体的参数更新,亦或是 DeepEP,我们需要关闭 LD_PRELOAD 的替换,用回 cudaMalloc。加上这项调整之后,slime 中的 offload 逻辑为:
slime 中的 megatron offload 逻辑
在 SGLang 社区的帮助下(感谢 Tom 神!),我们针对 slime 的需求更新了 torch_memory_saver,实现了这个 offload 方案。
如何 offload nccl
在彻底 offload 了 Megatron 中的 GPU tensor 后,我们发现还会有大量的显存残留,这是 nccl 导致的。在 pytorch 中,每个参与通信的 nccl group 都会分配一份不小的 buffer。对于较大的 MoE 模型,由于引入了各种并行策略,这种问题尤为明显,可能会占到 10GB 以上。
上述 LD_PRELOAD 的方案不太好处理 nccl 的问题,我们也不想去修改 nccl 源码,避免在维护 slime 之余还要维护一个 nccl fork。所以 slime 采用的方案是在 offload Megatron 的时候利用
destroy_process_group
来销毁 nccl group,然后在 load Megatron 之前重建。为此,我们模仿 VMM API,在对
dist.new_group
进行了 monkey patch,增加了一层
ReloadableProcessGroup
大致的一个流程图为:
通过 patch new_group 实现 offload nccl
这样我们就实现了通用的 nccl offload。不过因为我们需要重建 nccl group,这样的操作会对每轮训练的第一次通信速度有一定影响,但是我们认为从可维护性以及节省的显存上,这个方案有很大的优势。
结合上述两项优化, 我们将 Megatron 的残留显存从约 15~18GB 降低到了 3~5GB,从而将 MoE 模型的 mem_fraction 提升至了 0.7~0.8,明显提升了预留的 kv cache,提升了每个 server 能支持的并发度,结合上文的分析,实现了用更少的 GPU 启动更多的训练任务。
参数更新优化
参数更新是另外一个 RL 训练中特殊的环节,这方面 slime v0.1.0 提供了训练推理在不同进程条件下的最佳优化方案,Biao He 老师在这方面进行了大量的优化,推荐阅读他的这篇博文,我就不在此赘述了:
高效强化学习训练 - 优化 slime 中的权重同步
原始英文版见:Efficient RL Training - Optimizing Weight Sync in slime。
目前 slime 可以做到 7s 完成训推一体下 Qwen3 30B-A3B 模型 bf16 权重的参数同步。100s 完成 GLM4.5 355B-A32B 的 fp8 blockwise 量化 + 参数更新。
训练优化
对于 slime 的纯训练部分,我们认为 Megatron 已经提供了充足的优化,所以主要是保证了对 Megatron 全部并行策略的适配。
适配过程中的有一个比较有意思的 bugfix:我们发现在 SGLang 开启 mtp 后,Megatron 部分无法启动 DeepEP。后来发现是因为在开启 mtp 时,SGLang 会 disable overlap schedule,从而在被 offload 到 CPU 后仍用 nccl 而非 gloo 进行某个 metadata 的通信,和 DeepEP 产生了冲突。
性能优化 check list
slime 发布以来,我经常会被问到 slime 与其他框架的性能对比,有没有 benchmark。
我对 benchmark 的理解是,benchmark 不应该成为框架之间相互攻击的武器,而是成为查缺补漏的工具。为此,我们会逐渐推出 slime 关注的性能 benchmark,用于自我提升。
同时我也认为,在跑分之前,可以从很多定性的角度来分析一个框架对性能的重视程度。这里提供一个基础优化的 feature check list:
是否支持 MoE 的训练(目前大规模实验集中于 MoE)
是否支持内部的 sglang mem fraction 或 vllm gpu utilization 能调至 0.7 以上(保证 kv cache 空间)
是否支持 fp8 或更低精度量化的推理(降低推理访存,提升速度)
是否支持训练和推理均开启 deepep(优化 MoE all2all 通信)
是否支持投机采样(提升推理时延与吞吐)
是否有高效的训练 backend,如 megatron, torchtitan,并支持各项必要的并行(复用成熟的训练优化)
slime v0.1.0 对上述的所有优化都进行了初步的尝试,当然提升的空间依然很大,所以我们也希望这个版本能成为未来 slime 版本或者不同框架之间性能对比的 baseline。我们更是欢迎所有在性能上和我们有相同追求的朋友来试用 slime,参与 slime 社区~
新算法支持
为了更好地训练 MoE 模型以及进行 fp8 rollout,我们实现了 GSPO 与 TIS。同时社区的大佬也帮忙实现了例如 reinforce++, reinforce++ baseline 这样的算法。
正确性验证
slime v0.1.0 增加了端到端 CI:我们会对每个 PR 运行单机的 GLM4 9B 和 Qwen3 30B-A3B 训练,通过一些严格的检查来保证更新的正确性。例如,我们会明确要求
第一个 rollout 的重算 log prob 和 reference model 的 log prob 完全相等 (slime/backends/megatron_utils/data.py #L205 );
以及每个 rollout 内的第一个训练步的 ppo_kl 严格为 0 (slime/backends/megatron_utils/model.py #L47 2) 。
这样的精确验证是训练框架中很少能做到的,也是我非常自豪的一点。
来源:趣闻捕手一点号