🤖 machine learning

Compute-Optimal Is Not Cluster-Optimal: Systems-Aware Scaling for Sparse Mixture-of-Experts

本文介绍了 MOSAIC，这是一个系统感知的协同设计框架，它证明了只有通过将模型架构与硬件约束进行联合优化，而非仅仅通过传统的计算最优缩放法则，稀疏混合专家模型的最佳稀疏性才会显现。

原作者： Soumajyoti Sarkar, Yuxin Tang, Sheng Zha

发布于 2026-08-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Soumajyoti Sarkar, Yuxin Tang, Sheng Zha

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图构建一个终极、最强大的机器人大脑。在人工智能的世界里，这个“大脑”是一个庞大的计算机程序，被称为大语言模型（LLM）。为了让这些大脑变得更聪明，科学家们发现了一个简单的规则：如果给它们喂入更多的数据并让它们变得更大，它们就会变得更擅长理解世界。这个规则被称为“缩放法则”（scaling law）。长期以来，构建这些大脑的配方一直分为两个独立的步骤。首先，一位数学家会根据可用的计算能力（称为“算力”）来计算出大脑的最佳尺寸。然后，另一位工程师会尝试将这个大脑的大小压缩到特定的计算机芯片上，以使其运行得尽可能快。

你可以把这想象成规划一次公路旅行。第一步是根据你的汽油预算决定你想开多远。第二步是挑选汽车。但如果由于你选的汽车太重、太笨拙，导致它极其耗油，这意味着你实际上无法开到原计划的距离，那该怎么办？这就是这篇论文要解决的问题。它认为，你不能先选定一个大脑尺寸，然后再去担心汽车。你必须将大脑和汽车进行协同设计，因为大脑的形状会改变汽车行驶的效率。如果你忽略了汽车的引擎，你最终可能会得到一个“完美”的大脑，但它却太重而无法移动，让你只能困在起跑线上。

问题所在：“完美”却无法适配的大脑

这篇论文介绍了一种名为 MOSAIC（通过系统感知训练进行模型优化的协同设计）的新思维方式。来自亚马逊 AGI 基础团队（Amazon AGI Foundations）的作者们注意到，旧有的做法正在浪费资源。传统上，研究人员会计算一个“计算最优”（Compute-Optimal）的模型。这是指在固定的计算能力下，能获得最佳结果的模型尺寸，前提是假设每一分算力都被完美地利用了。

然而，作者们发现了一个隐藏的陷阱。他们关注的是一种特殊的 AI 架构——混合专家模型（Mixture-of-Experts, MoE）。想象一个巨大的图书馆，与其由一个巨大的图书管理员阅读所有的书，不如拥有数千个微小的专业专家。当一个问题进来时，一个智能路由会将问题发送给仅有的几个专家（例如 100 个专家中的 2 个）来回答。这使得模型非常快速且高效，因为它不需要为每个问题都唤醒整个图书馆。

旧的数学理论认为，为了获得最佳效果，你应该尽可能扩大图书馆的规模，并唤醒尽可能少的专家。换句话说，“完美”的模型应该是极其稀疏的（大部分是空的）。论文指出，如果你只看“我们需要多少次计算？”这种纯粹的数学逻辑，答案总是：“让它尽可能稀疏。”数学上的最优稀疏度始终处于可能性的边缘。

转折点：汽车引擎至关重要

故事在这里发生了变化。作者们意识到，虽然数学说“要极度稀疏”，但计算机芯片的物理现实却说“没门”。

当你拥有一个拥有数千个微小专家的模型时，计算机必须做大量的额外工作，仅仅是为了决定“唤醒哪些专家”以及“如何在专家之间传递信息”。这就像有一个巨大的办公大楼，经理整天都在各个隔间之间奔波递送备忘录，导致实际工作的时间寥寥无几。模型变得越稀疏，计算机在这些“奔波”任务（称为通信成本）上浪费的时间就越多，而用于实际思考（计算）的时间就越少。

论文认为，旧有的“计算最优”数学模型是有缺陷的，因为它假设计算机可以将 100% 的力量都投入到思考部分。而在现实中，一个超稀疏的模型可能只有 8% 的算力用于实际思考，而一个稍密一点的模型可能会有 15%。

解决方案：MOSIC

为了解决这个问题，团队构建了 MOSAIC。他们不再问：“对于固定的数学计算量，最好的模型是什么？”而是问：“在我们的特定计算机集群上，在固定的时间内，我们实际能运行的最优模型是什么？”

他们结合了两个要素：

缩放法则（Scaling Law）： 一个基于模型大小和形状来预测其智能程度的模型。
性能模型（Performance Model）： 一个模拟器，它可以预测模型在真实硬件上的运行速度，并将“奔波”成本考虑在内。

当他们运行 MOSAIC 时，结果令人惊讶。那个“完美”的模型并不是位于稀疏度极端边缘的模型。相反，最好的模型是一个内部解（interior solution）——即一个皆大欢喜的中庸之道。它足够稀疏以保证效率，但又不会过于稀疏，以至于让计算机因为管理那些微小专家而陷入停滞。

他们的发现

团队在使用了 NVIDIA B200 GPU 的真实硬件上进行了测试。他们训练的模型规模从 7 亿到 180 亿个激活参数不等（总参数量高达 790 亿）。

旧方法： 如果你仅仅遵循旧的数学理论，你会选择一个稀疏度约为 0.985（意味着 98.5% 的专家处于睡眠状态）的模型。但在他们的特定集群上，这个模型会非常慢，导致训练时间过长，并且它的智能程度实际上会低于一个稍密一点的模型。
MOSAIC 方法： 系统发现，那个“甜点位”（sweet spot）是稀疏度约为 0.96 的模型。这个模型并非“理论上”的最优，但它是“实践上”的最优。它运行得更快，能更有效地利用计算机的算力，并且实现了更低的误差率（意味着它更聪明），甚至超过了那个“理论上”最好的模型。

事实上，他们证明了那个在纸面上看起来最好（拥有最多“模型 FLOPs”）的模型，在现实生活中反而是训练速度最慢的。最终胜出的模型，是那个平衡了数学逻辑与硬件现实的模型。

为什么这很重要

这篇论文表明，我们不能再将 AI 模型的设计与运行这些模型的计算机设计视为两个独立的步骤。你不能仅仅选定一个模型尺寸，然后寄希望于计算机能够处理它。模型的形状会改变计算机的行为。

作者们谨慎地指出，他们的发现是针对他们所使用的特定计算机芯片（NVIDIA B200）以及他们开发的特定软件的。他们并没有发现一个适用于宇宙中每台计算机的通用法则。然而，他们证明了在 AI 训练的前沿领域，忽视“系统”（即硬件以及数据如何移动）会导致次优的结果。

通过使用 MOSAIC，他们展示了“最好的”模型并不是图表上的一个固定点，而是一个取决于你的特定硬件预算的移动目标。最有效的模型是那个像手套一样完美贴合你的集群的模型，而不是那个在纸面上看起来最完美的模型。这是一种从“计算最优”（数学所言）向“集群最优”（硬件所能实际交付）的转变。

简而言之，如果你想构建最聪明的 AI，你不仅需要一个更大的大脑，你还需要一个能完美适配其躯体的大脑。

技术摘要：计算最优并不等于集群最优

问题陈述
当前的规模化预训练工作流通常将算法缩放决策与系统实现分离。一个“缩放定律”（scaling law）阶段用于选择模型架构和数据预算，以在固定的计算量（FLOPs）约束下最小化损失；随后是一个脱节的“系统”阶段，旨在针对硬件效率优化实现方案。作者认为这种分离对于稀疏混合专家（MoE）模型而言是存在缺陷的。传统的缩放定律优化的是模型 FLOPs ( $C_{model}$ )，却忽略了所选架构的系统效率。两种具有相同预测损失和 $C_{model}$ 的架构，由于模型 FLOPs 利用率（MFU）、通信成本和内存占用率的差异，可能会产生截然不同的训练吞吐量。因此，基于 FL-OPs 预算得出的“计算最优”配置在特定的硬件集群上可能是不可行的或次优的，因为它未能考虑到集群在固定训练窗口内实际能够维持的“可交付”计算量。

方法论：MOSAIC
作者引入了 MOSAIC（模型优化与系统感知训练协同设计），该框架将模型架构与系统协同设计表述为一个统一的优化问题。MOSAIC 共同选择模型架构、训练 Token 预算以及分布式执行布局（并行策略），旨在硬件约束的预算下最小化预测损失。

该框架集成了三个核心组件：

四维联合缩放定律： 作者为稀疏 MoE 模型拟合了一种新的缩放定律，其中包含了总参数量 ( $N_{tot}$ )、训练 Token 数 ( $D$ )、稀疏度 ( $S$ ) 以及专家拆分因子 ( $G$ )。 $G$ 被定义为稠密前馈层宽度与单个专家宽度的比值（ $G = d_{ff}/d_{expert}$ ），捕捉了专家的粒度。该定律公式化为：
$L(N_{tot}, S, D, G) = \frac{a}{N_{tot}^\alpha} + \frac{b}{D^\beta} + \frac{c}{(1-S)^\lambda} + \frac{j}{(1-S)^\delta N_{tot}^\gamma G^\eta} + e$
这种形式允许基于离散的几何选择而非抽象的参数计数来预测损失。
算子级性能模型： 为了估计系统效率，作者开发了一个解析性能模型，用于预测给定架构和并行布局下的 MFU 和内存占用。该模型考虑了：
- 计算成本（稠密与分组 GEMM、注意力机制）。
- 通信开销（专家并行的 all-to-all 通信、流水线气泡、数据并行集合通信）。
- CPU 分发延迟（对于小微批次至关重要）。
- 针对节点间落后节点（stragglers）和节点内重叠（overlap）的硬件特定校准因子。
  该模型通过高达 180 亿活跃参数的训练运行进行了验证，其 MFU 预测的平均绝对百分比误差（MAPE）低于 15%。
联合优化 (MINLP)： MOSIC 求解一个混合整数非线性规划（MINNLP）问题，在要求模型 FLOPs ( $C_{model}$ ) 不超过可交付模型 FLOPs ( $C_{deliverable}$ ) 的约束下，最小化预测损失。
$C_{deliverable} = C_{peak} \times \text{MFU}(Z, P_\ell, B) \times \eta_{good}$
此处， $C_{peak}$ 是原始硬件上限，而 MFU 取决于架构。优化过程在离散的可实现几何结构 ( $Z$ ) 和执行计划 ( $P_\ell$ ) 网格中进行搜索，以找到能在集群实际能力范围内使损失最小化的配置。

核心贡献

重新定义计算预算： 论文表明，对于稀疏 MoE 模型，在固定的 $C_{model}$ 预算下进行优化会导致“边界寻求”行为，即最优稀疏度位于数据支持的上界（即最稀疏的模型）。这是因为缩放定律预测在固定 FLOPs 下更稀疏的模型具有更低的损失，却忽略了系统成本。
系统感知的协同设计： 通过将缩放定律与性能模型耦合，MOSAIC 将预算从 $C_{model}$ 转向了 $C_{deliverable}$ 。这揭示了最优稀疏度是一个由集群特定约束（内存、带宽、并行开销）决定的内点，而非边界点。
性能模型验证： 作者提供了一个经过验证的算子级性能模型，能够准确预测 MFU 并识别稀疏 MoE 的最优并行布局（张量、专家、流水线、数据并行），并证明了“最佳”布局高度依赖于具体的几何结构和集群规模。
实证验证： 该框架在 AWS p6 节点（NVIDIA B200 GPU）上进行了测试，运行规模高达 250B 总参数。结果证实，基于模型 FLOPs 的配置排名与基于硬件计算效率（墙钟时间）的排名存在显著差异。

结果

边界 vs. 内点最优： 在固定的 $C_{model}$ 下，拟合的损失随稀疏度的增加而单调递减，将最优值推向边界 ( $S_{max}$ )。然而，在固定的硬件预算 ( $C_{deliverable}$ ) 下，损失曲线呈现出一个内点极小值。极度稀疏模型的系统成本（通信、内存）最终会抵消其理论上的损失收益。
损失与吞吐量的权衡： 在一个 32 节点、20 天预算的实验中，实现最多模型 FLOPs 的配置（一种更稠密、更粗粒度的专家几何结构）并未实现最低损失。相反，损失最优的配置（一种更稀疏、更细粒度的专家几何结构）虽然维持了较低的 MFU，但实现了更低的最终损失，因为它是在保持有利损失轨迹的同时，唯一能符合可交付 FLOPs 约束的配置。
现实世界验证： 在 8 节点集群上的阶段性预训练运行验证了该框架。测得的 MFU 排序与预测的排序一致。此外，当绘制损失相对于“峰值等效硬件计算量”（考虑了 MFU）的曲线时，配置的排名与基于原始模型 FLOPs 的排名发生了翻转，证实了“计算最优”模型并非“集群最优”模型。

意义
本文主张在前沿语言模型的训练方式上进行根本性转变：从脱节的“缩放定律选择”与“系统调优”阶段，转向统一的架构-系统协同设计。作者指出，仅依赖模型 FLOPs 预算对于稀疏 MoE 模型是不够的，因为它未能计入架构的系统效率。通过引入 MOSAIC，他们证明了最优的模型几何结构并非一个普适常数，而是取决于特定的硬件集群、训练窗口以及并行约束。这种方法确保了所选架构不仅在理论上是损失最优的，而且在目标硬件上是切实可行且高效的。

技术摘要：计算最优并不等于集群最优

类似论文