仅用于站内搜索,没有排版格式,具体信息请跳转上方微信公众号内链接
↑↑↑关注后\“星标\“kaggle竞赛宝典
kaggle竞赛宝典
作者:大润发杀鱼工,文章摘自包包的算法笔记
训练DeepSeek-v3671B!超大模型RL训练的挑战和难点!
因为之前实习一直太忙,导致公共媒体荒废了一段时间…趁着回学校参加毕业答辩&毕业的契机,也想总结一下自己实习两个月的心得。另外本文所提的只是这个过程中的的一些坑点小结,贡献并非来源个人,感谢团队以及开源社区对671B任务的支持!
从开始在verl组实习,刚入职LD就给我提了需要一起推动671B超大模型的RL训练的目标。一种表情包式的,“啊,你说我吗?”的心情在内心翻滚。当时学习RL也不到2个月。一下让我支持这么大的任务,既兴奋又焦虑,到底能不能做好呢?从3月底入职到现在,经常会听到别人问我们,verl什么时候可以支持671B的训练呢?我想,我们一直在全速的推进,只是这中间潜在的问题并非视野扫过即可见。
下面我将用我自己的视角,来给大家阐述一下,对于超大模型RL训练的挑战和难点。
首先,与传统的单纯训练与推理流程相比,强化学习(RL)的训练框架具有明显的异构性。一方面,它需要借助推理框架对prompt进行rollout;另一方面,还需依赖训练框架完成log计算与梯度更新。这种双重依赖意味着整个任务流程需要协调至少两个不同系统,显著增加了系统集成与工程实现的复杂度,工作量几乎翻倍。
在具体实现中,目前主流的训练框架包括FSDP和Megatron(verl暂不支持DeepSpeed);推理框架则有vLLM、SGLang等可选。对于671B规模的模型,训练框架的选择非常关键,需要在通信效率、显存管理能力以及并行策略支持等方面进行全面考量。
目前verl的框架搭配支持:
框架搭配
fsdp
Megatron
vllm
支持
支持
SGLang
支持
支持
从训练角度来看,671B模型通常需要256张以上的GPU参与训练。以FSDP分片计算结构为例,其在大规模场景中存在通信效率和显存利用率问题。在一些开源实践中,即便使用160张卡训练405B的DENSE模型,也频繁出现OOM情况,若未进行精细的显存管理,训练过程常常难以维持稳定性。相比之下,Megatron凭借其成熟的模型并行策略、优化手段,以及NVIDIA内部TransformerEngine框架的深度集成支持,始终是超大规模模型训练中的首选。因此,综合分析后,Megatron是671B模型开源训练的优选方案。
在推理方面,vLLM与SGLang均已在大规模应用场景中得到广泛验证,并在工业界获得良好口碑。verl框架对这两种推理引擎均提供统一参数接口,完全兼容DeepSeekV3模型结构,框架间切换成本较低,工程落地灵活性高。因此,从推理部署角度出发,无论选择vLLM还是SGLang,均可作为671B模型的推理后端。
假如每位开发者都拥有1024张GPU卡,那么成功运行一个671B的模型并非难事。只需按照DeepSeek论文中提供的并行配置,或参考NV实践过的脚本,即可较为顺利地完成部署。然而在大多数实际场景中,可用的资源往往没有这么多。因此,为了尽可能在资源受限的条件下完成训练任务,就必须将更多节省显存的特性集成到verl的训练流程中。这些功能的支持也可分为训练侧与推理侧两个部分。
在训练侧,由于verl的早期开发是在去年,当时对Megatron的支持仍停留在0. 4版本,且需依赖patch实现功能扩展。随着NVIDIA协助将verl升级至Megatron0. 11,首要任务是废弃原有对Megatron的非标准调用逻辑,并重新构建对参数、梯度、优化器等组件的offload特性支持。然而在调研过程中发现,Megatron并未提供原生的offload实现,因此需手动实现相关功能。进一步地,由于offload时缺乏参数flat操作,导致显存碎片化问题尤为突出。最终我们通过在训练脚本中添加环境变量以限制碎片大小,从而缓解该问题。此外,虽然禁用cachepool也能缓解碎片化,但vLLM在0. 8.2版本左右并不支持该操作(sad)。这一限制进一步说明,在RL阶段实现过程中,训练框架与推理框架之间必须高度兼容,才能避免资源管理瓶颈。
在推理侧,显存占用主要来源于模型参数与KVcache。尽管vLLM早期已支持wake/sleep操作用于模型&数据的显存加载与卸载,但如果加载时机不合理,仍会引发OOM。以0. 8.2版本为例,其在wakeup时将参数和KVcache一并加载至GPU,而此时KVcache实际尚未被使用。这种不合理的内存调度导致了显著的资源浪费,最终引发显存溢出。为解决该问题,vLLM从0. 8.3版本起引入了分步式wakeup策略,verl也在第一时间进行了适配支持。
综上所述,verl的开发难点并不完全在于其自身框架的实现,而在于必须确保上下游训练与推理框架能够提供完善的接口支持,才能使verl更高效地解决训练过程中面临的实际问题。
模型权重的转换,我认为主要会分为两个方面。首先是模型的checkpoint加载和存储。第二个则是训练框架和推理框架之间的模型权重转换。
首先,verl的模型加载是从训练框架侧加载的。对于checkpoint加载和存储来说,一般fsdp的框架可以基于huggingface的模型直接加载。但是对于Megatron来说,它的模型必须要来源于Megatron本身的模型结构。这也就表明如果我们想跑671B的模型,那么就需要首先将huggingface上671B的模型转换为Megatron可以加载的形式。但是Megatron本身也有两种。首先是legacycheckpoint的格式,这种老版本的模型存储是这个模型结构会和Megatron的并行策略强绑定,如果你在第一次使用固定策略进行存储后,那未来你只能使用相同的策略来进行训练,否则将无法加载模型权重。
NV后来也是发现这种存储方式的弊端,所以Megatron后续又开发了基于dist-ckpt存储方式的结构。这种存储结构的好处在于存储的并行方式和之后训练加载的并行方式可以不相同,解决了之前比较生硬的加载逻辑。目前,对于671B模型的训练来说,由于模型比较大,NV内部的支持也只能让huggingface的模型先转换为老legacycheckpoint的模型,然后再转换为dist-ckpt。这个过程中就需要256张显卡才能完成这个任务,总的来说并且也比较复杂。不过由于使用了dist-ckpt,模型在存储的时候压力也会减小一些。
从训练框架和推理框架之间的权重转换,是RL场景下特有的训练阶段。在verl的具体流程中,训练过程中需要先使用推理框架进行相关推理。然后当获得prompt+response的结果之后,再将这些模型权重切换到训练框架中。我们知道权重实际上是以字典的形式进行存储的,kv的结构分别是{name,tensor}。但是由于Megatron和vllm&sglang之间对于权重的key命名并不相同。这也就导致我们在切换的时候实际上是需要对两边进行一个reshard的操作,并且也有许多需要解决的问题。
更细节的举例来说,比如在Megatron的attention计算中,QKV可能会被调整为一个Tensor以加快训练,但是在推理模型中他们实际上应该是三个Tensor。那这些细节都需要被一一对应上,才能够完成不同框架之间的权重对应。
进一步的假如推理框架和训练框架之间的并行策略并不相同,那如何从训练框架将数据导出,然后传给推理框架。也是一个问题。目前的解决策略,实际上是首先在训练框架端的每一张显卡上获得模型每一个Tensor的完整size。然后再将这个完整的Tensor传给推理框架。推理框架可以再进一步去识别每张卡上应该获得这个Tensor的某一个部分。这样做的好处是,让不同框架侧在内部解决并行策略带来的Tensor差异,在框架间进行权重传递时,保证按照模型正常的shape进行,直接跳过了并行策略之间的对齐工作。
权重的转换有点类似于拷贝操作,常规的做法,可能是需要先开辟一块额外的空间,然后将原来的数据调整搬运到新空间中,最后释放前面的空间。但是在模型权重的reshard时,这个操作对于显存的代价是不可接受的。目前的做法则是前面有提到的,由于整个模型结构会有若干个Tensor,我们按照注意的Tensor以生成器的lazy_load方式去进行传递。就可以尽可能的减少在权重转换过程中的额外显存开销。显存的峰值永远都只会是一个Tensor大小的波动。
模型的精度问题实际上比较好理解,从算法层面上来说,RL的训练应该是没有问题的。但是在实践的过程中往往会因为框架之间内部存在一些对齐问题,或者说人为写出了一些bug导致的。为了验证新开发特性的有效性,一般的开发中都需要先进行精度对齐的操作,就是先用一个已经确认是准确的框架,跑一遍相关算法,然后再使用我们开发的版本去再跑一次,以保证精度上的对齐。
但实践起来,还是比较的骨干。
做过精度对齐的朋友应该知道这实际上是一个非常困难的操作,一旦发现模型的训练精度没有直接对齐,将会导致后续非常麻烦的debug操作。这样定位问题往往快则几个小时,多则好几天时间。并且即便对齐本身没有问题,但运行benchmark和新特性也需要一段时间。
在做moe相关的训练时,实际上有一个比较明显的问题,找不到相关可以运行收敛的模型。即便到今天,对于dpskv3的结构来说,依然没有可以使用的小模型可以用于RL训练。即便目前有的字节或者kimi推的dpskv3结构,都是针对推理适用的,在其框架中是直接assertnottrain。导致即便有心想要对齐,依然会拔剑四顾心茫然。
模型的效率问题,实际上就是虽然说671B模型最后如果可以跑起来,那它只证明了verl这个框架在跑大规模的模型是可行的,系统的健壮性毋庸置疑。但如果要真的想让这个工作成为一个工业界可用的任务,能够使得其他公司能够在verl上面跑更大模型。就必须要让模型训练的效率得到有效提高,换句话说,也就是verl的训练速度/MFU不能非常差。
所以对于DeepSeek这种moe的模型来说,在训练框架&推理框架中加入对于专家并行的支持,或者容纳更多的并行策略,都是比较重要的事情。虽然我们也知道在训练框架或者是推理框架中,针对模型效率提高有着非常多的技巧。例如在训练框架中使用多种并行方式或者overlap操作。在推理侧使用多种并行策略,pd分离,ep并行等。那如何将他们有效地包含在verl的框架中,使得这些不同的策略能够朝一致的方向去提高模型的效率也是需要探索的。
简单的例子,在不同的环境和训练模型下:
(1)在训练侧,机内的TP可能比PP好,机内的EP也可能比TP好
(2)推理侧,有时候跨机TP可能比TP+PP效率来的高,有的时候PD分离往往不如不用PD分离等…即便最后verl成功的将671B模型成功跑起来,我们的并行策略搭配也不一定就是最高效的。模型的训练效率优化,一定是一个需要花费大量时间去探索的事情。(谢谢,赏了我一口饭吃)
更多的现实问题来源于verl或者说RL训练流程。我觉得verl这个框架类会似于一个粘合剂,它将不同的训练框架和推理框架拼接在一起。所以一旦这个框架产生了一些问题或者是一些bug,如果我们发现这本身并不是verl代码带来的。类似前文提到的一些框架的特性,需要等待新的版本才能支持。那么verl对于这件事情可能也是无能为力的。我们能做的也只有尽快去对齐上下游框架的问题,协调他们来解决问题。(部分能patch的就先修了,以免耽误正常训练)
比如我所遇到的几个事情
(1)在sglang在框架初始化时,会有一个显卡显存平衡性检测阶段,但是在verl中由于RL和ray的特性,会经常遇到平衡检测无法通过,导致SGLang框架初始化直接失败的问题。那我们就只能联系SGLang的开发人员去提供一些解决方案。
(2)在verl使用vllm的权重转换中,vllm对于moe模型的模型权重加载支持,其实是存在潜在的bug(好久没看了不知道修了没),但是因为推理框架的使用方法和verl权重切换逻辑并不相同。vllm的开发者在使用过程中并不会遇到这个问题。所以我们只能够通过打补丁并且在slack中对接框架开发者的方式来尽快推动这个bug的修复。
对于我来说,实际上在verl实习的这两个月体验还是非常不错的。组里的同事都很强,往往在我解决不了问题,或者对于一些实践还缺乏经验时,他们总是能够给我非常多的指点。另外就是在于其他框架开发者对接时,也认识到了很多大佬,比如之前在知乎一直刷到的游凯超,或者是虽然早早在小红书就认识,但是后面发现居然还有工作交接的chenyang,以及SGLang社区,领英,AMD,美团,Qwen团队的更多大佬。
对verl来说,671B是我们肯定会拿下的一个工作,只是说首先于未探索道路的隐藏问题和面向开源的友好,导致我们需要一些时间才能完成他。