为什么 Dropless MoE 需要通信感知的 JAX 运行时

Dropless MoE 不只是模型设计，也是运行时问题

NVIDIA 展示了 JAX 与 Transformer Engine 组合优化 Dropless MoE 训练的路径。其核心启示是：不规则专家计算、专家并行通信、低精度缩放、激活存储和编译器调度必须作为一个系统评估，单独优化某个 kernel 并不能保证整体扩展性。

Token4AI Cloud2026-09-15约 4 分钟阅读

THE BRIEF

博客要点

据 NVIDIA Technical Blog，DeepSeek-V3 671B 在 GB200 上的一个 JAX 基线从每 GPU 103 TFLOPS 提升到 1,068 TFLOPS，基线累计 kernel 时间中 84% 被归因于 GPU 间通信。文章还称，在 GB300 NVL72 上扩展到 1,024 GPU 时达到 97% 效率。NVIDIA 将结果归因于 grouped GEMM、MXFP8 分组量化、融合的专家并行 dispatch/combine、主机激活 offload 和 XLA 多 stream collective。这些是厂商报告，尚未独立验证。

来源公告日期 · 2026-09-14

Dropless MoE 不设置固定专家容量，每个被路由的 token 都继续计算。它避免了 token 丢弃，也避免了为统一形状而 padding，但把负担转移给运行时：每个 batch 的专家 token 数都可能不同，计算、内存搬运、同步和通信必须共同处理动态形状。真正的瓶颈因此不是稀疏矩阵乘法本身，而是动态路由造成的系统级不规则性。

Grouped GEMM 让多个专家在一次调用中使用各自的实际 token 数，减少逐专家 kernel 启动和无效 padding。它不能消除极端负载倾斜：少数专家收到大量 token 时，矩阵形状、Tensor Core 利用率和尾部延迟仍可能由最忙的专家决定。评测必须记录专家负载分布，不能只看平均 token 数。

在 JAX 中，token 计数若被复制到主机，可能引入同步，破坏 CUDA Graph 捕获或触发重新编译。将元数据留在设备上能保持更稳定的执行图，但要求运行时正确管理 ragged buffer、边界和溢出。MaxText 中的 buffer factor 与周期性检查应理解为容量和正确性保护，而不是对任意路由偏斜的保证。

专家并行把 token 发往持有目标专家的设备，再把结果组合并恢复原顺序。融合 dispatch 和 combine 可以减少中间写入与等待；token 去重在同一目的地需要多个专家副本时可能减少网络传输。但收益取决于 Top-K、专家布局和跨节点比例。应同时测量消息数、有效 payload、元数据开销和通信计算重叠。

MXFP8 的收益伴随数值风险。专家组的激活分布可能不同，分组缩放需要检查溢出、下溢、累加精度以及梯度统计。短时间吞吐提升不能证明训练等价；至少要比较 loss 曲线、路由均衡性、梯度异常和最终任务指标，并固定优化器、数据顺序和随机种子。

Host offload 用主机带宽换取设备显存。若传输能与独立计算重叠，它可以帮助模型装入设备；若 PCIe 或其他互连成为关键路径，收益会反转。XLA 多 stream collective 也有类似边界：并发通信可能隐藏延迟，却可能增加带宽竞争、临时内存和依赖复杂度。

可复现实验应逐项拆分优化。先用小型 Dropless MoE 对比简单参考实现，再依次启用 grouped GEMM、通信融合、MXFP8、offload 和多 stream。记录编译后 step time、编译时间、显存与主机内存峰值、专家 token 直方图、各阶段耗时、网络字节和重叠比例；随后在均衡与恶意偏斜路由、不同 Top-K、序列长度、专家数和节点规模下重复，并完整保存硬件、软件、配置、profiling 与收敛信息。这样才能判断厂商结果能否迁移到实际集群。