MoE 模型的训练瓶颈,早就不在参数规模上了,而在反向传播那一步。

一套典型的专家并行(EP)流程里,每个 token 被路由到少数几个专家,token-专家配对再通过 all-to-all 通信送到对应 GPU。问题在于路由是数据依赖的:某些热门专家会突然涌入大批 token,某些专家则一直闲着。论文里的观察很直白:MoE 层和 collective 通信都得等最慢的那个参与者,最忙的 GPU 成了 straggler,其他卡陪着空转。

论文在相关工作里列举的既有方案(FasterMoE、FlexMoE、SmartMoE、SYMI、LAER-MoE、UltraEP 等)大多走的是「重新分配」这条线:搬专家、复制专家(replication)、做动态放置。京东团队新近放到 arXiv 上的 LRMoE(arXiv:2610.37720,作者均来自 JD.com,北京)指出这条路已经走到头:复制出来的专家副本同样要占显存,还会和 token 相关的 buffer 抢同一份预算;在固定内存约束下,只搬不砍,总需求可能反而超过集群的总服务能力。

换个角度:反向需求不是常量,是可调度的变量

LRMoE 的核心洞察是把"哪些反向工作值得做"当成一个可以主动决策的变量,而不是前向路由确定之后的既成事实。

具体做法是:前向阶段照常跑完整路由和计算(保留全部已路由的前向工作,前向精度不受影响),但先把每个专家的反向保留配额定下来;等梯度回流时,再决定每个专家具体保留哪些 token-专家配对。

LRM:一个能分解的贡献度打分

论文定义了一个 Load Recalibration Metric(LRM)。对配对 u=(t,e),它在受保护专家投影上的权重梯度贡献是一个外积 G_u = (α_u g_t) φ_u^T。因为是 rank-1,用 Frobenius 范数平方衡量"丢掉它会造成的扰动",就能分解成三个因子的乘积:

R_u = α_u² · ‖g_t‖₂² · ‖φ_u‖₂²

  • 路由强度 α_u²:该 token 被路由到这个专家的权重有多大
  • 反向信号 ‖g_t‖₂²:回传到该专家输出的梯度能量
  • 专家状态幅度 ‖φ_u‖₂²:保存下来的投影输入的规模

论文特别强调第三项不能省。因为 G_u 是外积,两个配对可能拿到几乎相同的上游梯度范数,但路由权重或保存激活差得很远,真实贡献完全不同。

关键细节:砍掉之后还得重新打包

这里有个容易踩的坑,论文专门用一节来讲:如果只是把被丢弃的行置零、保持张量形状不变,矩阵乘的尺寸和显存占用基本没变,省不下来。LRMoE 做的是"物理压实"——把保留的梯度行和对应的保存激活重新打包成连续的、按专家分段的 buffer,用实际保留的行数驱动变长的 grouped GEMM,不做 padding。

论文图里的例子很直观:某批专家负载从 [4,2,3] 变成 [3,1,2] 之后,原本需要额外副本的专家就不再需要副本了,逻辑专家集合一个没动,减少的只是冗余物理副本。

实测数据

实验环境是 8 张 NVIDIA Hopper,EP8,PyTorch 2.11.0 / CUDA 13.0 / BF16,模型用 Qwen3-30B-A3B 和 GLM4.5-106B-A12B,序列长 8192,每卡 microbatch 4,每卡显存预算统一 80 GiB。基线包括 Megatron-LM 的固定归属、原生 UltraEP,以及作者自行复现的 SYMI(NSDI 2026)。

结论分两块。

一是收敛。丢弃 20% 反向配对(论文记作 C20)时,Qwen3-30B-A3B 最后 1000 步的平均 PPL 是 1.015141,完全不丢的基线 C0 是 1.015190;丢到 30%(C30)是 1.016270,比基线高约 0.106%。消融实验里 LRM 的平均 ΔPPL 是 0.012474 / 0.003464,而只用梯度范数是 0.097654 / 0.011987,绝对值梯度是 0.107485 / 0.016662,随机丢弃则是 0.775797 / 0.079282,差距在一个数量级以上。

二是速度。论文用"达到固定 loss 目标所需的累计 MoE 层训练时间"这个口径,报告相对 SYMI 降低 21.2%-44.8%——这个区间横跨两个模型:Qwen3-30B-A3B 上是 42.5% 与 43.2%,区间下沿的 21.2% 来自 GLM4.5-106B-A12B。具体到 Qwen3-30B-A3B 达到 0.02 损失:SYMI 是 64.6 分钟 / 8440 步,LRMoE 是 36.7 分钟 / 8693 步——步数反而更多,墙钟时间却更短,说明省下的执行时间足以覆盖多跑的那些步数。这个口径只统计 MoE 层内的执行,不含层外开销,不等同于端到端训练时间。

显存侧,三项小计(基础专家放置 + 额外副本 + token buffer)从 Megatron-LM 的 12.88 GiB 降到 5.44 GiB,比固定归属低 57.8%;token buffer 单项从 12.46 GiB 降到 4.97 GiB。反向执行的小计延迟则是 29.64ms 降到 12.69ms。

另外在 20 倍批大小压力下,LRMoE 仍保持 100% 的有效吞吐,而 UltraEP 是 58.4%、SYMI 44.9%、Megatron-LM 40.0%。

我的看法

三点值得留意。

第一,收益高度依赖"重新打包"这一步。省下显存的不是"不训练某些 token",而是"不把它们的梯度搬来搬去"。只做选择不做压实,收益会大打折扣。

第二,论文自己留了边界:被保留的更新本质上是近似的,附录专门讨论了它与整体梯度误差、与最终训练行为的关系。选择保证针对的是加性的局部目标,不等于端到端收敛等价。20% 丢弃几乎无损、30% 开始出现可见差距,这条线目前看仍偏保守。

第三,这是一个纯工程思路,但它把问题的归因换了位置:MoE 训练的算力浪费,一直默认是路由不均导致的;这篇论文把它重新归因到"你为多少工作预留了显存"。当显存成为硬约束时,先砍掉价值最低的那部分工作,比想办法把活摊匀更有效。

MoE 训练的下一个战场,可能不再是"怎么把负载做均匀",而是"怎么判断哪些负载根本不值得算"。