计算最优并非集群最优：面向稀疏专家混合模型的系统感知扩展

# 计算最优并非集群最优：面向系统的稀疏混合专家扩展  
来源：https://arxiv.org/html/2608.10605  
\nameYuxin Tang\[email protected] \addrAmazon AGI Foundations  
\nameSheng Zha\[email protected] \addrAmazon AGI Foundations  

###### 摘要  
在大规模预训练中，算法、架构和系统决策通常是在互不连接的阶段中做出的。首先是一个扩展律阶段，它选择架构和训练方案，在计算约束下优化损失；然后是一个独立的系统阶段，针对硬件效率优化实现。在本工作中，我们开发了 MOSAIC，它将模型架构与系统协同设计表述为一个优化问题。MOSAIC 将一个预测性扩展律与一个经过校准的性能模型耦合起来，该性能模型可估计模型 FLOPs 利用率（MFU）、通信成本、内存占用以及最优并行布局。我们针对稀疏混合专家（MoE）语言模型实例化该框架，其中专家数量、路由稀疏性以及 MoE 层的其他维度会影响损失和系统效率。我们在文本数据上训练的稀疏 MoE 模型上拟合了一个扩展律，其扩展维度包括稀疏因子，即在前向传播中每个 token 不激活的模型参数比例。我们工作中的扩展律扫描范围涵盖从 104 亿激活参数到 27 亿激活参数（原文如此，疑为 1.04 亿到 27 亿——译者注），总模型规模达到 790 亿参数。我们表明，在校准的稀疏范围内，一个与效率无关的模型 FLOPs 预算不存在内部最优稀疏性。拟合的损失随模型越稀疏而单调递减，计算最优值位于数据支撑的上边界。MoE 模型中的最优稀疏性反而是在集群的系统约束下出现的，正如 MOSAIC 所捕获的那样。我们的结果主张转向面向前沿语言模型训练的统一架构与系统协同设计。  

关键词：LLM 预训练、扩展律、硬件感知的模型协同设计  

## 1 引言  
计算最优扩展律（kaplan2020scaling; hoffmann2022training）是将固定计算预算转化为预训练中模型规模和数据的核心工具。它们在实践中是标准做法，因为一旦训练计算预算确定，架构（模型规模 \(N\)）和数据（token 数量 \(D\)）的决策都遵循对预测损失曲线的单次优化，而且该曲线可以外推至远超其拟合区间，使其在前沿规模下也是可预测的工具（bahri2024explaining; hestness2017deep）。对于稠密模型，hoffmann2022training 设计了关于损失与模型规模 \(N\) 和训练 token 数 \(D\) 的联合扩展律，该扩展律允许闭式形式的模型 FLOPs¹ 最优分配。更新的规定性扩展律在额外数据约束下推荐最优扩展维度（lovelace2026prescriptive），但这些约束都没有计入所规定架构的系统效率。这一点很重要，因为两个预测质量相同的架构，在计入其可实现模型 FLOPs 利用率（MFU）² 和最优并行布局后，其训练吞吐量可能相差数倍（rajbhandari2022deepspeed）。随之而来的是，架构的系统侧建议在很大程度上绕开了扩展律，它们将模型规模视为给定，并在该规模下针对硬件效率优化架构。例如，anthony2024codesign 表明，在给定规模下，通过将 transformer 的内在维度（如隐藏大小、注意力头数和词表大小）对齐 GPU 计算核的结构，可以在匹配精度的同时提高吞吐量。另一条互补的工作线如 FlashAttention（dao2022flashattention）以算术等价的方式重新安排操作，通过 IO 感知、内存高效的调度来计算相同的注意力。因此，训练方案通常是在两个互不连接的阶段中构建的：一个用于模型和数据的扩展律阶段，以及一个硬件感知调优阶段，如 Kimi-K2 报告（team2025kimi）中分别处理稀疏 MoE 扩展律和系统调优。这种分离在大规模下影响重大，因为模型通过数据、张量、流水线和专家并行被切分到一组 GPU 上，如 MegaScale（jiang2024megascale）。由扩展律选择的架构应当考虑这些系统效率，而不仅仅是模型 FLOPs（bian2025inference）。  
本文的目标是将系统阶段并入扩展律阶段，使单一优化同时选择架构、token 预算和集群并行执行布局，并将该框架实例化用于稀疏 MoE 语言模型。我们关注此前用于语言模型训练的稀疏 MoE 架构（fedus2022switch; sarkar2024revisiting; abnar2025parameters），其中稀疏性通过将 transformer 块的稠密前馈层替换为一组 \(E\) 个专家，并将每个 token 路由到其中一部分来实现。这将决定每 token 模型 FLOPs 的激活参数 \(N_{\text{act}}\) 与总存储参数 \(N_{\text{tot}}\) 解耦，专家数量 \(E\) 是控制总参数数量的旋钮之一。现有的稀疏 MoE 模型扩展律（clark2022unified; tian2026towards; krajewski2024scaling）在忽略系统效率的情况下最优化这些架构旋钮，包括 abnar2025parameters，他们表明稀疏因子 \(S\)（大致是每个 token 不激活的 \(N_{\text{tot}}\) 的比例）引入了一个经典的仅含 \((N,D)\) 的扩展律无法看到的额外最优点。我们在第 6.1 节和图 1 中表明，仅凭模型 FLOPs 预算不足以确定有意义的 MoE 稀疏性。  
最近的一些研究开始在 MoE 扩展律本身中浮现这些系统侧压力。ludziejewski2025joint 拟合了一个稀疏 MoE 扩展律，发现扩展专家数量可以在固定每 token 模型 FLOPs 下改善损失，但会增加存储参数和显存压力，而模型 FLOPs 预算从不为此收费，因此实际最优的 \(E\) 是损失曲面与内存感知预算相交之处，而不是在固定 \(N_{\text{act}}\) 下最小化损失的 \(E\)。wan2026holistic 进一步推进，将总参数和每 token 模型 FLOPs 视为独立自由度，并认为仅每 token 模型 FLOPs 不是选择稀疏 MoE 架构的充分公平性指标。在这些研究的基础上，我们拟合了一个新的扩展律，其变量被选择用来暴露系统效率。  
传统上，在模型 FLOPs 预算下最优稀疏 MoE 几何形状的推导分两个阶段：一个针对目标参数、稀疏性和数据预算的与几何无关的凸规划，然后匹配到最近的离散架构几何形状。我们的扩展分析表明，在任何模型 FLOPs 预算下，更稀疏的 MoE 模型都能达到更低的损失，且最优稀疏性位于边界，这迫使约束力从模型 FLOPs 转为硬件可交付模型 FLOPs，即集群在训练窗口内实际能够交付的有效计算。我们求解了一个单步离散几何混合整数非线性规划，称为 MOSAIC。它直接在离散架构网格上最小化预测损失，解决了边界稀疏性问题，并实现了图 1 中动机所描述的集群特定稀疏 MoE 几何形状。  
我们将本工作聚焦于稀疏 MoE 架构-系统协同设计的训练侧，将硬件成本下的测试时扩展律（sadhukhan2026kinetics）留作未来工作。与我们研究最接近的是 sun2026hardware，他们将 roofline 分析与架构选择耦合起来。然而，仅凭 roofline 会遗漏那些翻转 MoE 推荐的额外开销，即专家分发与合并、细粒度专家下的核启动开销、作为 EP 拓扑函数的全对全带宽，以及数据并行和流水线并行耦合的暴露集体通信（gale2023megablocks; jiang2024megascale）。  
参见图注  
（a）模型 FLOPs 选择边界 MoE 稀疏性。在固定 \(C_{\mathrm{model}}=3.3\times10^{21}\) 下，拟合损失在校准稀疏范围内递减，因此 \(S^{\star}=S_{\max}\)（左）。我们不封顶任何模型维度，在直到 \(S_{\max}\)（此处为 0.985）的每个稀疏度下，专家数量和激活参数都在预算下被优化。在具有 4 个 AWS p6-B200 节点、为期 5 天的硬件预算 \(C_{\text{peak}}\) 下，该处方不可行，因为超出 \(S=0.96\) 可行边界的配置交付的 \(C_{\text{deliverable}}\) FLOPs（红色区域）少于固定处方 \(S^{\star}\) 所需的量（右）。  
参见图注  
（b）MOSAIC 在硬件预算下产生一个内部可行稀疏性。由 \(C_{\mathrm{deliverable}}\) 引出的硬件感知损失 \(\mathcal{L}^{\star}_{\mathrm{hw}}(S)\) 在内部稀疏性 \(S^{\star}_{\mathrm{cluster}}\) 处取得最小，且比 \(S^{\star}=S_{\max}\) 处低 0.031 nats。每个稀疏性配置在同一个 4 节点、5 天硬件包络下，按其自身可达到的峰值 MFU 所支持的 token 数量进行评估。  
图 1：MOSAIC 的系统感知扩展动机。  

##### 贡献。  
为了运行 MOSAIC 工作流并在可交付模型 FLOPs 约束下获得最优稀疏 MoE 扩展维度，我们讨论以下贡献和观察：  
- •我们拟合了一个新的四维联合 MoE 预训练扩展律，它暴露了系统侧旋钮。通过我们的扩展分析，我们表明在模型 FLOPs 预算范围内以及在校准稀疏范围内，拟合损失最优稀疏性位于数据支撑的上边界。这种边界寻求行为促使我们从模型 FLOPs 预算转向可交付模型 FLOPs 预算。  
- •我们开发了一个算子级分析性能模型，用于衡量 MoE transformer 的逐块前向和后向成本。我们针对高达 180 亿激活参数的训练运行网格验证了性能模型，预测 MFU 的平均绝对百分比误差每次扫描保持在 15% 以内。  
- •我们将这两个组件耦合在 MOSAIC 内部，并在固定 GPU 小时预算下，在 NVIDIA B200 GPU 集群上搜索架构网格。我们表明最低损失模型配置不是产生最多模型 FLOPs 的那个。我们在基础设施上对高达 2500 亿总参数的分阶段预训练运行验证了 MOSAIC 优化的排序选择。我们表明，损失排序从模型 FLOPs 轴翻转到峰值等效硬件计算轴，正如 MOSAIC 所预测的那样。  

我们首先在第 2 节描述 MOSAIC 的抽象公式化。为了应用 MOSAIC 进行稀疏 MoE 模型优化，我们首先在第 3 节提供我们运行中使用的训练方案的技术概述。然后我们描述用稀疏 MoE 架构实例化 MOSAIC 的组件：（1）第 4 节中的稀疏 MoE 扩展律，以及（2）第 5 节中的性能模型。我们在第 6 节具体化 MOSAIC 在稀疏 MoE 上的应用，其中扩展律和性能模型共同说明为什么 MoE 稀疏性在模型 FLOPs 预算下没有内部最优点，以及为什么稀疏 MoE 的正确预算是硬件可交付 FLOPs。最后，我们在第 7 节讨论实验和结果。我们在表 1 中提供了本文其余部分使用的共享符号快照。  

表 1：MOSAIC 共享符号快照。

| 符号 | 含义 |
|------|------|
| \(L_{\text{layers}}\) | transformer 块的数量。 |
| \(d\) | 残差流宽度（\(d_{\text{model}}\)）。 |
| \(d_{\text{ff}}\) | 稠密 FFN 的隐藏宽度，稀疏 MoE 块取代之。 |
| \(d_{\text{expert}}\) | 单个专家的隐藏宽度。 |
| \(E\) | 每个稀疏 MoE 块中的专家总数。 |
| \(K\) | 每个 token 的 Top-\(K\) 激活专家数。 |
| \(G\) | 专家拆分因子，\(G=d_{\text{ff}}/d_{\text{expert}}\)。图 14 在附录 B 中可视化该构造。 |
| \(N_{\text{act}}\) | 每个 token 的激活参数（FLOPs 相关）。 |
| \(N_{\text{tot}}\) | 模型总存储参数（模型容量）。 |
| \(S\) | 稀疏性，\(S=1-N_{\text{act}}/N_{\text{tot}}\)（每个 token 不激活的参数比例）。 |
| 路由与负载均衡 | |
| \(H\in\mathbb{R}^{T\times d}\) | 一批 \(T\) 个 token 在 MoE 之前的隐藏状态（算法 1）。 |
| \(\phi\in\mathbb{R}^{E}\) | 每专家路由偏置（LFLB 控制器）。 |
| \(\mu_{b}\) | LFLB 控制器的偏置更新率。 |
| \(\rho\) | Gumbel 路由的提议-选择比，给出提议大小 \(M=\min(\rho K,E)\)。 |
| \(s'_i(x)\), \(s_i(x)\) | 干净 sigmoid 路由器分数，以及加入均衡偏置的选择分数 \(s_i(x)=s'_i(x)+\phi_i\)（第 3.3 节）。 |
| \(\mathcal{T}_{K}(x)\) | 为 token \(x\) 选择的 Top-\(K\) 专家集合。 |
| MOSAIC 决策变量 | |
| \(Z\) | MOSAIC 的架构决策变量，一个可实现的稀疏 MoE 几何形状，带有诱发的扩展律维度 \((N_{\text{tot}},S,G)\)（第 2 节）。 |
| \(P_{\ell}\) | 执行计划，\(P_{\ell}=(P_{\text{TP}},P_{\text{EP}},P_{\text{CP}},P_{\text{PP}},P_{\text{DP}},B_{\text{micro}},A_{\text{ckpt}})\)，给出沿张量（TP）、专家（EP）、上下文（CP）、流水线（PP）和数据并行（DP）组的并行度，以及微批大小和激活检查点模式。 |
| \(B\) | 批处理包络，\(B=(B_{\text{global}},T_{\text{seq}})\)，给出每个优化器步的微批数 \(n_{\text{mb}}=B_{\text{global}}/(P_{\text{DP}}\,B_{\text{micro}})\)。 |
| \(D\) | 训练 token 数量。 |
| \(C_{\text{model}}\) | 模型 FLOPs，\(C_{\text{model}}=6\,N(Z)\,D\)，其中激活参数 \(N(Z)\) 是稀疏 MoE 的 \(N_{\text{act}}\)。 |
| \(T_{\text{train}}\) | 集群上的墙钟训练窗口。 |
| \(C_{\text{peak}}\) | 原始硬件上限，\(C_{\text{peak}}=F_{\text{peak}}\,N_{\text{dev}}\,T_{\text{train}}\)，其中 w