让长序列 MoE RL 训练更好调：Megatron-Lite / bumblebee 的优化实践

长序列训练不只出现在 pretrain。进入 RL 阶段以后，prompt、rollout、reward、工具调用和多轮交互都会把上下文长度拉高。很多团队关心的问题也随之变化：他们需要用较少启动卡数，稳定跑通 128K、256K，甚至更长上下文的算法实验。

RL 的资源形态和 pretrain 很不一样。Pretrain 常见目标是大规模、稳定、长时间吞吐；RL 更常见的是大量实验并行。用户手里有更多 GPU 时，往往希望同时开几组 reward、rollout、数据配方或超参实验，不愿把全部卡都压进同一个训练作业。长序列 RL 训练系统首先要给出一个低门槛、可复用、少调参的启动点。

本文以 Qwen3.5-35B-A3B / 32 x H100 / 128K 为主线。这个 case 代表长序列 MoE 用户最常见的一类诉求：几十张卡能不能尽快跑起来，并且给后续算法实验留下足够简单的默认 recipe。

文章先讲这个 case 怎么从 Megatron-style baseline 走到更小 recipe，再回头看工程问题：这些优化在 Megatron 上也可以实现，但接入 5D parallel、FSDP2、full recompute 和 chunked EP overlap 的工程半径很大。bumblebee 的 primitive + Skill 设计，把这些工作拆成了更适合人和 agent 协作的小块。

1. 问题：长序列 MoE RL 难在系统耦合

1.1 RL 长序列的资源目标不同

35B / 32 x H100 / 128K 看起来只是一个训练配置，实际会牵出一组强耦合的调节选项：

这里没有把 local sequence length 当成独立调节选项。它是全局序列长度、CP 切分、TP/EP 布局和 batch 形状共同作用后的结果。这个结果落到 16K、32K 还是 64K，会直接影响 activation、attention、MoE buffer、kernel 形状和 CPU overhead。

local seq 和计算效率关系很紧。单卡序列太短时，计算块被切碎，kernel launch、runtime submit、MoE dispatch 调度和 CPU 侧 overhead 更容易露出来；单卡序列变长后，GPU 每次拿到的连续工作更大，很多固定调度成本会被摊薄，expert GEMM 和 attention 也更容易进入稳定区间。长序列 RL 的调优，很多时候是在显存压力、local seq、通信暴露和计算效率之间找一个可复用的平衡点。

1.2 配置失败通常来自多个底层因素叠加

一个配置 OOM，往往来自几个底层属性一起被推高：

• 静态显存：参数、梯度、优化器状态、all-gather buffer；
• 动态显存：activation、MoE buffer、logits / loss 中间张量；
• 计算效率：local seq、expert token 数、batch 形状；
• 通信暴露：attention / CP 通信、MoE dispatch / combine、参数 all-gather / reshard；
• CPU overhead：小 kernel、小 batch slice、频繁 launch 和调度。

一个配置跑通但很慢，也常常是动态显存问题解决了，计算效率或通信暴露又变成瓶颈。对长序列 RL 来说，难点就在这里：每次调整一个并行选项，背后都会同时改变显存、local seq、通信形态和 kernel 形状。后文的优化会按这几类底层因素逐个拆开。

1.3 Megatron-style all-in-one 系统的真实成本

Megatron 是一个能力很完整的训练系统。它支持大量模型、并行策略、优化器、checkpoint、recompute、kernel fusion 和历史配置路径。很多优化在 Megatron 上都可以实现，底层能力也常常来自 Megatron 和 Transformer Engine 生态。

成本主要来自 all-in-one 设计带来的认知负担。模型结构、并行状态、optimizer/checkpoint 语义、recompute 调度、性能开关和历史兼容路径交织在一个大系统里。支持新 feature 或定位性能问题时，开发者需要理解很长的调用链和配置面；对 agent 来说，这种大范围上下文会明显增加任务拆解和验证难度。

以本文涉及的两类优化为例：

• FSDP2 接入 Megatron 5D parallel，远超换一个 optimizer 的范围。它需要处理参数 shard、grad reduce-scatter、optimizer state、PP/TP/EP/CP process group、checkpoint reshard、full recompute 下的参数 materialization 时机，以及和现有 distributed optimizer 语义的对齐。
• chunked EP overlap 也远比把 token 切成几段复杂。full recompute 路径下，forward、recompute forward、backward、delayed wgrad、dispatch/combine 的依赖关系都要重新梳理。少改一点可能没有 overlap，改多一点又容易破坏数值或调度语义。

这些工作可以在 Megatron 上做，代价是实现半径和验证半径都很大。下面先讲优化过程，最后再回到 bumblebee 的设计：把能力拆成 primitive 以后，这些复杂改动更容易被拆成渐进开发和 agent 协作任务。

2. 从 baseline 到更小 recipe

2.1 Baseline：可用点已经不便宜

在常规 Megatron 风格的调参空间里，35B / 32 x H100 / 128K 的代表性可用高点如下：

这个点已经经过 PP、EP、CP、重算等组合筛选。常规路径可以工作，但继续在同一个高维空间里移动可用点，边际收益和试错成本都不理想。后面要让 recipe 更容易复用，一个重要步骤是把静态显存从 PP/EP 等并行布局里解耦出来。

2.2 路径选择：接受 full recompute，拉大 local seq

长序列 MoE 训练的首要问题是显存，但显存会继续牵出计算效率、通信暴露和 CPU overhead。对抗 OOM 至少有两条合理路线。

路径 A：少重算，缩小单卡序列

第一条路径是尽量不重算，或者只做少量 selective recompute。它的好处很直接：单步里的重复计算更少。

代价同样明显：activation 留在显存里，显存预算更苛刻。为了装下 128K 全局序列，往往要把单卡序列压小，也就是把 CP 开大。CP 能切 sequence，但会引入 attention 侧 collective 通信，并且受 KV head / group 等模型结构约束。对 Qwen3.5 这类模型，CP 不能无限放大。

这条路可以继续优化，比如用 pipeline 相邻 microbatch 的计算窗口覆盖一部分 CP 通信。但它会形成连锁反应：为了装下 activation 要压小 local seq；local seq 小了以后，计算效率和 CPU overhead 都更敏感；更大的 CP 又增加 attention 侧通信和调度约束。

路径 B：接受重算，拉大单卡序列

第二条路径接受更激进的重算，把 activation 常驻显存压下去，反过来把 local seq 拉大。

local seq 变大以后有几个好处：

• GPU 每次处理的连续工作更长，CPU launch / dispatch overhead 更容易被摊薄；
• CP 可以留出余地，不必一开始就顶到模型结构约束；
• expert GEMM、attention、runtime submit 等成本更容易进入稳定区间；
• 很多原本需要细抠的小 overhead，会被更大的计算块覆盖。

这条路径也有代价。full recompute 会把一部分 forward 计算重新做一遍，在我们的训练口径里大约是 30% 额外计算量。这部分计算不能靠调度消掉，只能靠更大的计算块和更少暴露的系统 overhead 把它赚回来。同时，MoE 通信也不会因为重算自动消失。

本文选择路径 B。长序列 RL 需要一个好启动的默认点，先把动态显存和 CPU overhead 两个变量压下去，用户就不用一上来就在 PP/TP/EP/CP/recompute/offload 里盲扫。后续追求极限性能时，仍然可以在少重算和多重算之间继续细调。

沿路径 B 走，还必须补两件事：静态显存要从 PP/EP 等并行调节选项里解耦出来，MoE 通信要尽量藏进计算窗口里。

2.3 显存线：activation、logits、static state 分开处理

Step 0：常规并行和重算先到 120+

沿路径 B，先用常规并行和 full recompute 把 35B / 32 x H100 / 128K 跑起来：

TP=2, PP=4, EP=8, CP=2, full recompute
127.53 TFLOPs/GPU, peak 42.91 GB

这个点证明基础路径可行，也暴露了问题：即使有 distributed optimizer，静态显存仍然会和动态 peak 叠在一起；继续只靠 PP/EP/CP 调显存，会重新回到复杂耦合空间。

Step 1：用 linear CE 消掉 logits peak

长序列 + 大 vocab 下，如果完整物化 logits，再做 cross-entropy，loss 前后的临时 tensor 会成为 peak memory 的大项。

以 Qwen3.5 text vocab 248,320 为例。baseline 使用 TP=2 时，每个 vocab-parallel rank 约 124,160 个 vocab 列。CP=2 时每 rank local tokens 是 131,072 / 2 = 65,536。如果 loss 前物化 fp32 logits：

local_tokens * partition_vocab * 4 bytes
= 65,536 * 124,160 * 4
= 32.55 GB decimal ~= 30.32 GiB

如果不按 TP 切 vocab，同样 local tokens 下会翻到约 65.10 GB decimal。这个量级说明 logits / loss 中间张量已经是能直接改变 peak 的主项。

处理方式是 linear CE：沿 vocab 维 streaming 计算 loss 和 backward，不把完整 logits 留在显存里。这项优化有成本，loss 层本身会变慢；在我们的测量里，loss 层耗时大约变成原来的 1.3 倍。好处是长序列下显存 peak 明显降低。叠加这一类 loss-side 优化后，35B 代表性训练点进入约 160 TFLOPs/GPU 区间：

这里重点解释为什么 loss 前后的 logits peak 必须被单独处理；每个数字不必包装成单一优化的严格 on/off ablation。否则长序列下的动态显存会反复把配置推回 OOM 边缘。

Step 2：用 FSDP2 处理静态状态

linear CE 处理的是 loss-side dynamic peak，但参数、梯度、优化器状态仍然是另一块更难处理的静态显存。

baseline 里的 distributed optimizer 已经做了优化，但它没有让静态显存问题完全消失。它更接近 ZeRO-1 语义：优化器状态被切分，模型参数和梯度仍然按 PP/TP/EP 等模型并行布局常驻。到了 128K 长序列，这部分静态显存会和 activation、logits、MoE buffer、all-gather buffer 一起叠到 peak 上。

从训练状态切分看，Megatron distributed optimizer 和 FSDP2 可以先按 ZeRO 语义理解：

• distributed optimizer 更接近 ZeRO-1：optimizer state 已经按 data-parallel 维度切分，但参数和梯度仍按模型并行布局常驻；
• FSDP2 更接近 ZeRO-3：参数、梯度、optimizer state 都进入全局 shard，计算当前 FSDP unit 时再 all-gather 需要的参数。

如果 nGPU 指参与训练的全局 GPU 数，n_param 指全模型参数量，optimizer state 在 distributed optimizer 下已经避免每卡完整重复，可以粗略估成：

optimizer state per GPU ~= n_param * 12 / nGPU

真正需要继续往下压的是参数和梯度。由于 Qwen3.5 MoE 参数占大头，distributed optimizer 下参数和梯度仍主要跟 PP、EP 的切分绑定：

dist_opt parameter + grad per GPU ~= n_param * 6 / (PP * EP)

FSDP2 的目标是把参数和梯度也推向全局 shard：

FSDP2 parameter + grad per GPU ~= n_param * 6 / nGPU

在 35B、32 x H100、128K、PP=2, EP=8, CP=4、full recompute 的同主并行度对照中，FSDP2 基本没有牺牲吞吐，同时把 peak 降了约 8.88 GB：

这个结果说明，把参数和梯度继续 shard 以后，静态显存余量明显变大。PP 和 EP 也不必再默认承担省静态显存的职责，可以回到更清楚的位置：

在这个思路下，后续可以尝试 PP=1, EP=8, CP=4。纯 FSDP2 no-chunk 口径下，PP=1 的历史高点约为 180.18 TFLOPs/GPU、60.54 GB peak。再叠加后文的 chunked EP overlap 后，同一主配置进入约 37-38 GB peak、接近 190 TFLOPs/GPU 的区间。

到这里，路径 B 里的显存代价被拆成三块：activation 靠 full recompute，logits 靠 linear CE，static state 靠 FSDP2。剩下的主要问题是 MoE all-to-all 通信暴露。

2.4 通信线：chunked EP overlap

在长序列 + MoE 下，计算已经很密，但 expert parallel 的 token dispatch / combine all-to-all 仍然会暴露在时间线上。

直接能想到的办法是 1F1B overlap：用 pipeline 相邻 microbatch 的计算窗口去盖 dispatch / combine。但本文选择的显存路径默认 full recompute，旧的 1F1B overlap 不能直接支持这条路径；如果为了复用它而退回少重算路径，又会把前面避开的动态显存、local seq 和 CPU overhead 问题带回来。

所以这里重新设计了一个支持 full recompute 的 MoE EP A2A overlap。做法很直接：把 token 维切成多个 chunk，用通信流和计算流错位执行，让相邻 chunk 的 dispatch / combine 尽量落到 expert compute 窗口里。原来串行暴露的通信，有机会被相邻 chunk 的计算覆盖。

这里先解释两个后文会反复出现的术语：

• no-chunk：一个 MoE layer 把当前 microbatch 的 token 一次性 dispatch 到 expert，完成 grouped GEMM 后再一次性 combine 回来。时间线上可以简化成 dispatch -> grouped GEMM -> combine。
• chunk2：沿 token 维把同一个 MoE layer 的输入切成两个 chunk。每个 chunk 都有自己的 dispatch、expert compute 和 combine。调度器用通信流和计算流错位执行，让前一个 chunk 的 combine 或后一个 chunk 的 dispatch 尽量落进相邻 chunk 的 expert compute 窗口。

更大的 chunk 数也可以尝试，比如 chunk3、chunk4，但 chunk 越多并不天然越好。chunk 太多会增加 kernel inflation、runtime submit 和调度开销，也可能让 expert GEMM 变小，计算效率下降。本文重点解释 no-chunk 到 chunk2 的基本机制，以及长序列下为什么 overlap 收益会变大。

Kernel 排布：收益来自真实时间窗

forward 的 no-chunk 排布可以简化成：

dispatch -> grouped GEMM -> combine

chunk2 之后，通信流上 c0 dispatch / c1 dispatch / c0 combine / c1 combine 串行排布，计算流上 c0 GEMM / c1 GEMM 串行排布。依赖关系是：

c0 dispatch -> c0 GEMM -> c0 combine
c1 dispatch -> c1 GEMM -> c1 combine

在 16K forward 的关键 rank 上，no-chunk measured total 是 10.36ms，chunk2 是 7.88ms。这个收益不能理解成“切成两块就盖住一半通信”。真实 trace 里，c0 dispatch / c1 dispatch 都早于 expert window，forward 里真正可见的 overlap 是 c0 combine 被压进 c1 GEMM；c1 combine 仍然形成尾部。

backward 的语义顺序反过来：先 combine-bwd，再 expert backward compute，最后 dispatch-bwd。在对应 16K backward 窗口里，no-chunk 是 13.18ms，chunk2 是 11.77ms。这里不能只把计算画成一整块 “expert backward”，还需要拆成 dgrad 和 delayed wgrad。delayed wgrad 必须发生在最后一个 chunk 的 dgrad 之后，并且可以和尾部 dispatch-bwd 形成 overlap。

进一步把 recompute forward 和 backward 接起来以后，收益还会来自重复计算的减少。fused 路径可以让 backward 直接消费 recompute expert output，recompute forward 里的部分 combine 可以被移除；同时，recompute 阶段也不必再执行原本用于生成 MoE layer 输出的 fc2。这相当于把 recompute forward 里一段原本只为 backward 准备中间结果的尾部计算省掉。

这一步需要重新整理 forward / backward 的依赖关系。它改的是数据生产和消费的位置：recompute forward 只保留 backward 真正需要的 expert 中间结果，后续的 combine 和 fc2 交给 fused backward 路径消费或跳过。这样既让通信和计算窗口更好地贴合，也减少了 full recompute 路径里一部分重复计算。

Sequence 越长，overlap 收益越明显

为了隔离 MoE EP A2A overlap 的局部趋势，我们先看 35B-A3B sparse MoE 单层 proxy。这个 proxy 覆盖一个 sparse MoE layer 的 forward + backward，不包含 attention、dense MLP、optimizer 和其它 transformer layer，不能直接当作端到端 train step 收益；它主要用来解释 overlap 在不同 sequence length 下的趋势。

趋势很稳定：sequence 越长，通信窗口越大，chunked EP overlap 的收益越高。chunk 数、EP size、非对称 chunk 权重仍然要谨慎处理。chunk 太多会带来 kernel inflation 和 runtime submit，EP 太大也可能被跨机通信衰减吃掉收益。这里的结论很具体：长序列 MoE 里，EP A2A 需要进入调度层面处理；具体 chunk policy 仍然要根据 sequence、EP size 和 kernel 形状选择。

2.5 合体结果：从 4D 盲扫到 recipe

走完显存线和通信线以后，这条路径才站得住：

1. 不容易 OOM：显存账被 full recompute、linear CE、FSDP2 分别处理；
2. 性能足够高：MoE all-to-all 暴露被 chunked EP overlap 处理。

显存线只解决能不能跑，通信线决定跑起来以后是否够快。两条线合在一起，才有一个可用的长序列 RL 默认点。

35B 的关键路径可以按真实调优顺序理解：

这张表主要用来说明调优顺序：baseline 先钉住；linear CE 处理 logits；FSDP2 在同主并行下证明静态显存可下降；然后用 FSDP2 带来的静态余量选择更好的 PP=1, EP=8, CP=4；最后叠加 chunked EP overlap 处理 nsys 里暴露出来的通信时间。

特别需要区分的是，37-38 GB peak 来自 FSDP2 叠加 chunked EP，而非纯 FSDP2 no-chunk 的显存收益。

旧空间里，用户要同时调 PP/TP/EP/CP 四个并行选项，而且每个选项都影响显存、通信和 kernel 形状。本文路径的目标，是把默认调参压成小得多的 recipe：

变化可以概括成三条：

1. TP 从默认 recipe 里去掉；
2. PP 和 EP 不再被当作省静态显存的主要工具，静态显存由 FSDP2 独立处理；
3. 默认只需要按模型规模定 EP，再按目标单卡序列定 CP；追求极限性能时再调 PP。

这张图是总结图，并非 measured cell。它回到最开始的底层属性：动态显存、CPU overhead、静态显存、计算效率、通信暴露都要有默认处理方式。只有这些问题分别有解，recipe 才不会变成另一组复杂调节选项。

3. 为什么 bumblebee 适合这类优化和 agent 协作

3.1 Megatron 也能做，接入成本不同

linear CE、FSDP2、chunked EP overlap 都可以被搬到更重型的训练框架里，底层 kernel 和并行语义也大量依赖 Megatron / Transformer Engine 生态。

差别主要在接入成本和验证成本。

FSDP2 接入 Megatron 5D parallel，需要同时面对 PP/TP/EP/CP/DP process group、参数 all-gather、梯度 reduce-scatter、optimizer state、distributed checkpoint、recompute、MoE expert 参数、embedding/head 共享路径，以及已有 distributed optimizer 的语义。任何一处处理不一致，都可能表现成 loss、grad norm、peak memory 或 step time 的异常。

chunked EP overlap 也类似。它要动的是 MoE dispatch/combine 和 expert compute 的时间关系。full recompute 下还要考虑 recompute forward 和 backward 的依赖、delayed wgrad 的位置、不同 chunk 的 buffer 生命周期、通信流和计算流同步点。这个优化看起来像局部 schedule，实际牵动的是 MoE primitive 的完整 forward/backward 语义。

这些复杂度不会因为换一个项目名就消失。bumblebee 把改动边界缩小，让开发者和 agent 能把注意力放在当前 primitive 上，无需一次性吃下整个 all-in-one framework。

3.2 Primitive 边界压低改动半径

这里的 bumblebee 是 Megatron-Lite 的内部代号，也可以理解成这套轻量训练实现继续沉淀后的 library 形态。模型由一组 primitive 组合出来，避免继承一个不断膨胀的大模型类。Runtime 管训练协议，Primitive 管可替换能力，Model 负责选择和组合。

放到本文三个优化里，对应关系很清楚：

• linear CE 是 loss primitive 的替换；
• FSDP2 是 optimizer / state sharding primitive 的替换；
• chunked EP overlap 是 MoE communication primitive 的替换。

每个替换都有相对清楚的输入、输出和验证口径。模型通过组合 primitive 获得能力，优化也沿着 primitive 边界逐步落地。

3.3 Paired baseline 让局部优化更容易验证

bumblebee 保留 paired baseline。旧 primitive 和新 primitive 可以互换，同一套模型、同一份输入、同一组并行配置下，对比 peak memory、step time、loss、logits、grad norm 和关键中间 tensor。

这对复杂优化很重要。FSDP2 改的是参数和 optimizer state 的生命周期；chunked EP 改的是通信和计算的调度顺序。它们都可能同时影响性能、显存和数值。paired baseline 让每次替换都有清楚参照，开发者不用在完整训练栈里猜变化来自哪里，agent 也可以把任务拆成“读 Skill、改 primitive、跑 paired test、汇报差异”这样的小闭环。

3.4 Agent-native：小边界、Skill 和渐进式任务

Agent 很难一次理解三万行级别的 all-in-one framework。它更适合处理边界清楚、输入输出明确、验证方式固定的局部任务。

bumblebee 的 agent-native 设计可以落到实际流程里：

1. 每个 primitive 都可以配一份 Skill，说明它做什么、怎么组合、怎么验证、有哪些 trade-off；
2. 每次开发只需要处理一个局部替换，再用 paired baseline 验证；
3. 模型可以渐进式生长，不要求一次支持全部 feature。

这就是 bumblebee 适合 agent 协作的地方。它没有降低 FSDP2 或 chunked EP 本身的理论难度，但把复杂优化拆成了更小、更可验证、更容易回滚的工程任务。

4. 小结和后续

长序列 MoE RL 训练难调，原因很具体：PP/TP/EP/CP/recompute/offload 会同时改变显存、通信、计算效率和 CPU overhead。本文选择的路径是先接受 full recompute，把 local seq 拉大，用更稳定的计算块降低动态显存和 host overhead 的复杂度；再用 linear CE 处理 logits peak，用 FSDP2 处理静态状态，用 chunked EP overlap 处理 MoE all-to-all 暴露。

最终，用户面对的默认空间从 PP/TP/EP/CP 的 4D 盲扫，收敛成“按模型规模定 EP，按目标单卡序列定 CP，必要时再调 PP”的更小 recipe。这个 recipe 没有消灭所有调优问题，但把问题拆成了更清楚的几类：activation、logits、static state、MoE communication、local seq 和计算效率。

工程层面，这篇文章讲的是复杂优化怎样落地。FSDP2 接入 5D parallel、full recompute 下的 chunked EP overlap，都可以在 Megatron 上实现，但工程复杂度很高。bumblebee 把这类改动落到 primitive 边界里，再配上 paired baseline 和 Skill 化知识，让人和 agent 都能更快理解、修改、验证和组合。

后续还有几类工作值得继续推进：

1. 扩展显存估计器。当前分析重点覆盖 activation、logits、static state 和 MoE buffer，后续还需要把 FSDP all-gather buffer、当前 unit 参数、workspace、fragmentation、rankwise peak 纳入统一估计。
2. 细化 FSDP2 和 distributed optimizer 的适用边界。FSDP2 更省静态显存，但引入参数 all-gather / 梯度 reduce-scatter；distributed optimizer 更简单、通信少，但参数和梯度仍被 PP/EP 绑定。
3. 把 FSDP2 + 5D parallel + chunked EP 沉淀成更通用的 pretrain / RL 长序列调优指南，特别是 100B 和 1T+ 模型下 PP、EP、CP、FSDP、TP 的组合边界。
4. 继续优化 chunked EP，包括减少显存碎片、尝试更融合的 expert kernel、把 dispatch / expert compute / combine 的调度推进到训练版 mega MoE kernel。
5. 继续细调 RL perf。当前默认把 local seq 拉到 16K-32K 来摊薄 CPU overhead 和 kernel launch，但 RL 端仍可以根据具体任务，继续寻找吞吐、显存和重算之间更好的 trade-off。

致谢

感谢 Zijie Yan、Hongbin Liu 和 Kunlun Li 在这项工作中的支持、讨论和 review。