看到“专家池从 8 扩到 128,训练速度只慢了 5%”这个数字时,我第一反应是:这是在逗我?做了这么长时间 MoE 训练的人都知道,专家越多,通信这张网就越复杂,能把额外开销压到 5% 这个量级,已经不是算法问题了,是工程能力问题。更别说题目里还挂着“万亿参数”和“开源”两个词,前者说明这套框架要应付的是大规模集群训练,后者意味着这套积累不是藏在内部代码里,而是可以拿过来直接做解构和扩展。
这篇内容我想结合万亿参数 MoE 训练的实际工程经验,拆一下 Olmo-core 3 这个框架为什么能做出这种效果,以及它背后的设计取舍。适合三类人看:一是正在做或者准备做 MoE 预训练的工程师,二是关注训练框架选型的研究团队,三是想把大模型训练从“跑通”推到“跑得更快”的同学。我会把 MoE 训练里最容易被忽视的通信成本、专家负载均衡、框架包装层的设计逻辑、以及实操时怎么从 8 个专家往 128 个专家迁移,一次说清楚。
1. 8 个专家变成 128 个专家,为什么“只慢 5%”让人意外
1.1 MoE 稀疏计算?计算量不是主要成本
先把一个常见的误解拆干净:很多人以为专家池从 8 扩到 128,等于模型前向的计算量涨了 16 倍,所以慢 5% 是吹牛。但 MoE 架构的核心是“稀疏激活”,不是“稠密计算”。
在 Mixture of Experts 模型里,每一层会布置多组前馈网络,也就是专家。输入 token 经过一个路由网络(router),根据特征打一个分配权重,选出排名最靠前的 k 个专家,然后只把这 k 个专家算一遍。比如num_experts=8, top_k=2,那么每个 token 只经过 8 个专家里的 2 个,计算量大约是单专家版本的 2 倍,而不是 8 倍。把专家数从 8 涨到 128,只要 top_k 不变,单次前向的计算密度基本不变,或者只随 top_k 微调。这就是“稀疏”带来的第一层红利。
所以“万亿参数”这个数字要拆开看:真正的稠密激活参数,可能只有几十亿到一两百亿级别,剩下的大头全是专家参数。专家参数平时沉在那里,要靠路由决定哪些被激活。训练时,总参数量的增长主要体现在显存占用和通信带宽上,而不是 GPU 的浮点运算量上。
1.2 慢的根源在通信与显存,不在浮点运算
既然算子计算量没有随专家数暴涨,那真正让训练慢下来的是什么?是通信。
先说显存。每增加一个专家,模型权重体积就增加一份前馈网络的参数。8 个专家时,一台机器可能装得下全部专家副本;128 个专家时,专家参数大概率要切到多张卡、甚至多个节点上。这就引入了专家并行(Expert Parallelism)。参数放得越散,token 的中间激活就要跨设备传送。
再说通信。MoE 层的前向后向里,每个 token 都要把隐藏状态发送到它被分配到的专家所在设备。这个操作对应到分布式训练里,就是大批次的 All-to-All 通信:每个节点都要同时给其它节点送数据、收数据。在这个交换的过程中,batch 里的 token 要按路由结果重新排列、发送、接收、还原。专家数量越多,路由的分布越散,这个“转盘调度”变得越复杂。
在我的经验里,普通 MoE 框架在专家规模扩大后,训练总耗时涨 10% 到 30% 都很常见。所以 Olmo-core 3 声称 8 到 128 只慢 5%,这个数字放在全行业都很有含金量。它背后一定不是单点优化,而是从通信重叠、路由调度、显存布局到训练范式的整体设计。
2. MoE 训练卡在哪:路由、All-to-All 与专家负载失衡
2.1 一次前向里,token 是怎么被“送往”不同专家的
要理解 MoE 训练为什么难搞,得先把一次前向的数据流说透。
假设模型有 32 层 transformer,其中每一层都有 128 个专家。输入的 token 经过注意力层之后,得到一组形状为[batch, seq_len, hidden_size]的隐藏状态。此时路由网络会对每个 token 计算一个概率分布,选出概率最高的 2 个或 4 个专家。这个决策结果通常会转成一个“排列索引”:哪些 token 去专家 3,哪些去专家 47,哪些去专家 99。
接下来框架要做的是把一个 batch 里的所有 token,按照路由索引,重新分装成若干个小包,转发给专家所在设备。每个目标设备把收到的 token 送入本地专家,做完前向计算,再把结果按原路径传回发起设备。这里最耗费资源的不是专家本身的矩阵乘法,而是“分组、排序、打包、跨设备投递、回传、还原”这套动作。
我把这种通信模式比作公交总站的调度:一批乘客到了站台,要先看路线、查班次、临时分队,然后分别上不同的车,到站后再按原队形集合。线路越多,调度中心要处理的交换信息就越多。虽然每个 token 真正坐的车只有 k 趟,但调度系统要处理的信息量跟网络里的专家分布强相关。
2.2 专家池变大之后,原来不明显的成本被放大
专家数从 8 涨到 128,不是“多配几组参数”这么简单,至少三个维度的成本会被放大。
第一,通信矩阵变稀疏,但非零事件变多。8 个专家分布在 8 张卡时,每张卡之间可能只有一两条交换链路。128 个专家分布在 32 张卡上时,任意一张卡都可能和其它卡发生 token 交换。如果每个 token 都路由到随机专家,通信模式会迅速逼近“全互联”,网络流量瞬间上涨。
第二,显存布局变复杂。8 个专家时期,单个专家权重可能只占一小块显存,每张卡可以完整放下很多专家。128 个专家时,如果你不想让某个 GPU 因为装载过多专家而爆显存,就得把同一个专家切到多张卡上,或者把专家均匀铺开。这两种做法都有代价:前者增加通信,后者仍然增加通信。核心矛盾是:专家规模越大,想要的负载越均匀,通信就越多。
第三,负载失衡被放大。路由网络在训练初期并不稳定,经常出现“热门专家”和“冷门专家”。当专家池小的时候,热门专家排队等待的时长还能被其它任务吸收;128 个专家时,一小撮专家收到大量 token,其它专家空转,整个 All-to-All 的交换效率就会断崖下跌。这也是为什么所有工程化 MoE 框架都必须引入负载均衡辅助损失(load balancing loss)或容量因子约束,让大家都有活干。
2.3 要把开销压在 5% 内,工程上得拆掉哪几颗“雷”
现在回到标题里的那个数字。想做到 8 到 128 只慢 5%,肯定绕不开下面这几颗雷。
第一颗是通信和计算的重叠。最笨的实现是先启动 All-to-All,等所有 token 都到了目标设备,再开始专家计算。这样通信时间完全暴露在关键路径上。高效的框架一定会在做完路由决策后,把一部分 token 先发送出去,同时让已经到达的设备立刻开始算,把通信当作流水线的一部分“藏”在计算后面。我称之为“边发边算”,这也是现在主流大规模训练框架都在做的方向。
第二颗是梯度同步与参数更新的错峰。MoE 层多出来的通信不仅在数据交换,还在反向传播后的梯度同步。如果每个专家梯度都等在同一个通信组里求和,网络基本就堵死了。常见的做法是把专家梯度按设备分组、按层分组,再用异步的通信算子完成聚合,让优化器更新在部分梯度到位后就开始。
第三颗是低精度通信。对 MoE 这种通信密集型的场景,激活和梯度转一圈,流量可能比参数本身还大。使用 BF16、FP16 甚至量化后的通信包,能大幅度压缩带宽开销。当然代价是需要额外的反量化处理,以及对数值稳定性的仔细验证。这部分往往是框架源码里最“晦涩”的地方,却是收益最明显的地方。
我自己测过不同框架的 MoE 性能,结论很一致:算法决定了你能训练多大模型,框架工程决定你能训得多快。所以看到“5%”这个数字,我首先想的不是“专家扩容容易”,而是“这个框架在通信、负载、重排、调度上做了一整套组合拳”。
3. Olmo-core 3 的框架设计:从训练范式包装到底层调度
3.1 同一个 MoE 引擎,要同时服务 SFT 和 GRPO 两类训练范式
我注意到这个框架的公开资料反复提到一个关键词:训练范式包装。典型的名字包括 SFTTrainerWrapper、GRPOTrainerWrapper。这个设计很关键,因为如今的训练已经不只有“预训练”一个动作了,大量团队会在基座做完后接着做 SFT,甚至用强化学习范式微调模型。分类看,SFT 和 GRPO 对训练循环的要求几乎完全不同。
SFT,监督微调,是标准的输入输出对拟合,损失函数通常是交叉熵。它的数据流很简单:读数据、前向、算 loss、反向、更新。GRPO,组相对策略优化,是一种强化学习训练范式:给模型一组提示,让模型先生成多个候选回答,再用一个组内的相对胜负关系作为奖励信号,更新策略网络。它涉及到采样、奖励模型或规则评估、优势计算等步骤,训练循环明显更复杂。
如果团队用同一套底层 MoE 引擎同时支持这两种范式,就得把“训练流程”抽象成可插拔的组件:底层分布式调度、模型结构、梯度逻辑保持一致,上层把数据流、损失计算、采样策略封装成不同训练器。这也是我看到 SFTTrainerWrapper、GRPOTrainerWrapper 这个名字时的理解:框架的核心价值不只是“快”,而是“同一套底层,兼容多种训练路线”。
这就解释了为什么研究团队选框架时会越来越看重训练循环的抽象层。像我们自己做实验,经常是早上跑 SFT,下午切 GRPO,如果框架把两种流程硬编码成两套工具,维护成本会很高。wrapper 就像一套接口协议,让开发者在上层写业务逻辑,底层优化由框架统一接管。
3.2 专家并行与张量并行共存时,谁来决定参数放哪
另一个让我关注的地方是并行策略。大模型训练框架绕不开并行维度:数据并行(DP)、张量并行(TP)、专家并行(EP)通常是组合使用的。
在一个 MoE 模型里,非 MoE 层,比如 attention、embedding,通常用张量并行切分。张量并行会把同一层的矩阵切到多张卡上,让每张卡持有部分权重和部分计算。MoE 层则不一样,每个专家通常是完整的、独立的模块,适合按专家维度做专家并行,让不同专家落在不同设备上。
问题是:在 128 专家、万亿参数量的规模下,这两种并行要同时操作,谁先通信、谁后计算、参数放哪,就不再是算子级问题,而是调度系统和并行拓扑设计的问题。框架需要根据硬件拓扑,决定每个专家放在哪个节点、每张卡上放几个专家、token 跨设备传输时走哪条链路、是否需要把通信和通信合并。这已经不是单个算子能控制的复杂度,需要整套框架层面来统一规划。
在 Olmo-core 3 里,我看到它在公开资料里强调“面向 MoE 优化”,这通常意味着它对专家并行的支持不是简单的把专家塞到设备里,而是把路由、通信、计算三条链路做成一张可调度的图。哪个专家离用户请求更近、哪次 All-to-All 可以批量执行、什么时候启动唤醒缓冲区,这些细节会决定最终的吞吐和延迟。
3.3 AI2 把它开源,对想做大模型训练的人意味着什么
说到 AI2 开源,我最直接的感受是:终于有机构愿意把万亿参数级别的“工程经验包”整个开放出来了。现在很多 MoE 框架,要么是实验室原型,要么是商业公司内部工具,能同时满足开源、MoE 原生优化、覆盖 SFT/GRPO 训练范式的选择并不多。
这对研究团队和中小规模的训练团队是实打实的好事。如果你是刚起步做 MoE 训练的团队,直接在大模型基础上去补通信调优和负载均衡,成本很高。但如果你能拿到一个已经验证过的框架,跟着它的设计思路做二次开发,就能把更多精力放在模型结构和数据配方上,而不是每天和各种通信 bug 搏斗。
另外,开源这件事对“可复现性”也很重要。很多论文里写“我们用 MoE 模型、32 个专家”,但没人告诉你规模上去后通信怎么掉了一地。开源框架相当于给了你一个基准线,你可以在同样代码下跑自己的实验,对照官方数值验证你的硬件配置、分布式设置和数据处理是否到位。“开源”这两个字,在这类庞大的训练系统里,价值远不止省去开发成本。
4. 从 8 个专家跑到 128 个专家:建议的实操路径
4.1 先跑通最小配置,再谈扩充专家池
我知道很多人拿到框架的第一反应是直接拉满配置,128 专家、万亿参数体量、几百张卡,上来就看效果。但实际上最稳的路径是先跑通小规模配置。你至少需要一条能正常收敛、能稳定跑完几步更新、能保存 checkpoint 的完整链路。
具体步骤大概是这样的:先把框架克隆到本地,准备好训练数据,然后修改模型配置。这里给一份示意配置,实际字段以你手上框架的具体接口为准,但思路是一样的:
# 示意配置:先从小处验证链路 model: name: "moe-small" hidden_size: 2048 num_layers: 12 num_experts: 8 top_k: 2 moe_layer_interval: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] train: per_device_train_batch_size: 8 gradient_accumulation_steps: 8 learning_rate: 5e-4 warmup_steps: 100 max_steps: 5000 sequence_length: 2048 parallel: data_parallel_size: 4 expert_parallel_size: 2 tensor_parallel_size: 2跑起来之后,先确认几个最基本的点:loss 能不能稳定下降;日志里的 step time 是否平滑;GPU 利用率是否在合理区间。如果这一步都出问题,比如显存爆掉、通信循环卡死、token 分布失衡,先解决这些小配置里的毛病,再去放大专家池。
小规模跑通还有另一个好处:它能给你一个 baseline。比如你现在用 8 专家,稳定训练速度是 1000 token/s。后面扩到 128 专家,所有效果都跟这个 baseline 比。不然你直接上大配置,根本分不清速度变化是来自专家扩容,还是来自数据加载、网络抖动、超参变化。
4.2 扩展之前要盯住哪几个指标
从 8 专家扩展到 128 专家前,我会在配置文件里打开一段时间的观测窗口,每天记录下面这些指标,别等碰到问题再去翻日志。
第一个是端到端吞吐,我习惯用 tokens per GPU per day 来看。它综合反映了计算、通信、显存三者的平衡,训练快不快最终看的就是这个指标。
第二个是 step 的时间构成拆解。框架一般会记录 data load、compute、communication、optimizer 各自的耗时。如果数据加载占了 30% 以上,说明 datapipe 需要换思路,这个问题跟专家数无关,却会干扰你对 MoE 扩容的判断。
第三个是显存曲线。记录峰值显存和平均显存,特别是专家参数和中间激活的占比。有时候扩专家池后显存没炸,但把过大的 activation 撑到了显存边缘,会造成换页或者程序抖动。
第四个是专家热度分布。统计每个专家接收到的 token 数量,判断是否存在明显的“头部专家”和“冷门专家”。如果某几个专家接收的 token 量是平均数的 2 倍以上,说明路由已经在失衡,辅助负载均衡损失需要调大,或者容量因子需要变化。
我建议用一张表把这些指标固定下来,隔两天做一次对比:
| 指标 | 8 专家 | 128 专家 | 变化 | 预警线 |
|---|---|---|---|---|
| 端到端吞吐(tok/gpu/day) | 20万 | 19万 | -5% | 超过-10%需排查 |
| 单步耗时 | 1.6s | 1.68s | +5% | 超过+10%需关注 |
| All-to-All 通信占比 | 12% | 20% | +8% | 超过30%需优化拓扑 |
| 峰值显存 | 40GB | 52GB | +30% | 接近卡容量需切分 |
| 热门专家 token 占比 | 18% | 30% | +12% | 超过35%需调负载均衡 |
这张表不是为了好看,是为了在扩展专家池的过程中快速定位瓶颈:如果是通信占比涨得夸张,优先调网络拓扑或通信重叠;如果是显存涨到接近瓶颈,优先调整专家分片策略;如果是路由倾斜,优先动 loss 和容量约束。
4.3 128 专家不是终点,超参与环境都要重新审视
很多人会忽略一个点:从 8 专家扩到 128 专家,不只是“把 num_experts 改大”那么简单,训练超参也要重新审视。
首先是 top_k。8 个专家选 2 个,相当于四分之一的专家参与计算。到了 128 个专家,如果还是选 2 个,专家层的信息聚合可能太稀疏,模型容量利用不充分;如果盲目把 top_k 调到 4 或 8,计算量又会明显上升。这里需要重新跑一组小规模消融,观察吞吐和 loss 之间的平衡。
其次是辅助负载均衡损失的权重。专家池越大,路由网络背负的压力也越大。如果没有负载均衡约束,很容易出现几个专家承接绝大部分 token,其余专家长期闲置。但约束太狠,路由又会趋于均匀化,失去专家分工的意义。我通常会先以很小的权重起步,比如 0.001,然后根据训练日志里的专家热度方差逐步调整。
再一个是学习率和 warmup。专家变多,模型的有效表达能力变强,很多情况下需要适当降低学习率或者拉长 warmup。直接沿用前面的策略,loss 曲线可能看起来能降,但最终收敛质量和稳定性会变差。
还有一个很容易被忽略的工程细节:checkpoint 的体积。专家参数翻倍之后,保存一份 checkpoint 的时间会成倍增加,严重时候还会造成训练空窗期。我建议启用异步 checkpoint,并把保存频率从固定 step 改成基于时间间隔或验证集效果,避免训练节奏被写盘拖垮。
5. MoE 训练中我踩过的坑和排查清单
5.1 看到训练变慢,先拆时间线再下结论
做过大规模训练的人基本都遇到过:某天训练速度突然从 2000 token/s 掉到 1700 token/s,第一反应是框架不行。但我在实际排查里发现,时间线拆开后问题经常在别处。
我的建议是,只要训练变慢,先做三件事。第一,把 step 里的时间拆出来,看 data loader、前向、反向、梯度同步、优化器更新各占多少。第二,看 GPU 的利用率曲线,如果是锯齿状频繁跳动,大概率是等待同步;如果利用率很高但吞吐上不去,可能是路由或通信模式有问题。第三,检查日志里的 loss 曲线,看它是不是还在正常下降。如果 loss 还在降,很多“慢”其实是心理预期或硬件自然波动造成的,不用急着动配置。
我自己曾经因为把 All-to-All 通信和异步调试信息混在一个日志里,误以为通信占了 60% 的时间,实际核对了 timeline 后,发现是数据预处理阶段在 token 划分上做了太多额外的重复拼接,无端消耗了 CPU 和内存带宽。排查链路不清晰,很容易被表面现象带偏。
5.2 专家不均衡:损失曲线正常,模型能力却是偏的
这是 MoE 训练最隐蔽的坑之一。有些时候你盯着 loss 看,一切正常,但训练结束后的评测效果并不好,尤其是某些领域任务表现差。最后查下来发现,优质 token 都集中到了少数几个专家里,其他专家参数没有充分参与训练,等于你的模型只用了专家池的三分之一。
要避免这种情况,我习惯每训练几千步就输出一次专家统计:每个专家接收的 token 数、平均 top-1 概率、以及所有专家权重梯度的 L2 范数。如果梯度基本只在热门专家上运动,说明冷门专家没有贡献,这个时候要适当提高负载均衡损失权重,或者检查容量因子是否设置得太小,导致某些专家被频繁丢弃 token。
还有一招是换一种路由初始化。在启用大专家池前,对路由网络做一次特别初始化和预热,让 token 一开始就分布得相对均匀,比训练中后期再纠正要省力气得多。这有点像组织里的任务分配,开头就把人和任务分均匀,后期出问题的概率小很多。
5.3 “只慢 5%”的适用边界:硬件、网络、数据规模缺一不可
最后想提醒所有读到这篇内容的人,“8 到 128 只慢 5%”是一个有前提条件的结论,不要直接照搬到自己的场景。这个数字要成立,通常依赖几个条件:节点数量足够多、节点间网络带宽足够高、batch size 足够大、每一步训练里要交换的数据量足够充足,让通信时间有机会被计算时间覆盖。如果你的集群只有单机多卡,或者互联还是普通的万兆以太网,那么通信开销可能完全盖过计算,慢 5% 就会变成真实的 20%。
所以正确姿势是:在拿到框架后,先按自己的硬件环境固定一组可控变量(卡数、数据规模、上下文长度、batch size),跑一个 8 专家的 baseline,再扩展到 128 个专家做 A/B 对比。用自己的时间线数据和指标表验证“5%”这个结论在自己的环境里成不成立,而不是拿着别人的数字就往上套。
基于我自己的经验,还有一个特别实用的判断方法:先用同一个框架在相同 batch size 下测 64 专家和 128 专家,如果两者差距很小,说明通信重叠做得很好;如果差距一下被拉开,瓶颈大概率出现在路由或者网络拓扑上,而不是专家本身的浮点运算。
MoE 训练不是把专家数加到很大就万事大吉,它是一个典型的木桶工程,哪块板最短,性能就丢在哪块。框架层面能把专家池扩展造成的损耗压到 5%,说明它在通信调度和训练范式抽象上做了大量看不到的工夫。我个人的体会是,想真正利用好这类开源框架,第一步反而是静下心来读懂它的并行策略和 wrapper 设计,然后在小规模配置上建立自己的 baseline。拿到“只慢 5%”这种数据时,不要只记住结果,把它放在自己的硬件、数据和超参条件里重新验证一遍,你的收获会比看十篇评测文章都大。