MoK技术解析：密集与稀疏计算平衡如何优化大规模训练-平芜编程栈

最近大模型圈子里出现了一篇画风相当清奇的论文：Mixture-of-Kittens，简称 MoK。别看名字软萌得像猫咖促销，它讨论的其实是 MoE（Mixture 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_k=1 依然表现良好的情况。综合下来，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 中，共享组不需要 gate，router 只需要覆盖路由组内的专家。这使得：

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 普遍优于标准 MoE；MoK-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, shared=False): 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, k=top_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_k=2 更干净，计算量和通信又更省；
学习率要比同规模 MoE 略低。共享路径让梯度更新更密集，相同学习率下的有效 step 步长会被放大，我一般会乘 0.7~0.85 做补偿；
共享组和路由组的初始化标准差可以分开设。共享组承担共性的统计平滑任务，初始方差略小有助于稳定性；路由组保持常规 MoE 的初始化量级。

4.3 训练稳定性和显存优化

MoK 训练中的常见稳定性问题和标准 MoE 有区别，需要单列：

共享路径的梯度累积更密集，因为每个 token 都会贡献，所以共享参数的有效梯度尺度天然比路由参数大。如果出现梯度爆炸，优先对共享组单独做 clip，而不是一棍子全部 clip；
Activation 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 低，这是资源受限团队更容易接受的特性。