最近大模型圈子里出现了一篇画风相当清奇的论文Mixture-of-Kittens简称 MoK。别看名字软萌得像猫咖促销它讨论的其实是 MoEMixture of Experts在大规模训练里最硬核的那个问题——在 NVL72 这套 72 卡液冷机架级系统上怎么把密集计算和稀疏计算的配比调好才能让超大参数模型在固定算力预算下练得更快、效果更稳。我第一遍读的时候以为只是一个架构小 trick读完才发现它把预训练语言模型的 scaling law、训练环境选型、增量训练这些环节全都串起来了值得专门拆一篇解读。MoK 全称是 Mixture-of-Kittens: Striking the Right Balance Between Dense and Sparse Computing核心就一件事给 MoE 的专家层加一个“变量共享”开关。一个模型里既有所有 token 都会经过的共享专家dense 路径又有按 token 内容路由的专家sparse 路径两者比例可以连续调节。这个开关直接解开了我一直没太想明白的问题MoE 明明有十几倍参数训练和推理计算量为什么还能压得比同尺寸 dense 模型低因为真正算的只有 top-k 个专家。那如果把一部分专家变成人人可用的共享专家是不是还有更多操作空间MoK 就是把这个直觉系统化并且用 NVL72 上的大规模实验证明了一遍。1. 论文到底在解决什么问题1.1 MoE 看似省算力实际优化空间很大标准 MoE 的 FFN 层被拆成 E 个专家每个 token 由 router 选择 top-k 个专家完成计算。总参数量可以做得很大但每个 token 激活的参数很小所以业界常用 7B 激活配 8 专家的结构跑出 40B 以上的总参数在相同推理成本下获得接近更大 dense 模型的能力。但真正训过 MoE 的人都会有类似体感它并没有想象中高效。所有专家对任意 token 都是“稀疏被选”没有一个稳定的共享基底。结果就是不同 token 走不同专家各专家的梯度噪声相对独立训练早期波动明显路由模块和专家模块强耦合容易出现专家退化部分专家学到的东西高度重合所有计算都是 sparse训练 FLOPs 几乎只由激活参数量决定缺少一个灵活的旋钮来匹配不同 batch size、数据配比下的最优算力分配。换句话说MoE 把“参数量”和“计算量”解耦了但这把解耦只做了一半。它把参数量做大却把计算量固定在 top-k 机制上没有考虑“我应该让多少参数被所有人共享”这个维度。MoK 补上的正是这一半。1.2 dense 和 sparse 不是二选一而是一条连续轴很多人的直觉是把模型分成两类dense 模型每个 token 走全部参数MoE 模型每个 token 走一小部分参数。MoK 的做法把这层二值判断打碎变成一条连续轴。它把 FFN 层分成若干组每组单独决定是共享还是路由。共享组内所有 token 都经过该组全部参数不参与 router 选择路由组内维持标准 top-k 机制。于是共享比例 0 时退化成标准 MoE共享比例 1 时退化成 dense MLP中间值就变成“一部分人做公共基础能力一部分人做专门化处理”的混合团队。生活化类比就是公司架构dense 模型是全员参与所有项目MoE 是每个项目只调几个精兵MoK 则是在精兵之外保留一支稳定的基础支持团队所有人都能拿到基础服务再由路由小组处理高难特化需求。这种结构的价值在于基础功能被反复使用、持续打磨专门能力则保持低计算成本和高弹性。2. MoK 技术原理拆解2.1 变量共享机制到底怎么工作MoK 的“变量共享”不是简单的 parameter sharing也不是 weight tying而是将专家分组后按组定义计算路径。论文用共享比例这个超参数决定某个组是 dense 还是 sparse。实现层面可以这样理解一个 MoK 层包含 G 组专家每组专家内部有若干个 expert。对于组 g如果它的共享标记为 1则组内所有 expert 被作为一个整体参与计算所有 token 都完整走一遍该组的 MLP如果共享标记为 0则组内复用标准 router每个 token 只激活组内 top-k 个 expert。于是总参数量 P 可以分为两部分共享参数数量 P_shared αP所有 token 都会利用路由参数数量 P_routed (1-α)P每个 token 只利用其中 top_k / E 的比例。单 token 的实际激活参数量可以用公式估算P_active αP (1-α)P × top_k / E当 α0 时就是纯 MoE 的激活量P × top_k / E。当 α0 时激活量会有小幅提升但因为共享部分承担了大量通用计算路由组的 top_k 往往可以显著调小论文实验里甚至出现 top_k1 依然表现良好的情况。综合下来MoK 可以在同 FLOPs 下塞入更大总参数量 P模型记忆容量更大计算开销却不变。2.2 共享比例、路由比例和计算复杂度的联动训练 MoK 模型的 FLOPs 不是简单看激活参数量要拆成共享路径和路由路径两段算共享路径 FLOPs 2 × tokens × P_shared路由路径 FLOPs 2 × tokens × P_routed × top_k / E总 FLOPs 是两者之和。这个公式带来的工程含义很直接如果想在固定 FLOPs 预算下提升模型效果可以调大 P用部分共享参数换模型容量如果想在固定参数量下减少计算成本可以提高共享比例压缩路由需求。MoK 论文花费大量篇幅讨论的 scaling law本质上就是在算这个三维空间的最优解。一个很容易踩的误区是共享比例越高计算量越高所以“共享就是省算力”是错误的。共享只省通信和路由 overhead不省 FLOPs。真正的作用是利用固定参数量去换取更合理的计算分配让通用特征提取被重复利用而不是每个 token 都去触发一组全新的专业参数。2.3 Router 在 MoK 里的职责收缩标准 MoE 的 router 需要为每个 token 在所有 E 个专家上计算 logits再取 top-k。MoK 中共享组不需要 gaterouter 只需要覆盖路由组内的专家。这使得gate 参数量下降路由计算开销减少每 token 的路由目标更集中负载不均衡的压力得到缓解。MoK 论文仍然保留 load-balancing auxiliary loss但对共享组不施加任何约束。实际训练中路由组的 load 均衡相对传统 MoE 更容易保持因为共享路径已经吸收了大部分高频共性输入路由组只需处理后验分叉明显的样本。3. 在 NVL72 上的训练实验规模与方法3.1 NVL72 系统不是普通训练环境NVL72 是 NVIDIA GB200 NVL72 机架级解决方案一个机架内包含 72 个 Blackwell GPU配合 Grace CPU 和 NVLink5/NVSwitch 互联。它的特点是把过去需要跨节点、跨交换机才能完成的 all-to-all 通信压缩到机架内单一高速域里完成。对 MoE 类模型来说token dispatch 的通信开销往往会吃掉大量训练效率NVL72 的设计正是冲着这类通信密集任务去的。论文选择在 NVL72 上训练不只是因为 GPU 多更关键的是可以用一个机架完成其他集群需要几十台机器才能承载的通信拓扑从而把 MoK 在“减少 all-to-all 通信”上的优势放大得很明显。共享路径不需要走远程专家调度token 的跨设备移动量比标准 MoE 低很多这在 NVL72 这种全互联环境下体现得尤其充分。我在自己经验里也有类似体感MoE 小规模实验时通信开销不明显一上多机多卡all-to-all 就成了瓶颈。MoK 这种“部分共享 部分路由”的设计天然减少需要远程搬运的 token 比例等于从架构层面给通信减压而不是单纯靠换网络设备。3.2 关键实验设计从 10B 到百亿参数的验证路径论文的实验路径很标准分三步走1B 级小规模消融验证共享比例、top-k 取值、路由组分组方式等超参10B 级全面对比与同计算预算的 dense 模型和标准 MoE 对齐比较 loss 和下游任务分扩展到大总参数、大规模 tokens 的预训练验证真实生产环境下的收益。在 10B 对比中论文报告的结论是固定训练 FLOPs 下MoK 普遍优于标准 MoEMoK-10B 的激活参数少于 MoE-12B但平均下游表现更稳部分数据集上还反超。这里的意义很实际——过去我们说“MoE 用更少计算换更好效果”MoK 把这个 advantage 做得更彻底因为它允许你用同样算力去换取更大总参数量。论文还做了大量关于训练后期行为的分析。MoK 的 loss curve 在高参数量区间下降更平滑这和我看到的一些 MoE 训练现象吻合标准 MoE 中不同专家在训练后期容易出现 specialization 过猛导致的震荡MoK 因为有一大块稳定的共享参数兜底模型容量被分摊得更均匀早期梯度方差也小一些。3.3 训练数据规模与 Tokens 的影响为了让结论具备说服力论文在数据规模上没有刻意降低标准。预训练数据量达到 1.057T tokens训练对象是总参数量 130B 级别的 MoK 巨型模型。这个规模的实验如果没有 NVL72 这类机架级系统光是通信稳定性和故障恢复就够团队折腾几周。对做增量训练或者持续预训练的团队来说MoK 还有一个隐藏收益因为共享参数承载了大部分通用知识增量训练阶段可以把共享部分冻结或降低学习率只微调路由组参数。这有点像给模型打了个补丁而不是整体重训。论文虽然没有细讲这个玩法但按 MoK 的结构推导这是很自然的后续方向。4. 从论文到工程怎么把 MoK 落进自己的代码4.1 Megatron 框架下的 MoK 层实现思路读论文是一回事落地是另一回事。MoK 层的实现并不复杂关键是把专家分组和路径选择逻辑改清楚。我梳理一下在 Megatron-LM 这类框架里实现 MoK 的核心步骤。class MoKExpertGroup(nn.Module): def __init__(self, hidden_size, experts_per_group, sharedFalse): super().__init__() self.shared shared self.experts nn.ModuleList([ MLP(hidden_size, hidden_size * 4, hidden_size) for _ in range(experts_per_group) ]) def forward(self, x): if self.shared: out sum(expert(x) for expert in self.experts) return out else: logits router(x) topk_idx torch.topk(logits, ktop_k, dim-1).indices out dispatch_and_compute(x, topk_idx, self.experts) return out几个工程要点共享组不应该走 MoE 的 all-to-all 调度它的权重在所有数据并行 rank 上都应当完整保留训练时只做张量并行或者不做并行避免每个 token 还要跨节点取参数路由组继续用 expert parallel把专家切到不同 rank维持 MoE 的扩展方式共享组的输出和路由组的输出最后要 scale 校正。论文中通常先按共享组数量和路由组激活量的统计比例来做归一化否则共享路径的输出量级容易被路由路径淹没Load balancing loss 只需要作用在路由组共享组不参与。实际开发时建议先用一个 FFN 层替代小规模 dense 模型做 smoke test验证共享路径的梯度方向是否正常再逐步扩展到真正的 MoK 层。4.2 超参数选择共享比例与 top-k 的搭配MoK 最核心的超参组合是共享比例 α 和路由组的 top_k。在我的实验经验里有几个相对稳定的甜点区间供参考10B 级模型共享比例 0.2~0.4 是比较安全的地带。低于 0.2基本退回 MoE共享路径起不到稳定基底的作用高于 0.4计算量上涨明显除非刻意扩总参数否则收益会递减路由组 top_k 建议先设 1。因为共享路径已经覆盖了通用特征路由组只需要挑一个最匹配的专家做特化效果往往比 top_k2 更干净计算量和通信又更省学习率要比同规模 MoE 略低。共享路径让梯度更新更密集相同学习率下的有效 step 步长会被放大我一般会乘 0.7~0.85 做补偿共享组和路由组的初始化标准差可以分开设。共享组承担共性的统计平滑任务初始方差略小有助于稳定性路由组保持常规 MoE 的初始化量级。4.3 训练稳定性和显存优化MoK 训练中的常见稳定性问题和标准 MoE 有区别需要单列共享路径的梯度累积更密集因为每个 token 都会贡献所以共享参数的有效梯度尺度天然比路由参数大。如果出现梯度爆炸优先对共享组单独做 clip而不是一棍子全部 clipActivation memory 会增加因为共享路径的中间激活必须保留到反向传播。可以用 activation checkpointing 把共享组的 FFN 中间结果重算实测显存能省 30% 左右混合精度训练下共享路径的 FP8 量化误差会累积得比路由路径快因为它是高频使用的部分。建议对共享组保留 BF16 权重只在路由组做 FP8 加速或者用更高精度做共享组的主权重更新。5. 常见问题与排错实录5.1 共享专家和 MoE 里的 shared expert 有什么差别很多人会问MoE 里不也经常加一个 shared expert 吗MoK 和它有什么本质差别区别在定位。标准 MoE 加 shared expert通常只加一个目的是让共享 expert 承担一部分通用 token 的处理缓解路由压力但它的参数量占比很小模型主体依然是纯稀疏。MoK 直接把一部分参数显式声明为共享计算路径共享比例可以是 20%、30% 甚至 50%而且这个共享不是“专家”概念下的一个单元是一整组 FFN 计算路径。换句话说标准 MoE 的 shared expert 是补丁MoK 的共享路径是架构的主干设计之一。两者在效果上的差异在大规模训练中会非常明显因为 MoK 改变了 FLOPs 的分配结构而 shared expert 只是让 MoE 的计算分配多了一个选项。5.2 路由不均衡问题是否还存在在 MoK 中路由组的负载均衡压力比标准 MoE 小但并没有完全消失。我在实际训练中遇到过路由组内 2~3 个专家被高频触发、其余专家几乎闲置的情况。这说明共享路径吸收的是共性特征路由组内部仍然存在专家 specialization 不均的问题。对策还是老三样但需要调整顺序先加 load balancing loss该项系数从 0.01 开始调如果仍然不均衡再用 expert dropout 或 token drop 作为软约束最后才考虑调整 top_k。MoK 的优点是路由组专家数量可以比标准 MoE 少因为共享路径已经承担了大部分轻量预测这样均衡压力天然小于同规模 MoE。5.3 训练 loss 突然变 NaN 怎么查MoK 训练的 NaN 问题大部分不是模型结构问题而是共享路径的数值溢出。因为共享路径的权重参与所有 token 的计算其梯度统计方差大在混合精度下特别容易出现中间激活超出 FP16/BF16 表示范围。排查顺序建议先冻结路由组、只看共享组前向确认是否存在 scale 异常检查共享组的 FFN 激活值如果常见范围已经超过 6 万就该考虑在共享组前加 RMSNorm 或者降低 FFN 中间层维度倍数用 BF16 跑一遍全模型如果用 BF16 就没问题只在 FP8 下爆那基本上是量化 scale 的问题需要给共享组单独设量化范围。我个人的偏好是共享组永远用 BF16 主权重不做 FP8 激进量化。这个看似保守的选择能省掉至少一半的 NaN 排查时间。5.4 显存和带宽不够的团队能复现吗MoK 论文在 NVL72 上跑但方法本身并不依赖特定硬件。共享路径让 token 的跨设备移动变少对小集群反而更友好。如果你的团队只有 8 卡 A100可以先用 1B 级 MoK 做验证把共享比例调到 0.5观察 loss 是否在相同 FLOPs 下优于 MoE。共享路径不需要 all-to-all通信压力比 MoE 低这是资源受限团队更容易接受的特性。6. 我读完这篇论文之后的几点体会MoK 最打动我的不是它刷了多高的 score而是它把“模型参数容量”和“单 token 实际计算成本”这两个维度彻底分开后又重新用共享比例把它们连接起来。过去我们习惯了要么全 dense、要么全 sparseMoK 提醒我们中间路径才是工程上最常见、也最容易被优化的状态。我目前正在自己的小规模预训练框架里复现 MoK计划先在 1B 级验证共享比例为 0.3、top_k1 的配置然后往 3B 级迁移。如果训练稳定性和下游效果都符合预期后续的增量训练阶段我也想尝试冻结共享路径、只微调路由组的方案这有望把持续预训练的成本再压低一截。最后说一个小技巧如果你只是想快速感知 MoK 在自家数据上的收益不用一开始就重写框架。把现有 MoE 模型里参数量最大的那组 experts 直接改造成共享路径跑一次 post-train 微调观察 loss 变化就能得到一个相当可靠的预判。改造的成本很低但信息量很大。Mixture-of-Kittens 这个名字看着像玩笑背后其实是实打实的计算分配哲学。大规模训练不是参数量游戏而是“让每一份 FLOPs 都花在刀刃上”的游戏。MoK 给了我们一个更好用的旋钮接下来就看工程社区怎么把它转起来了。