作者：sallon88
原文：https://zhuanlan.zhihu.com/p/2093494438676927321
前段时间组里接了个活儿，把DeepSeek V4.1 Flash的Agentic RL训练在我们自研的LoongSage（https://github.com/baidu-baige/LoongSage) 框架内拉起来。
一开始想着，框架本身已经支持DeepSeek V4的Agentic 训练了，而且SGLang社区也放出了day0 (https://www.lmsys.org/blog/2026-09-10-deepseek-v41/) 的支持，再支持个4.1能有多难。但是真正上手才发现，每走一步都是血泪。。现在整条链路总算通了(r2e gym + mini-swe-agent, 24k response长度，训推一致保持万分位)，趁着记忆还热乎，把这些没写在官方文档里的工程天坑盘一盘，希望给做同类方向的同学一些参考，避免在同样的坑上再踩一遍。。
权重的格式转换
为了保证训练的稳定性，训练和推理我们统一用bf16精度， dsv4.1引入了多模态，还需要将其转换为纯文本模型。但是官方和社区并没有现成的转换脚本，只能摸索着自已写。这里有几个地方需要特别注意：
转换后hf model config中的quantization_config删掉，要不然SGLang加载会有问题； 删掉vision_config并把原有的text_config内容提到上一级； 删掉mtp/dspark相关的配置(暂时不参与训练)；
少了这些后面SGLang加载的时候都会遇到各种莫名其妙的错误。
推理引擎跨机部署
权重转换好了，接下来就会撞到显存的问题。SGLang单机8卡部署的情况下，权重load后会占用显存130G以上，还需要预留kv cache, 共置场景下参数传输时还需要为megatron预留显存, 因些比较保险的办法是用两机部署SGLang, 即设置ep=16。
然后问题来了，tp不能开到16，最大只能是8，原因是liner_o_group_proj这个权重的切分方式决定的，不能继续往上切。tp到头了，剩下的并行度只能靠dp attention来补，于是dp=2成了唯一选择(不要考虑pp, 开了pp，router replay就不好使了) 。
本以为给dsv41开启dp attention这个成熟功能只是个常规操作，但是在这条路径上有两个隐藏bug(这两个我们都已经给社区提了修复PR)：
- 启动后一收到健康检查就挂
。开启dp attention后，收到prefill token数必须得是attention tp的整数倍，不是的话必挂。健康检查发来的探活请求token数肯定对不上，于是服务一起来，刚收到健康检查，就立即挂了，表现得就像是启动失败，很有迷惑性 - 关掉cuda grpah时decode永远卡死
。调试阶段习惯性的把decode的cuda graph关掉方便定位问题，结果直接假死。追查下去发现是漏塞了一个idle batch, 整个decode循环只能死等下去。。
由于前面做了权重转换，碰到这两个dp attention新路径上的bug，排查的时候很容易把方向带偏，一会儿怀疑权重一会儿怀疑网络，绕了不少弯路。
避免kv cache重算
最终dp attention开起来之后，又冒出一个更隐蔽的问题。LoongSage框架内置的router是不感知多dp的，它只向某一个SGLang engine发请求，engine在开启dp attention时会自动将请求负载均衡到自身的多个dp。这在agentic多轮的场景下会极其低效，原因在于kv cache的复用，当一条轨迹的多轮请求在多个dp间来回跳时，会带来极其严重的kvcache重算。 好在SGLang本身支持把请求指定到特定dp, 我们顺着这个机制做了适配，将LoongSage的sticky routing机制从engine下沉到了dp层面，才算解决kvcache的复用问题。
加速启动把engram搞没了
框架启动慢是另一个让人头疼的点。SGLang启动时，默认情况下每张卡都要读取完整的权重，几十张卡一起读，如果权重本身是在共享文件系统上，每启动一次可能都要在小时级别。一个很自然的优化是，框架拉起SGLang时把load format指定成dummy, 跳过权重加载，等megatron这边起来之后再通过权重同步把真实权重同步过去，这样启动时间能砍掉一大截(仅由megatron启动时读取自身rank所需权重)。但这里埋了个雷，权重同步是不传engram的(冻结参数且存放在内存中)，用dummy模式跳过了权重加载，engram也跟着一起跳过了，结果就是engram权重丢了。我们临时的解决办法是patch SGLang的代码，在dummy模式下单独把engram加载进来，避免它随其它权重一起被跳过。
engram引入的坑不止一个， engram通常是放到主机内存里，它的host内存布局默认是shared, 也就是本机一块全局share memory，多卡共用，单机这没问题，但是跨机的时候这套就不成立了，必须显式将engram布局设置成private, 也就是sharded模式, 让每个rank各自load自已负责的那一部分。
主机内存瓶颈
训练dsv4.1, 最大的瓶颈其实是主机内存。engram本身会占用大量主机内存, 但另一个更容易被忽视的因素是R3. SGLang启用专家回放时，需要提前为每个token预留内存存放routed experts. 在dsv4.1场景下每个token需要960B(40x6x4)， dsv4.1对kvcache的显存占用做了极致优化，每个token的kv cache占用差不多也是这个量级, 换句话说你如果每张卡预留20G显存放kv cache, 那就得同时预留20*8=160G主机内存存放routed expert. 所以共置部署时如果不小心把kv cache池开的很大，开启r3后，对应的主机内存占用会跟着水涨船高(8倍)，最后不是显存oom, 而是主机oom。我们是在配置中显式指定max_total_tokens, 把kv cache池的大小明确下来。
另一个内存占用的大头是cpu优化器。训练dsv4.1这种级别的大模型，基本都会开启cpu optimizer, 把优化器状态放到主机内存以节省显存。但是由于engram(SGLang一份，megatron也得有一份)以及上面r3的叠加影响，会更容易踩到主机内存oom的问题。根据我们的实测，原本可以正常跑dsv4的6机h20(单机2T内存)，是无法正常跑dsv41的， 每次都会挂在cpu optimizer的更新权重环节。经过一翻尝试，最终用megatron最新的chunked_optimizer_state_offload机制才能跑起来，这个功能当前仅存在于megatron的dev分支，它同样是把优化器状态放入主机内存，但是主机内存的峰值占用，可以降低三分之一，从18N（N为参数量， 主权重4N, 主梯度4N, 动量相关8N, bf16权重的pin memory 2N）降到12N。 另外，最好在PYTORCH_ALLOC_CONF中配上 pinned_max_cached_size:32。因为 pin memory 申请内存时会向上按 2 的幂次取整——你只需要 11G，但默认情况下它会向上取整申请 16G(这个向上取整是为了方便pin memory释放后下次申请复用， RL的场景几乎用不到)，这多出来的部分在本就紧张的主机内存上会成为压垮骆驼的最后一根稻草。上面的配置就是避免这样的情况发生。
混合精度下的续训问题
chunked_optimizer_state_offload 开起来之后，如果同时开了 fp32 lmhead，续训时会发现之前存的 checkpoint 加载不回来了。本质和之前 开启cpu optimizer 遇到的问题是一类——megatron优化器没有正确处理混合精度的场景，fp32 和 bf16 的状态对不上。我们在框架里通过 monkey-patch 的方式做了临时修复解决这个问题。
写在最后
从权重格式到显存占用，从SGLang的隐藏bug到主机内存的精打细算，这一路踩下来，最大的感受是 把一个新模型的RL训练真正跑起来，难点往往不在模型本身，而是在大量不容易被写进文档的琐碎工程细节里。目前完整的修复实现和训练配方都已经内置到LoongSage框架中(LoongSage/conf/dsv41_flash_bf16 at main · baidu-baige/LoongSage：https://github.com/baidu-baige/LoongSage/tree/main/conf/dsv41_flash_bf16)，如果你也在做类似的事情，也欢迎一起交流，少踩点坑。
附dsv41 swe训练的关键指标