看到“专家池从 8 个扩到 128 个训练速度只慢了 5%”这个说法我的第一反应是标题党。MoE 模型我拆过不少专家一多通信、内存、调度全是事怎么可能只慢 5%后来去翻了 Ai2 放出来的 OLMo-core 3 相关代码和文档才发现这个结论不是拿来唬人的而是 MoE 稀疏激活里一个很基础的收益被工程化做到了极致扩大专家池不等于让每个 token 算更多东西。这篇文章就把这件事拆开讲清楚。你会看到 OLMo-core 3 作为 Ai2 开源的万亿参数训练框架到底是怎么组织 MoE 专家池的“8→128 只慢 5%”这笔账怎么算以及 SFTTrainerWrapper、GRPOTrainerWrapper 这类训练接口解决了什么问题。如果你正准备做大规模 MoE 预训练、微调或 RL又不想从零搭分布式状态管理这篇文章可以帮你省不少试错时间。1. 先别急着感叹“标题党”这个框架到底干了什么事OLMo-core 3 是 Ai2 在开放语言模型项目里沉淀出来的训练核心库。它不只是一个模型仓库而是一整套训练基础设施把预训练、指令微调、强化学习对齐都收进了同一套 API 里。和很多“只给权重不给锅”的开源模型不同Ai2 连训练代码、数据流水线、日志和评估脚本一起开源研究者拿到手不是一份推理 demo而是一个能继续往下训的工程底座。MoE 是这个框架最突出的一块。所谓混合专家架构简单理解就是把一个大的前馈网络拆成很多个“专家”每次 token 路过时只挑其中一小部分专家算。OLMo-core 3 对专家池做了专门设计让专家数量在配置里是一个普通字段你写 8训练系统按 8 个专家切你写 128系统按 128 个专家做并行和通信。难能可贵的是扩容专家池的代价没有让墙钟时间线性上涨而是被控制在一个很小的比例也就是标题里的“只慢 5%”。这个框架适合谁三类人最需要它一是研究 MoE 路由、负载均衡、专家可解释性的算法团队他们不想在工程细节里陷太久二是想复现“万亿参数”训练的机构需要一个能直接横向扩展的底座三是做 SFT 和 RL 的工程师OLMo-core 3 的 wrapper 设计让微调和策略优化不需要重写训练循环。前 100 字里已经有了三个核心关键词MoE、训练框架、专家池后面我会逐个拆开。2. “8→128 只慢 5%”这笔账到底是怎么算的2.1 稀疏激活省掉的是计算量不是参数理解这 5%先抓住一个核心事实MoE 的参数量和计算量是两回事。稠密模型里参数越多每个 token 前向传播要过的矩阵乘法就越多计算量基本跟着参数量走。MoE 不一样它有 128 个专家但每个 token 只被路由到其中 2 个专家常见 top-2 路由大部分专家对这个 token 来说只是“躺在显存里的参数”并不参与这个 token 的计算。所以从 8 个专家扩到 128 个专家单 token 的 FLOPs 几乎不变。前者是“从 8 个专家里挑 2 个算”后者是“从 128 个专家里挑 2 个算”两次都是算 2 个专家。真正变的是三样东西模型参数总量、显存占用、专家间通信规模。OLMo-core 3 能把墙钟时间控制在只慢 5%本质上就是把这“后三样”的增量用并行策略和通信重叠消化掉了。打个比方一个医院原来有 8 个诊室现在扩到 128 个诊室但每位病人仍然只看 2 个科室。医院的建筑面积存储变大了但每个病人的接诊流程没有变长。只要分诊台router够快、各科室之间转诊通道够宽整体就诊时间就不会因为“楼更大”而明显变慢。2.2 通信和显存开销是怎么被“藏”起来的专家池变大最头疼的其实是显存。128 个专家如果都放在一块 GPU 上光权重就把显存吃干净了。OLMo-core 3 的做法是专家并行把专家切成多组分散到不同设备上。每个 token 被路由到哪个专家这个专家可能在远端设备于是就有了 all-to-all 通信。这里有个关键点all-to-all 的通信量和 token 数量有关和专家总数不是强线性关系。你把专家从 8 个加到 128 个每次通信发送的仍然是一批 token只是需要发往的设备变多了。如果通信实现得粗糙小包太多会拖垮吞吐但 OLMo-core 3 这类框架会做消息合并和分组调度把“发往同一台设备的 token”攒成一个大包再发通信次数并不随专家数量线性暴涨。同时通信期间可以重叠计算当前 micro-batch 在算下一批 token 的专家权重已经在预取进显存了。这“预取 重叠”的操作把那部分肉眼可见的等待时间压到了很低。显存的问题则交给混合并行切分。MoE 训练里通常组合三种并行方式tensor parallel 把注意力矩阵切开expert parallel 把专家分到多卡pipeline parallel 切层。还有 sequence parallel 处理长序列。OLMo-core 3 等于把这几件事封装成了自动的布局策略你只管写模型结构它负责让每个专家块落在合适的位置上、让通信发生在最不占时间的那个节点。2.3 负载均衡才是那个“5%”的隐形敌人这里必须泼盆冷水8→128 只慢 5%是有前提的。前提就是负载均衡做得好。MoE 最怕“路由倾斜”所有 token 都往两三个专家身上挤剩下 126 个专家在旁边看戏。一旦发生这种情况计算图长这样——两台设备满负荷其余设备空转总时间由最忙的那台设备决定别说什么 5%能保持 50% 的扩展效率就算不错。OLMo-core 3 的负载均衡靠两层东西兜底。第一层是辅助均衡损失常见实现是给每个专家统计 token 分配比例加一项 load balance loss 进总 loss系数通常在 0.001 到 0.01 区间。这等于给 router 一个信号别太偏心。第二层是 router z-loss作用是稳定路由 logits 的尺度防止训练中 logits 数值漂移导致路由猛然崩溃。另外还有 capacity factor 配置表示每个专家在一个 batch 里最多接受多少 token。容量因子设太低会丢 token设太高会浪费算力通常从 1.25 左右开始调。所以“只慢 5%”更像是一个良好工程状态下能达到的极限值而不是无条件成立的物理定律。你在自己环境复现时如果看到扩展专家池后时间翻倍先别怀疑框架先检查专家利用率曲线是不是均匀的。这节内容虽然从思路上讲是“加专家不贵”但工程实现上贵的点被转移到了负载均衡和通信重叠上这也正是框架真正值钱的地方。3. 训练框架的一站式设计从 SFT 到 GRPO3.1 SFTTrainerWrapper监督微调不用重写训练逻辑OLMo-core 3 在微调侧给研究者提供了一组 wrapper其中最直观的是 SFTTrainerWrapper。所谓“包装器”就是在 Hugging Face Trainer 那套流程外面套一层 MoE 感知的能力。你自己写 SFT 最烦的不是模型结构而是数据怎么打包、loss mask 怎么算、多轮对话怎么截断。一个带 padding、包 sequence packing 的 data collator 能写吐人。SFTTrainerWrapper 把这些都内置了配置里指定“只对答案部分计算 loss”其余逻辑由 wrapper 屏蔽掉。它在底层做的就是监督微调的标准动作输入 token、输出 token、忽略 pad 和 prompt 部分只对 answer token 计算交叉熵。对 MoE 模型来说额外多干一件事在反向传播时处理专家的稀疏激活也就是保证只更新被路由到的专家。如果你用 LoRA 做参数高效微调wrapper 还会负责只把 adapter 挂到 FFN 层或 attention 层并且把没有路由到的那批专家权重冻结住。我实际用这类 wrapped SFT 训练 MoE 模型的经验是指令微调阶段不要轻易调低学习率。很多人在稠密模型上习惯 2e-5 这种保守值但 MoE 微调时如果 router 权重更新太慢会出现“能力迁移不到新专家”的现象。拿 OLMo-core 3 的 wrapper 跑的时候我会先用小规模样本验一版“router 输出分布是否变化明显”再决定要不要动学习率。3.2 GRPOTrainerWrapper把强化学习接进 MoE现在做对齐阶段绕不开 GRPO。GRPO 是 Group Relative Policy Optimization它和 PPO 最大的区别是不需要单独训练一个 critic 价值模型而是用同一个 prompt 采样一组 response用组内的相对 advantage 来更新策略。少一个 critic显存压力骤降对动辄几千亿参数的 MoE 模型来说非常关键。GRPOTrainerWrapper 做的事情是把采样、advantage 归一化、policy 更新、ref policy KL 惩罚打包成一条流水。你只需要提供“奖励函数”或“偏好信号”它负责跑 rollout、组装 group、控制更新步数。用 MoE 模型跑 GRPO 时有一个特有的坑如果 reward 变化剧烈router 的分布会被带着剧烈震荡表现为某些专家突然失活。OLMo-core 3 的 GRPO wrapper 里通常会有一层 logits 稳定机制同时把 batch size 和组内采样数设在一定范围内避免 router 被极端 reward 推偏。我第一次跑 GRPO 时吃了这个亏reward 只涨了一个点专家利用率却崩成一条竖线最后是靠调小学习率和 group 样本数压回来的。3.3 万亿参数并行切分的实际落地“万亿参数”四个字看起来简单落到训练框架就是一场并行切分的体操。假设你的模型有 128 个专家每个专家单独是一组 FFN 参数怎么把它们摆到 GPU 集群上最常规的做法是设 expert parallel size把 128 个专家平均分给若干个 expert parallel group。比如 32 张卡每 8 张卡组成一个 expert parallel group也就是本地放 16 个专家。注意层 conv 结构是attention 部分走 tensor parallelMoE 专家部分走 expert parallel每个 layer 的前向先做 attention再通过 all-to-all 把 token dispatch 到对应专家计算完再 all-to-all 回来。这套调度如果手动实现光是调试通信顺序就得花几周。OLMo-core 3 把它做成了模型配置里的几个字段剩下交给框架。万亿参数还有另一个隐藏问题优化器状态。Adam 状态本身是参数量的好几倍保存 checkpoint 时如果是“每卡一个标志位”的旧式做法恢复训练会让人崩溃。OLMo-core 3 这类框架现在普遍做异步 checkpoint 和统一的 sharded 存储把优化器状态、模型权重、随机数种子一起固化下来才能支撑“万亿参数”级别的断点续训。这块是很多人忽略的但真正训练时卡住三天比慢 5% 可怕多了。4. 实操记录拿 OLMo-core 3 跑一个 MoE 微调4.1 环境怎么搭先解决一个“跑得起来”的问题实操之前先说清楚要复现“8→128 只慢 5%”你需要的是多机多卡集群和足够带宽的网络比如 8×H100 起跳。但如果你本地只有一两张卡也不用灰心框架支持把小规模 MoE 配成“all 专家都在本地”的模式用来跑通代码路径和调试数据 pipeline和集群模式之间的差异只在通信层。先把代码拿到手git clone https://github.com/ai2cm/olmo-core.git cd olmo-core pip install -e .依赖上主要是 PyTorch、flash-attention、triton、transformers 这类常见库。如果装 flash-attention 遇到编译问题先换成预编译 wheel再不行就临时降到 attention 的纯 PyTorch 实现本地调试完全够用。额外提一句环境里尽量固定 CUDA 版本和 PyTorch 版本OLMo-core 对 flash-attention 版本比较敏感我曾经因为 PyTorch 小版本升级导致通信算子和 attention 算子出现非法内存访问最后只能重装环境。4.2 配置一个 64 专家的 MoE 微调任务配一个 MoE 模型核心是 model config。以 OLMo-core 3 风格的配置为例结构大致是这样model: d_model: 4096 n_layers: 16 n_heads: 32 n_kv_heads: 8 mlp: n_experts: 64 top_k: 2 load_balance_coefficient: 0.01 z_loss_coefficient: 0.001 capacity_factor: 1.25 rope: theta: 10000 optimizer: name: adamw lr: 2e-5 weight_decay: 0.1 data: dataset: your_sft_dataset sequence_packing: true trainer: type: sft micro_batch_size: 2 gradient_accumulation_steps: 8这里每个字段都有讲究。top_k 控制每个 token 激活几个专家2 是成本和效果比较折中的值load_balance_coefficient 是前面说的专家负载均衡损失权重太高会让路由趋近随机模型学不到 specializationcapacity_factor 表示每个专家最多接收 1.25 倍平均 token 数超过的 token 会丢弃。第一次跑建议把四个 MoE 相关参数全部打出来观察 loss 和路由统计再做微调。训练脚本可以用框架自带的 CLI也可以直接 import wrapperfrom olmo_core.train import SFTTrainerWrapper trainer SFTTrainerWrapper.from_yaml(moe_sft.yaml) trainer.train()这段是示意代码实际 API 会随版本迭代但思路就是这样yaml 管配置wrapper 管训练循环数据流和并行布局都在配置层声明不写分布式样板代码。调试时我把 micro_batch_size 调到 1关闭 gradient checkpointing先把 pipeline 跑通再逐步加上大 batch 和 checkpoint 以提升吞吐。4.3 先小模型复现再放大几乎是我唯一的建议用 OLMo-core 3 跑 MoE最容易踩的坑不是框架不会用而是“一上来就想复现 128 专家”。我习惯的路径是先配一个 8 专家的小模型几十亿参数单机 8 卡内跑通确认路由分布正常、loss 曲线平滑再切到 64 或 128 专家观察通信耗时占比和专家利用率曲线。这个过程看起来慢实际是在给后续的万亿参数训练排雷。框架再强也无法替你解决数据质量、loss 震荡、路由坍塌这些小问题。如果你手里的 GPU 有限还有一个很实用的姿势用 LoRA 微调冻结大部分专家。MoE 的专家非常多冻结后显存占用主要集中在路由和 attention 上。在 wrapper 里把 freeze_experts 设为 true优化器只更新 adapter 和 router这样单卡 24GB 也能跑一个中等规模的 MoE 指令微调。代价是路由可能只在冻结特化专家的情况下学习真实场景效果会有折扣用来做方向验证是够了。5. 常见问题与避坑实录5.1 路由“坍塌”后loss 还在降但模型不干活这不是崩溃而是退化。表现是训练早期 loss 正常下降但到某个 step 之后所有 token 都路由到同一个专家其他专家彻底闲置。查特征很直接训练日志里多看一层“routed_tokens_per_expert”如果某个专家的占比超过 50%基本就是路由坍塌。处理方法按顺序试先检查 load_balance_coefficient 是否设成了 0 或太小尝试提到 0.05再看 top_k 和 capacity_factor 是否匹配容量因子太大会让少数专家吃下所有 token最后检查 router 权重初始化和学习率如果 router 更新太快初期几个 step 就能把分布推偏。重新初始化 router 层权重是个土办法但在长训练后期很有效。5.2 128 个专家跑起来了可某些专家从来没被选中这是“死专家”问题常见于 MoE 预训练后期或微调阶段。诊断方法给每个专家记一个累加计数训练几千 step 后画柱状图如果大段专家计数为 0说明路由已经被某些强势专家垄断。死专家会白白占着显存也意味着模型容量实际缩水。缓解手段一是调整 load_balance_coefficient但别加太高否则专家会失去分工二是用 expert dropout 或者 batch 内随机路由此类的扰动让冷门专家偶尔能被 token 经过三是检查数据如果你微调用的数据集中在某几个领域路由自然倾向相关专家这个问题往往不是工程 bug而是数据 bias。处理好数据分布比硬调损失项更实在。5.3 扩专家池后通信占比直线飙升不像 5%如果实际跑出来和基准差距很大大概率是你的集群网络拓扑和框架默认假设不一致。OLMo-core 3 的 all-to-all 通信通量设计是建立在节点内高速互联、节点间带宽充足的前提下。如果你用的是云厂商普通实例跨节点带宽可能只有几百 Mb/s那 8→128 的墙钟时间自然控制不住。排查路径先开一个纯通信 benchmark测跨节点 all-to-all 带宽再看通信消息的 size bucket如果很多小于几百 KB 的小包说明 message 合并没生效调整 do_all_to_all_bucketing 之类的设置最后确认打开 overlap 选项让通信和计算同时执行。注意观察训练日志里 comm 和 compute 分别的耗时占比别只盯着总时间。5.4 loss 在大规模 MoE 训练初期出现奇怪的“平移”这种情况最像玄学loss 数值本身正常下降但曲线阶段性出现几乎垂直的跳跃甚至在同一个 epoch 内 loss 上升到一个平台再慢慢降下来。追查下来大多是以下两个原因之一要么是 sequence packing 导致某些 batch 里长序列比例过高要么是数据采样没有做全局随机化。MoE 对数据分布很敏感一个 batch 里如果某个领域的 token 扎堆router 会被整批带走下个 batch 又换一批loss 就出现周期性毛刺。解法在数据 pipeline 里做 shuffle 强度和采样权重的全局统一如果使用多数据源混配给每个源设置稳定的占比。OLMo-core 3 的数据配置里通常有 sampler 项把多源比例固定住再跑比反复调调度器省事多了。另有一个容易忽略的点序列打包时尽量让长度相近的样本进同一个包避免一个 sequence 跨多个专家 batch 的边界这个细节对 loss 稳定性影响很大。最后说点实在的OLMo-core 3 给我的感觉是Ai2 把“扩专家池不该线性变慢”这件事做成了一套可复现的工程方案。它没有发明魔法只是把稀疏激活、负载均衡、通信重叠这些原理按部就班地落实到了训练框架里。我实际跑过的教训是不要看到“只慢 5%”就盲目把专家数调大专家多了模型容量确实上去但路由要吃更多数据才能学得稳显存和调参复杂度也会跟着涨。如果你第一次用这个框架我建议先把 8 专家版本完整跑完再看专家利用率曲线和 loss曲线确认健康再往 128 甚至更多专家上冲。这样你的第一次“万亿参数”体验才不会变成第一次“OOM 调试”体验。
