Training Communication-Efficient Mixture-of-Experts Language Models With Layer Re-Configuration
Training Communication-Efficient Mixture-of-Experts Language Models With Layer Re-Configuration
发表时间: 2026-08 · arXiv:2608.28511 (NVIDIA)
原文: https://arxiv.org/abs/2608.28511
Simeng Sun, Roger Waleffe {simengs,rwaleffe}@nvidia.com
速读
一句话结论 CE-MoE 打破"每个 token-mixing 层后面跟一个 MoE 层"的惯例，只保留少数几个更宽的路由 MoE 层，省下的深度换成更多 Mamba-2 与 dense FFN 层；在总参数与激活参数对齐的前提下，2B 到 31.5B 五个规模验证 loss 与基线持平而训练 GPU 小时降 30.5%–35.0%，31.5B 上省 33.3% 的同时下游均分与推理吞吐更高。
要解决什么问题 专家并行下每个 MoE 层要用 all-to-all 把 token 分发到持有对应专家的 rank、再收回输出，反向沿同样路由再走一遍梯度，这些集合通信会吃掉端到端训练时间中相当大的一块。一步训练的通信量满足 $V_{\text{step}} \propto L_E \cdot S \cdot \text{GBS} \cdot K\,(4 D b_x + 2 b_p)$，其中 $L_E$ 是 MoE 层数、$K$ 是路由 top-K、$D$ 是分发维度、$b_x, b_p$ 是每元素与每路由概率的字节数。改集合通信算法、做重叠、优化 kernel，减的都是暴露出来的开销，搬运量一点没少；公式给出四个真能减量的杠杆：通信精度、$D$、$K$ 和 $L_E$。已有做法在 $L_E$ 上只走到隔层放一个 MoE（Switch Transformer 那种），仍守着 1:1 的交替。
怎么做的 核心思路是把 token-mixing 与 channel-mixing 的层数解耦，允许偏向 token-mixing 的非均匀布局：约每四个 token-mixing 层才配一个 MoE 层。构造分三步。定层组成：总层数不变，固定 attention 数以控住长上下文开销，令 dense FFN 数等于 MoE 数，则目标比例 $\alpha$ 下 $b = \lfloor L/(\alpha+2) \rfloor$，剩下全给 Mamba；以 Nemotron-3 Nano 的 23 Mamba + 23 MoE + 6 attention、$\alpha=4$ 为例得到 8 MoE + 8 dense FFN + 30 Mamba，all-to-all 调用数从基线的 $23q$ 降到 $8q$，比 Switch 式的 $12q$ 再少一截。定层序：以 Mamba-2 为脊，把块打包成虚拟块后分阶段贪心排布，让 attention 与 channel-mixing 块尽量均匀地摊在脊上。定模型形状：固定 hidden size，在专家数、top-K 与各处 FFN 宽度上网格搜索，只收总参数与激活参数都与基线对齐的候选，再挑通信量代理 $L_E K$ 最小的那个；31.5B 例子里选中的配置把 $L_E K$ 降了 59.4%。这套布局有个前提：连续堆 token-mixing 会不会塌。370M 玩具模型堆七个连续 token-mixer，自注意力变体出现 loss 尖峰并发散，Mamba-2 平稳收敛；用稳定秩 $r_{\text{stable}}(X) = \lVert X \rVert_F^2 / \sigma_1^2$ 量残差流，发散的自注意力早期层稳定秩接近 1、每 token 的 $\ell_2$ 范数大出几个数量级——token 差异靠幅度而非方向承载，这正是不稳定的来源。主 token-mixing 路径因此交给线性循环块。
效果如何 实验全用 hybrid Mamba-2 模型，在 H100 集群上以 BF16、序列长度 8192、全局批 768 训练，规模从 2B-A0.5B 到 31.5B-A3.5B。成本用端到端 GPU 小时而非理论 FLOPs，质量看验证 loss 加 11 个下游基准。对比基线三类：同规模常规 hybrid-MoE（标准交替布局）、LatentMoE（压低分发维度那条路线）、21 个受约束的随机层序。同批量下五个规模 GPU 小时降 30.5%–35.0%，验证 loss 只差千分之几；层序消融中分阶段贪心比随机层序均值低约 1.8 个标准差，说明起作用的不只是层组成，还有层序。14B-A2B 同 token 数省 34.1%，下游均分 56.93 → 58.41；31.5B-A3.5B 省 33.3%，58.37 → 58.65。省下的算力也可换成更多数据：多训 1.25 倍 token 后 14B 均分升到 59.36、仍省 26.1%，1T 长程设置省 21.0% 且均分 66.26 → 66.67。推理侧在单节点 H100 上跑静态合成基准，吞吐比基线高 28%–36%。作者自列的代价：同批量下每个专家每步反而少收 22% 路由 token；序列级负载不均衡信号高于全 MoE 基线，全局批不随专家数放大时早期 loss 尖峰更多；节省比例在 FP8/NVFP4、更大 world size 或开通信重叠后可能变化。
引言与核心贡献
在训练采用专家并行（Expert Parallelism, EP）的混合专家（Mixture-of-Experts, MoE）语言模型时，由于需要将Token路由到对应的专家，前向和后向传播中不可避免地需要跨节点通信（而不是仅限本地计算）。具体的，前向传播中每个MoE层使用 all-to-all 集合通信将Token激活分发到托管所选专家的节点，并将专家输出组合回原始Token节点；后向传播沿相同的路由模式进行梯度通信。这些集合通信会产生巨大的通信量，并消耗端到端训练时间的很大一部分【28，Scalable training of mixture-of-experts models with Megatron Core + 2026 + arXiv + https://arxiv.org/abs/2603.07685】。在无丢弃路由且无容量填充的情况下，一个训练步骤（前向+后向）的专家并行通信 量 $V_{\mathrm{step}}$ 的缩放公式如下：
$$ V_{\mathrm{step}} \propto L_E \cdot S \cdot \mathrm{GBS} \cdot K \left( 4 D b_x + 2 b_p \right) , $$其中 $L_E$ 是MoE层数，$S$ 是序列长度，$\mathrm{GBS}$ 是全局批大小，$K$ 是路由的 top-$K$，$D$ 是分发隐藏层维度，$b_x$ 和 $b_p$ 分别是每个隐藏层元素（例如BF16）和每个路由概率（FP32）的字节数。
系统级的优化（如集合通信算法、计算与通信重叠等）可以降低暴露的通信成本，而本文的研究目标是与其正交且互补的：从架构层面直接减少移动数据的总通信体积。基于上述公式，本文提出了一种通信高效的MoE模型（CE-MoE）。
核心创新点：
1. 异构层布局设计：打破了传统模型在每个Token混合层（如注意力机制、Mamba-2）后交替堆叠MoE层的常规做法，解耦了Token混合（token-mixing）和通道混合（channel-mixing）的深度。
2. 重分配专家容量：大幅减少路由MoE层的数量（降低 $L_E$），将专家容量集中在少数几个路由MoE层中，并通过增加额外的Token混合层和密集前馈网络（Dense-FFN）层来维持模型的整体深度。
3. 规模化验证：在总参数量从2B到31.5B的缩放阶梯上，在总参数和激活参数匹配的情况下，CE-MoE模型持续降低了训练成本，同时在验证损失和下游基准测试上与全MoE基线模型表现持平。在31.5B规模下，CE-MoE减少了 $33.3\%$ 的GPU小时消耗，同时提高了下游平均得分和推理吞吐量。
背景知识与设计原则
相关工作涵盖多个层级的MoE训练优化。在内核级别，SonicMoE【12，SonicMoE: Accelerating MoE with IO and tile-aware optimizations + 2025 + arXiv + https://arxiv.org/abs/2512.14080】通过感知分块的优化提高了MoE效率。在系统和框架级别，DeepEP【8 ，DeepEP: An efficient expert-parallel communication library + 2025 + GitHub + https://github.com/deepseek-ai/DeepEP】提供了高吞吐低延迟的all-to-all内核，Megatron-Core【28 ，Scalable training of mixture-of-experts models with Megatron Core + 2026 + arXiv + https://arxiv.org/abs/2603.07685】 和 MegaScale-MoE【15，MegaScale-MoE: Large-Scale Communication-Efficient Training of Mixture-of-Experts Models in Production + 2026 + EuroSys + https://dl.acm.org/doi/10.1145/3767295.3769325】提供了可扩展的系统栈。在架构-系统边界 ，Shortcut-connected MoE【2，Shortcut-connected expert parallelism for accelerating mixture-of-experts + 2024 + arXiv + https://arxiv.org/abs/2404.05019】改变了依赖结构以实现通信与计算的重叠。此外，LatentMoE【10 ，LatentMoE: Toward optimal accuracy per FLOP and parameter in mixture of experts + 2026 + arXiv + https://arxiv.org/abs/2601.18089】 和Multi-Head LatentMoE【6，Multi-head LatentMoE and head parallel: Communication-efficient and deterministic MoE parallelism + 2026 + arXiv + https://arxiv.org/abs/2602.04870】通过降低有效分发维 度 $D$ 来减少流量。
现有大模型通过交替排列减少MoE深度。现有的架构通常通过交替堆叠MoE和密集FFN块来减少路由MoE的深度【17，GShard: Scaling giant models with conditional computation and automatic sharding + 2020 + arXiv + https://arxiv.org/abs/2006.16668】【11 ，Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity + 2022 + JMLR + https://arxiv.org/abs/2101.03961】【30 ，ST-MoE: Designing stable and transferable sparse expert models + 2022 + arXiv + https://arxiv.org/abs/2202.08906】【26 ，MAI-Thinking-1: Building a Hill-Climbing Machine + 2026 + Technical report + https://microsoft.ai/pdf/mai-thinking-1.pdf】，例 如Switch Transformer在每隔一个FFN层放置MoE。然而，这些设计仍然保留了标准Transformer中Token混合和通道混合块之间 $1:1$ 的规则交替（例如 $\mathtt{attn} \to \mathtt{dense\_mlp} \to \mathtt{attn} \to \mathtt{moe}$）。相比之下，本文更激进地减少了MoE深度，解耦了Token混合与通道混合的数量，产生了偏离标准Transformer块设计的异构层布局。
CE-MoE的异构性源于分配比例的倾斜。CE-MoE分配了更多层用于沿序列维度交换信息，为每个MoE层分配四个Token混合层，或者为每个通道混合层分配超过两个Token混合层。基于这种侧重Token混合的层组成，本文采用了一种简单的分阶段贪婪过程来构建CE-MoE层模式，将注意力和通道混合块分布在由线性循环块构成的骨干上，同时允许连续的Token混合块存在。
连续Token混合层引发的秩崩溃风险与评估。CE-MoE的布局引入了相邻的Token混合块，这引发了对训练稳定性的担忧，因为现有研究表明，没有MLP或跳跃连接的纯自注意力堆叠会以双指数速度收敛到秩1【9，Attention is not all you need: Pure attention loses rank doubly exponentially with depth + 2021 + ICML + https://proceedings.mlr.press/v139/dong21a.html】，从而丧失表达能力。本文通过实证表明，即使有跳跃连接，自注意力堆叠在训练中也是不稳定的。在包含7个连续Token混合层（自注意力或Mamba-2【7 ，Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality + 2024 + arXiv + https://arxiv.org/abs/2405.21060】）的 约 $370\mathrm{M}$ 参数微型模型实验中，注意力变体出现了损失尖峰并最终发散，而Mamba-2变体则平稳收敛。
量化Token多样性与秩变化的指标定义。为了量化分析，本文测量了残差后表示 $X \in \mathbb{R}^{N_{\mathrm{tok}} \times D}$ 的逐层统计量，其中 $N_{\mathrm{tok}}$ 是采样的评估Token数，隐藏维度 $D = 1024$。设 $\sigma_i$ 为 $X$ 的奇异值，$\| \cdot \|_F$ 为Frobenius范数，$p_i = \sigma_i^2 / \sum_j \sigma_j^2$ 为归一化谱分布，$H(p)$ 为其香农熵。本文报告了稳定秩 $r_{\mathrm{stable}}(X) = \|X\|_F^2 / \sigma_1^2$ 和有效秩 $r_{\mathrm{effective}}(X) = \exp(H(p))$。当任一指标接近1时，表明残差流的大部分能量集中在单一的奇异方向上，Token表示变得几乎共线。
自注意力网络存在严重的低秩瓶颈。实验结果表明，自注意力网络具有比Mamba-2低得多的稳定秩和有效秩，且在早期层崩溃最为严重。在较高学习率下，早期自注意力层的稳定秩接近1，有效秩降至个位数；此外，其每个Token的 $\ell_2$ 范数比Mamba-2大几个数量级。这表明Token的差异越来越多地由标量幅度而非多样化的方向来承载，这成为不稳定的主要来源。
线性循环块是异构布局的更优选择。上述结果表明，当Token混合器在重复堆叠下能保持健康的秩时，连续的Token混合块是可行的。在微型实验中，线性循环块（Mamba-2）表现出比自注意力更好的稳定性，因此是异构布局中主要Token混合路径的更好候选者。虽然门控注意力【24，Gated attention for large language models: Non-linearity, sparsity, and attention-sink-free + 2025 + arXiv + https://arxiv.org/abs/2505.06708】和归一化GPT【20 ，nGPT: Normalized Transformer with representation learning on the hypersphere + 2024 + arXiv + https://arxiv.org/abs/2410.01131】等机制可能有助于稳定纯自注意力网络，但这超出了本文范围。后续部分聚焦于混合模型（Softmax注意力与线性注意力），其中线性循环块提供了构建侧重Token混合的CE-MoE层模式所需的架构灵活性 。
核心方法设计
构建通信高效的混合MoE（CE-MoE）架构旨在在保持基线混合MoE模型下游质量的同时，大幅减少路由MoE层的使用。该过程分为三个步骤：确定层组成、生成层模式、调整模型形状。
确定层组成
打破1:1的层比例限制。标准Transformer架构将注意力与密集FFN（或MoE）的组合视为单层，隐含地强制执行 $1:1$ 的Token混合与通道混合比例。CE-MoE通过减少MoE层，并将深度重新分配给额外的Token混合层以及适量的密集FFN块，放宽了这一比例。
层分配的数学表达。假设基线的Mamba、MoE和注意力块的组成分别为 $(a^\star, b^\star, c^\star)$，总层数 $L = a^\star + b^\star + c^\star$。CE-MoE使用 $(a, b, c, d)$ 表示Mamba、MoE、注意力和密集FFN块，同时保持总层数 $L = a + b + c + d$ 不变。CE-MoE通过设定目标Token混合与MoE的近似比例 $\alpha$ 来设置块数量，并保持 $c = c^\star$ 以控制长上下文注意力成本。在设定 $d = b$ 的情况下，计算公式为 $b = \lfloor \frac{L}{\alpha + 2} \rfloor$，剩余的块全部分配给Mamba。
实际配置案例分析。以Nemotron-3 Nano基线组成为例【21，NVIDIA Nemotron 3: Efficient and Open Intelligence + 2025 + arXiv + https://arxiv.org/abs/2512.20856】，其配置 为 $(a^\star, b^\star, c^\star) = (23, 23, 6)$（总 $L = 52$）。若目标比例 $\alpha = 4$，则 $b = d = \lfloor 52 / (4 + 2) \rfloor = 8$，$a = 30$。假设 $q$ 表示一个路由MoE层中每个专家并行组内每个Rank在一次微批次（前向+后向）中调用all-to-all的次数，基线交替布局、隔层（Switch风格）变体和CE-MoE分别产生 $23q$、$12q$ 和 $8q$ 次MoE all-to-all调用。相对于基线，Switch风格减少了 $47.8\%$ 的通信，而CE-MoE减少了 $65.2\%$。在考虑路由 top-$K$ 增加 $1.5\times$ 的情况下，CE-MoE的通信体积也能降低约 $47.8\%$。
生成层模式
采用分阶段贪婪策略排序。层组成确定了块的多重集，但未确定顺序。CE-MoE使用算法1中的分阶段贪婪过程分配层顺序，其中M、E、$\star$ 和 - 分别表示Mamba、MoE、注意力和密集FFN块。该算法将自注意力和通道混合块分布在Mamba-2骨干上，同时避免不良的局部排序。
算法执行的四个阶段。
首先，算法运行在四个连续阶段中：
1. 计算常规虚拟块的名义Mamba预算 $\ell = \lceil n_M / n_{\mathrm{vb}} \rceil$。
2. 阶段1和阶段2贪婪地创建包含 $m = \ell$ 的注意力承载块（集合 $A$）和常规非注意力块（集合 $N$）。
3. 阶段3将任何剩余的层存储在残余块（集合 $R$）中，其Mamba数量可能较小。
4. 阶段4执行 Distribute(A, N)，将每个虚拟块视为原子，均匀地将 $N$ 块分配以划分 $A$ 块序列，从左到右打破平局。最后将 $R$ 中的块追加到末尾。
在每个虚拟块内部，应用相同的间距原则：通道混合块分布在Mamba骨干上，注意力（如果存在）位于Token混合的尾部，并且该块以通道混合结束，以便下一个虚拟块可以以Token混合开始。
# Algorithm 1 Generate CE-MoE layer pattern
Require: Composition (a, b, c, d) over (Mamba, MoE, attention, dense FFN), b >= c
Ensure: Layer pattern over M, E, *, and -
1: (n_M, n_E, n_A, n_D) <- (a, b, c, d); n_vb <- max(n_E, n_A, n_D, 1); l <- ceil(n_M / n_vb)
2: # A virtual-block tuple (m, 1_e, 1_a, 1_d) uses indicator variables; 1[·] is indicator function; || denotes concatenation
3: A, N, R <- [], [], []
4: while n_A > 0 and n_M >= l do # Phase 1: emit attention-bearing virtual blocks
5: (1_e, 1_d) <- (1[n_E > 0], 1[n_D > 0])
6: A.append((l, 1_e, 1, 1_d))
7: (n_M, n_E, n_A, n_D) <- (n_M - l, n_E - 1_e, n_A - 1, n_D - 1_d)
8: end while
9: while n_M >= l and n_E > 0 and n_D > 0 do # Phase 2: emit regular non-attention virtual blocks
10: N.append((l, 1, 0, 1))
11: (n_M, n_E, n_D) <- (n_M - l, n_E - 1, n_D - 1)
12: end while
13: if n_M + n_E + n_A + n_D > 0 then # Phase 3: collect residual counts into shortened tail block(s)
14: append residue virtual blocks with m > 0 covering the remaining counts to R
15: end if
16: blocks <- Distribute(A, N) || R # Phase 4: place N to evenly partition A, breaking ties left-to-right
17: P <- empty string
18: for all virtual blocks (m, 1_e, 1_a, 1_d) in blocks do
19: s <- Insert(M^m, E, 1_e) || *^{1_a} || -^{1_d} # Insert returns M^m if 1_e=0, else adds E after floor(m/2) Mamba
20: P <- P || s
21: end for
22: return P运行示例与消融实验验证。对于组成 $(a, b, c, d) = (30, 8, 6, 8)$ 的运行示例，过程产生 $n_{\mathrm{vb}} = 8$ 和 $\ell = 4$。贪婪分配后，$|A| = 6$，$|N| = 1$，$|R| = 1$；分配操作随后产生全局虚拟块顺序 AAANAAAR。为了验证该过程的有效性，在2B参数和140B Token下，本文将分阶段贪婪的CE-MoE模式与21个共享相同块组成的受限随机层模式进行了比较。受限随机模式排除了明显糟糕的布局（如连续的通道混合块）。结果显示，分阶段贪婪的CE-MoE在验证损失上匹配基线，并优于所有受限随机替代方案（比均值低约1.8个标准差），这表明层排序本身携带着有用的归纳偏置。
确定模型形状
参数预算匹配与网格搜索。在固定层模式后，本文搜索模型形状配置，使得CE-MoE和基线在匹配的总参数和激活参数预算下进行比较。大部分宽度和路由比例保持在基线值附近，而每个专家的MoE-FFN宽度和专家数量被允许增加，以补偿较少的路由MoE层。在固定隐藏大小的情况下，本文网格搜索专家数量、路由 top-$K$、每个专家的MoE-FFN宽度、共享专家宽度和密集FFN宽度。在预算匹配的候选者中，选择具有最小 $L_E K$ 的配置（当 $D$ 和通信精度固定时，这是公式1中的通信成本代理）。
31.5B规模的形状搜索实例。在31.5B示例中，基线使用隐藏大小 $H = 2688$，每专家MoE-FFN宽度1856，共享专家宽度3712，专家数 $|\mathcal{E}| = 128$，$K = 6$，$L_E = 23$，得出 $L_E K = 138$。搜索空间跨越了1,708,200个候选者，其中5,817个满足接受标准。最终选择的CE-MoE形状使用密集FFN宽度8096，每专家MoE-FFN宽度3520，共享专家宽度3520，$|\mathcal{E}| = 192$，$K = 7$。相对于基线，总参数和激活参数仅分别变化了 $-0.001\%$ 和 $-0.257\%$，而 $L_E K$ 降低了 $59.4\%$。
表 1: 31.5B 运行示例的宽度和路由搜索空间。所有范围均相对于基线形状。
| 参数 | 搜索范围 |
|---|---|
| 隐藏大小 H | 1.0× (固定) |
| 密集FFN宽度 | 3.0-5.0× |
| MoE-FFN宽度 (每专家) | 1.5-4.0× |
| 共享专家宽度 | 0.95-1.15× |
| 专家数量 |E| | {1.25, 1.5} ×; 可被EP rank整除 |
| 路由密度 K/|E| | 0.5-1.5× |
| 接受标准 | 总参数 ±1%; 激活参数 ±1.5% |
表 2: 模型配置。最后一列报告了 LatentMoE 专家的输入/输出维度；– 表示无潜在投影。
| 规模 | Tokens | 模型 | 层数 | H | |E|/K | f_expert | f_dense | d_latent |
|---|---|---|---|---|---|---|---|---|
| 2B-A0.5B | 140B | Baseline | 12M+3A+0D+12E | 1024 | 128/6 | 640 | - | - |
| CE-MoE | 16M+3A+4D+4E | 1024 | 160/7 | 1520 | 1920 | - | ||
| 4B-A0.77B | 215B | Baseline | 15M+3A+0D+15E | 1280 | 128/6 | 768 | - | - |
| CE-MoE | 20M+3A+5D+5E | 1280 | 192/7 | 1504 | 3168 | - | ||
| 7B-A1B | 310B | Baseline | 16M+4A+0D+16E | 1536 | 128/6 | 1024 | - | - |
| CE-MoE | 20M+4A+6D+6E | 1536 | 192/7 | 1792 | 4576 | - | ||
| 14B-A2B | 560B | Baseline | 20M+4A+0D+20E | 2048 | 128/6 | 1280 | - | - |
| CE-MoE | 26M+4A+7D+7E | 2048 | 192/7 | 2432 | 3840 | - | ||
| 31.5B-A3.5B | 300B; 1T | Baseline | 23M+6A+0D+23E | 2688 | 128/6 | 1856 | - | - |
| CE-MoE | 30M+6A+8D+8E | 2688 | 192/7 | 3520 | 8096 | - | ||
| LatentMoE | 23M+6A+0D+23E | 2688 | 512/22 | 1856 | - | 672 | ||
| CE-LatentMoE | 30M+6A+8D+8E | 2688 | 640/24 | 3712 | 6336 | 768 |
实验配置
评估指标体系。本文从三个维度评估CE-MoE：(1) 训练成本：使用端到端训练成本（GPU小时）作为主要指标，通过剔除预热迭代和延迟异常值来估计稳态GPU小时数。(2) 模型质量：报告验证损失和11个标准化语言模型基准测试的下游性能（涵盖MMLU、编程、数学和常识推理等）。(3) 推理速度：在输入序列长度2048和8192下，报告不同输出序列长度的每秒Token数。
实验环境与模型参数。所有实验均使用混合Mamba-2模型，层组成表示为 $\mathtt{M}+\mathtt{A}+\mathtt{D}+\mathtt{E}$（分别对应Mamba-2、自注意力、密集FFN和MoE层）。在每个模型规模内，基线和CE-MoE变体使用相同的隐藏大小，并在总参数和激活参数预算上匹配。硬件配置为配备第4代NVLink的H100 HBM3集群。软件与训练配置采用BF16精度，序列长度8192，默认全局批大小（GBS）为768。优化器采用Adam，参数为 $(\beta_1, \beta_2) = (0.9, 0.95)$，权重衰减0.1，梯度裁剪1.0，并使用WSD学习率调度。MoE层使用all-to-all Token分发、分组GEMM、MoE排列融合、带有专家偏置的Sigmoid路由器分数，以及系数为1e-4的序列级辅助负载均衡。每个MoE层还包括一个无门控的共享专家MLP。
实验结果分析
同批次缩放阶梯实验验证
CE-MoE在严格控制变量下持续降低训练成本。本文首先在相同全局批大小（GBS=768）下比较基线和CE-MoE，跨越了从2B-A0.5B到31.5B-A3.5B的五个规模。图5显示，CE-MoE在五个规模上将稳态GPU小时数减少了 $30.5\%$ 到 $35.0\%$，同时最终验证损失与基线非常接近。这种相同GBS的设置对CE-MoE而言是保守的，因为其专家数量增加到基线的 $1.5\times$，而路由 top-$K$ 仅增加到 $1.1\times$；在均匀路由下，每个CE-MoE专家每步接收的路由Token减少了 $22\%$。在这一机制下CE-MoE仍能匹配损失，证明异构层模式本身提供了有用的建模能力。
模型质量与扩大批次实验
节省的通信成本可转化为更多训练数据或更高性能。本文探究了是否可以将节省的训练成本再投资到更大的全局批次中（将14B和31.5B实验的GBS增加到960）。表3显示，在14B-A2B规模，相同Token的CE-MoE运行比基线少用 $34.1\%$ 的GPU小时，并将下游平均得分从56.93提高到58.41；而更大GBS的CE-MoE运行消耗了 $1.25\times$ 的Token，平均得分进一步提高到59.36，同时仍比基线节省 $26.1\%$ 的GPU小时。在31.5B规模，相同Token的CE-MoE减少了 $33.3\%$ 的GPU小时，平均得分从58.37提升至58.65。在LatentMoE对比中，CE-LatentMoE消耗 $1.25\times$ 的Token，平均分从60.77升至62.36，成本降低 $20.9\%$。在1T Token的长周期设置中，CE-MoE在减少 $21.0\%$ 成本的情况下训练了 $1.25\times$ 的Token，平均得分从66.26升至66.67。
表 3: 训练成本和模型下游质量。报告了14B-A2B和31.5B-A3.5B的匹配Token比较，以及选定的在1.25× Token上训练的GBS=960的CE-MoE/CE-LatentMoE变体。
| 模型 | 14B-A2B | 14B-A2B | 14B-A2B | 31.5B-A3.5B | 31.5B-A3.5B | 31.5B-A3.5B | 31.5B-A3.5B | 31.5B-A3.5B | 31.5B-A3.5B |
|---|---|---|---|---|---|---|---|---|---|
| 消耗Tokens | Baseline 560B | CE-MoE 560B | CE-MoE 700B | Baseline 300B | CE-MoE 300B | LatMoE 300B | CE-LatMoE 375B | Baseline 1T | CE-MoE 1.25T |
| 训练成本 | |||||||||
| 归一化GPU-h | 1.00 | 0.66 | 0.74 | 1.00 | 0.67 | 1.00 | 0.79 | 1.00 | 0.79 |
| GPU-h 减少 | 34.1% | 26.1% | 33.3% | - | 20.9% | 21.0% | |||
| 下游质量 | |||||||||
| MMLU-Pro | 29.26 | 30.76 | 33.08 | 30.17 | 31.92 | 35.25 | 37.41 | 42.06 | 44.56 |
| MMLU 5-shot | 59.53 | 60.38 | 60.32 | 58.83 | 58.05 | 61.27 | 62.90 | 65.97 | 67.19 |
| HumanEval | 41.46 | 39.21 | 42.16 | 42.44 | 42.29 | 43.54 | 45.85 | 47.59 | 49.15 |
| MBPP | 53.33 | 53.11 | 55.39 | 55.41 | 54.14 | 57.39 | 58.23 | 61.03 | 62.04 |
| GSM8K | 64.44 | 69.75 | 73.24 | 67.63 | 69.83 | 72.71 | 71.87 | 81.20 | 78.47 |
| MATH-500 | 55.00 | 60.80 | 60.00 | 60.40 | 59.40 | 59.20 | 66.00 | 66.60 | 70.80 |
| MATH-Hard | 31.32 | 35.64 | 34.76 | 35.22 | 36.10 | 36.96 | 42.50 | 47.25 | 48.24 |
| RACE | 74.16 | 73.88 | 73.88 | 73.78 | 74.74 | 76.08 | 76.08 | 80.96 | 78.76 |
| ARC-Challenge | 72.70 | 74.57 | 74.49 | 72.53 | 73.04 | 77.99 | 77.39 | 83.11 | 84.30 |
| HellaSwag | 76.11 | 75.90 | 76.52 | 76.94 | 77.01 | 77.73 | 78.09 | 79.46 | 79.56 |
| WinoGrande | 68.90 | 68.51 | 69.14 | 68.67 | 68.67 | 70.40 | 69.69 | 73.64 | 70.32 |
| 平均分 | 56.93 | 58.41 | 59.36 | 58.37 | 58.65 | 60.77 | 62.36 | 66.26 | 66.67 |
推理速度评估
路由MoE层数减少显著提升了推理吞吐量。在单节点H100上进行的静态合成推理基准测试显示（批大小为4），对于31.5B模型，在输入序列长度2048和8192、输出序列长度从128到2048的范围内，CE-MoE的吞吐量比基线提高了 $28\% - 36\%$。尽管CE-MoE使用了更宽的专家FFN，但其路由通过的MoE层数更少，这降低了每个生成Token的路由器开销和同步点数量。
残差流分析
CE-MoE通过不同的轨迹达到可比的表达能力。本文比较了2B-A0.5B基线和CE-MoE模型在相同训练步骤下的逐层残差后表示。图7显示，基线与最终层的余弦相似度曲线相对平滑地增加，而CE-MoE遵循阶梯状轨迹：较少的专家层产生较大的表示更新，而它们之间的Token混合层提供较小的细化。有效秩曲线表明，尽管CE-MoE具有较大的层间秩变化，但两个模型最终具有相似的有效秩；秩在Token混合主导的跨度内下降，但在专家层附近恢复，表明这些跨度不会导致不可逆的秩崩溃。
结论与展望
本文提出了通信高效的混合MoE（CE-MoE）架构，通过减少MoE层数 $L_E$ 并将容量重新分配到额外的Token混合层以及更少但更宽的专家和密集FFN块中，在保持参数预算匹配的前提下降低了专家并行通信。缩放实验表明，CE-MoE在保持验证损失的同时，持续降低了GPU小时训练成本；在更大规模下，它能以更低成本训练更多Token，在下游基准上保持竞争力，并表现出优于全MoE基线的初步推理吞吐量增益。
未来的工作需要克服目前的一些限制：降低 $L_E$ 带来的GPU小时节省可能会在其他计算或通信精度（如FP8或NVFP4）、更大的集群规模或重叠通信策略下发生变化。此外，CE-MoE表现出比全MoE基线更高的序列级负载不平衡信号，如果全局批大小不随专家数量相应扩展，更容易遭受早期训练损失尖峰的影响；更好的路由正则化或负载均衡可能会改善效率-质量的权衡。
参考文献引用汇总
- 【2，Shortcut-connected expert parallelism for accelerating mixture-of-experts + 2024 + arXiv + https://arxiv.org/abs/2404.05019】（在背景知识中引用，说明架构-系统边界的通信优化工作 ）
- 【6，Multi-head LatentMoE and head parallel: Communication-efficient and deterministic MoE parallelism + 2026 + arXiv + https://arxiv.org/abs/2602.04870】（在背景知识中引用，说明通过改变架构和并行方式减少通信随激活专家数线性扩展的工作 ）
- 【7，Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality + 2024 + arXiv + https://arxiv.org/abs/2405.21060】（在秩崩溃风险章节引用，指代Mamba-2模型 ）
- 【8，DeepEP: An efficient expert-parallel communication library + 2025 + GitHub + https://github.com/deepseek-ai/DeepEP】（在背景知识中引用，说明提供高吞吐低延迟all-to-all内核的系统框架 ）
- 【9，Attention is not all you need: Pure attention loses rank doubly exponentially with depth + 2021 + ICML + https://proceedings.mlr.press/v139/dong21a.html】（在秩崩溃风险章节引用，证明纯自注意力堆叠会双指数收敛至秩1 ）
- 【10，LatentMoE: Toward optimal accuracy per FLOP and parameter in mixture of experts + 2026 + arXiv + https://arxiv.org/abs/2601.18089】（在背景知识中引用，说明通过降低有效分发维度减少流量的方法 ）
- 【11，Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity + 2022 + JMLR + https://arxiv.org/abs/2101.03961】（在背景知识中引用，作为交替堆叠MoE和密集FFN块的现有架构代表 ）
- 【12，SonicMoE: Accelerating MoE with IO and tile-aware optimizations + 2025 + arXiv + https://arxiv.org/abs/2512.14080】（在背景知识中引用，说明内核级别的MoE效率优化 ）
- 【15，MegaScale-MoE: Large-Scale Communication-Efficient Training of Mixture-of-Experts Models in Production + 2026 + EuroSys + https://dl.acm.org/doi/10.1145/3767295.3769325】（在背景知识中引用，作为提供可扩展系统栈的研究 ）
- 【17，GShard: Scaling giant models with conditional computation and automatic sharding + 2020 + arXiv + https://arxiv.org/abs/2006.16668】（在背景知识中引用，作为交替堆叠MoE和密集FFN块的现有架构代表 ）
- 【20，nGPT: Normalized Transformer with representation learning on the hypersphere + 2024 + arXiv + https://arxiv.org/abs/2410.01131】（在秩崩溃风险章节引用，作为可能稳定纯自注意力网络的机制之一 ）
- 【21，NVIDIA Nemotron 3: Efficient and Open Intelligence + 2025 + arXiv + https://arxiv.org/abs/2512.20856】（在核心方法设计章节引用，提 供Nemotron-3 Nano作为基线组成的具体例子）
- 【24，Gated attention for large language models: Non-linearity, sparsity, and attention-sink-free + 2025 + arXiv + https://arxiv.org/abs/2505.06708】（在秩崩溃风险章节引用，作为可能稳定纯自注意力网络的门控注意力机制 ）
- 【26，MAI-Thinking-1: Building a Hill-Climbing Machine + 2026 + Technical report + https://microsoft.ai/pdf/mai-thinking-1.pdf】（在背景知识中引用，作为交替堆叠MoE和密集FFN块的现有架构代表 ）
- 【28，Scalable training of mixture-of-experts models with Megatron Core + 2026 + arXiv + https://arxiv.org/abs/2603.07685】（在引言和背景知识中多次引用，说明EP集合通信消耗大量训练时间，并作为系统栈的代表 ）
- 【30，ST-MoE: Designing stable and transferable sparse expert models + 2022 + arXiv + https://arxiv.org/abs/2202.08906】（在背景知识中引用，作为交替堆叠MoE和密集FFN块的现有架构代表 ）
💬 评论讨论
欢迎在这里分享您的想法和见解！