Compute-Optimal Is Not Cluster-Optimal: Systems-Aware Scaling for Sparse Mixture-of-Experts
Soumajyoti Sarkar, Yuxin Tang, Sheng Zha
cs.LG, cs.AI
2026-08-11
亚马逊 MOSAIC 把 MoE 架构和系统布局一起优化,证明算力最优会把稀疏度顶到硬件边界,集群约束下的真最优损失再低 0.031 nats。
预训练一个稀疏 MoE(Mixture-of-Experts)大模型,常规做法分两步:先用 scaling law 选架构(专家数、稀疏度),在算力预算下把 loss 压到最低;再让系统团队为这套架构挑并行布局、塞进硬件。两步互相不说话。
问题在于,纸面算力不等于硬件真交付的算力。论文给出一个尖锐的观察:两个预测 loss 完全相同的架构,一旦算上各自能跑到的 MFU(Model FLOPs Utilization)和最优并行布局,实际训练吞吐能差好几倍。对稀疏 MoE 这点尤其致命,scaling law 单独看会说越稀越好,loss 随稀疏度单调下降,于是「算力最优」永远顶在稀疏度的天花板,而那个天花板在真实集群上往往根本跑不动。
MOSAIC(Model Optimization via Systems-Aware TraIning Co-design,亚马逊 AGI 团队)把架构和系统放进同一个优化。它一次性选定模型几何(专家数 E、专家切分因子 G、稀疏度 S)、训练 token 数 D,以及完整的并行布局 P(张量/专家/上下文/流水线/数据并行,外加 micro-batch 和激活 checkpoint),约束是一个固定的集群规模和训练窗口。
关键的偏离在于:算力预算随架构变化,等于 Cdeliverable = Cpeak · MFU(架构, 布局) · goodput,而不是一个常数。于是约束从「Cmodel ≤ 某个固定 FLOPs 数」变成「Cmodel ≤ Cdeliverable」,前者里的算力是 6·Nactive·D 的纸面值,后者是硬件真正吐得出来的算力。
喂给这个优化器的有两块。
第一块是扩展过的 scaling law L(Ntot, S, D, G),新加了专家切分因子 G 这个轴(G = dff / dexpert,描述把一个稠密 FFN 切得多细)。拟合覆盖 1.04 亿到 27 亿激活参数(总参数最高 790 亿),稀疏度 0.5 到 0.981,G 取 {1,2,4,6,8}。G 的指数 η≈0.95 是整张拟合表里仅有的两个被数据真正钉死的系数之一,含义很实在:在固定显存下,更细粒度(G 更大)的专家确实能稳稳削弱稀疏带来的 loss 惩罚。
第二块是一个性能模型,从模型规格和设备配置直接预测 MFU、内存占用和最优并行布局。它把一次迭代拆成 Transformer 计算、流水线气泡、词表阶段不均衡、集合通信、optimizer step 几块,再对两类开销做显式修正:小模型上 CPU 派发 kernel 的 launch 开销,大集合通信上跨节点 straggler 和跳数开销。
为什么用「几何阶梯」上的离散单步搜索,而不是先连续优化再取整?因为边界问题:模型算力的最优值就贴在边界上,取整救不回来。
| 指标 | 数值 |
| Scaling law 拟合 R² | 0.99877 |
| 留出 RMSE(每 G 分层 top 10%) | 0.00779 loss |
| 拟合覆盖激活参数 | 1.04 亿 – 27 亿(总参最高 790 亿) |
| 性能模型 MFU 误差(700M/5.4B/18B 扫描) | 9.6% / 9.0% / 13.2%,均 < 15% |
| 布局排名 Top-5 重叠(5.4B/18B vs 700M) | 100% / 80% |
最直接的对比在图 1:在 Cmodel = 3.3×10²¹ 这个预算下,纯模型算力的最优稀疏度是天花板 Smax = 0.985。换成 4 个 AWS p6-B200 节点跑 5 天的真实硬件,可行边界只到 S = 0.96,再往上配置交付的算力反而更少。MOSAIC 在硬件约束下找到一个内部的集群最优稀疏度,其预测 loss 比那个边界最优还低 0.031 nats。
图 8 给出这种「反向排序」的证据:在 512 卡(64 个 p6 节点)规模上,每交付一个模型 FLOP 所需的 wall-clock 随稀疏度单调上升,最高达到最稠密配置的 1.70 倍,和 loss 随稀疏度下降的排序正好相反。这种相反正是为什么必须把算力预算做成随架构变化的东西。
专家切分因子 G 的效应(图 6):预测 loss 随 G 单调下降;实测跑到 G=4 还在降,到 G=8 在运行噪声内走平甚至略反转。
对任何要训稀疏 MoE 的人:脱离系统单独看 scaling law 选架构,可能拿到一个集群根本跑不高效、甚至按真实 GPU-小时算更差的配置。架构和系统是耦合的,一起优化能在更少的 GPU-小时里达到同一个 loss。这个框架的实用之处在于它输出并行布局,不只给模型维度。
诚实地讲,这是一篇协同设计与优化框架的工作,不是新模型,也不刷 benchmark SOTA。它真正立得住的两个产出:一个是把 G 加进 scaling law、并发现 η≈0.95 被可靠识别的经验结果;另一个是「纯模型算力下不存在内部最优稀疏度」这个概念性结论,直接解释了为什么需要协同设计。