利益相关：参与过Kimi-K2的接生，自己的孩子怎么看都顺眼。

自从Kimi K2发布以来，很高兴得到了开源社区大量的关注。注意到尽管我们的模型结构近乎完全继承了DeepSeek-V3（下文简称 DSv3），依然有很多小伙伴深入探究两个模型仅存的一点“不同”背后的原因。作为Moonshot Infra侧推理小透明一名，今天想从推理角度来简单讲一下Kimi K2的config为什么“长成”现在的样子。

提前叠甲：内容涉及到很多训练相关的内容，里面会掺杂一些个人理解，如有不准确的地方请其他同事纠正。

1. K2模型结构的设计宗旨。

在启动K2训练之前，我们进行了大量模型结构相关的scaling实验，结果是，所有当时propose的、与DSv3不同的结构，没有一个能真正打败他的（顶多旗鼓相当）。因此，问题就变成了，我们要不要为了与DeepSeek不同，强行选择一个没有优势但不一样的结构，最终的答案是no。原因也很简单：DSv3的结构是经过验证，在large scale上依然有效的，而我们的“新结构”还并没有经历过足够大规模的验证。在已经有muon优化器和更大参数量两个巨大变量的前提下，我们并不想引入没有明确收益的额外变量来“标新立异”。于是，就有了第一个约束条件：完全继承DSv3的结构，调整适合我们的模型结构参数。

而第二个约束条件就是成本，包含训练成本和推理成本。原因很简单，作为一家小公司，我们的训练和推理资源都是非常有限的。在DSv3推出之后，我们经过认真评估，认为它的训练成本和推理成本，都比较接近我们当前能承受的上限。因此，我们需要将K2的训练和推理成本，尽量控制在与（我们自己训推）DSv3持平的水平。

综上所述，模型结构的设计问题就变成了： 在给定DSv3结构的框架之下，如何选择合适的参数，使得模型在训练、推理成本与DSv3相当的前提下，获得明显更低的loss。其中训练成本方面本文不会详细展开（才不会承认因为我也是一知半解），我们会在我们之后发布的tech report中介绍K2的训练方案与优化，敬请期待：）

2. 具体改动和动机

正如很多人对比两份config文件所观察到的，我们在模型结构参数上，具体的改动主要包含： (1) expert 数量 (2) attention head 数 (3) 前面的 dense 层数只有 1 层 (4) 无分组的简化版 router。接下来我会按这个顺序，从模型推理的角度介绍它们背后的考量。本次推理方案完全沿用 DeepSeek 的 tech report [1] 和 OpenDay [2] 中提到的 EP+DP 方案，理论分析暂不含通信（假设通信可被推理层面尽可能overlap掉）。

2.1 num_experts = 384

这条结论来自 pretrain 团队的 sparsity scaling law，是 K2 项目从 pretrain 阶段开始的主要驱动力之一（另一个当然是 MuonClip）。简而言之，我们验证了在固定 # activate params 不变的前提下，单纯增长 MoE 总参数量，scaling law 依然成立，且不论训练 loss 还是验证 loss，结论始终保持，也就是无需担心增大总参数量会过拟合。因此，num_experts=384 承担了降低模型 loss 的核心任务。

对推理的影响：

• prefill 阶段：如果 prefill 节点数能做到和 num_experts=256 时一样大，且 prefill 的 seqlen 足够长，耗时基本无明显增长，因为此时的prefill是compute bound 的任务，我们的激活参数量不变，MoE 环节的总 FLOPS 也不变。
  
• decode 阶段：由于需考虑线上实际的 TBOT 指标，我们无法无限增大推理时的 batch size（虽然现在已经被狂喷慢慢慢了orz）。因此可以粗略地认为 MoE 阶段的GEMM仍是 memory bound，那么 MoE 参数量增大到 1.5×，相关计算的耗时也就变成了 1.5×。以 EP=128 为例（128 是我们的384 和DSv3的256的最大公约数，方便比较），对 DSv3 来说，每个 EP rank 上会存放 2 个路由专家和 1 个共享专家，大约 7.5 GB 的 MLP weights（不计算 EPLB方案 的冗余专家）；而 K2 则需要大约 10 G，大了2.5G。
  

2.2 num_attention_heads = 64

MoE 阶段实打实地变慢（贵）了，我们就得考虑能否从其他环节找补回来，attention 的 head 数是第一个想到的切入点。原因在于 MLA 的论文中，DeepSeek 为了让 MLA 尽可能充分利用访存带宽，相比 MHA 常见 attention heads ≈ layer 数的设计，把 head 数翻倍，也确实带来轻微涨点，但也带来两个问题：prefill 和 decode 实际上都变贵了。相比之下，如果我们将attention head数量重新变回64：

• prefill 阶段：
  (1) MLA attention 计算量为 2hs²(d_nope + d_rope + dv)，h 是 head 数，s 是序列长度，三个d分别为128,64,128。而整个模型其他部分基本都是矩阵乘法，计算量的公式为 2Ns，其中 N 是所有矩阵乘法相关参数的参数量。注意到，attention计算量与seqlen成平方关系，而其余矩阵乘法则为线性关系，即随输入序列长度增大，attention 在 prefill 阶段耗时占比越大。而K2 的目标场景（agent、vibe coding）中，长序列是标准的使用形态，正好被这个问题直戳痛点。而将 head 数砍半，则可以一定程度上削减这个快速增长的平方项对整体耗时造成的影响。
  (2) 除此之外，attention 前后还有 QKVO projection，这几个矩阵乘的参数量与 head 数线性相关。大家应该也能看到 DSv3 的激活参数 37 B，而K2 只有 32 B，差的 5 B 就来自这里。粗略来看，DSv3 激活的 37 B 中，QKVO projection 占 10 B，K2 只有 5 B，随着参数量的减少，这几个projection在prefill阶段的FLOPS消耗也随之减少，K2再胜一局。
  
• decode 阶段：
  attention core 的计算耗时主要取决于 KV cache 大小，这一点 K2 和 DSv3 完全相同，平手。但 QKVO projection部分 与 prefill阶段 类似，实打实地把 10 GB 的访存量降到了 5 GB。更关键的是，在 DP Attention 下，QKVO projection 会在所有 rank 上 replicate，因此这 5 GB 的差距不会像TP那样，随并行度增大而摊薄。因此不管EP size多大，每个rank都有5GB的仿存减少。回顾前面 MoE 的差距，EP128 下我们总参数量增大到 1TB，每个 rank 才多了 2.5 GB 访存，而这里 head 数从 128 降到 64，就能省下 5 GB，瞬间感觉自己很赚。
  

综上，降低 head 数可以瞬间把 MoE 参数增大亏掉的部分全部补回来，还有富余。我们最担心的只剩下这样对模型效果是否有明显的负影响。算法同学通过充分对比实验，确认了把 head 数还原到接近层数的“baseline”设置对 loss 的负影响要远远小于 MoE 参数增大的正影响，于是，num_heads=64 就这么愉快地决定了。（留一道思考题：减少 attention head 数，还可以为 Speculative Decoding 留下了更多提速空间。）

2.3 first_k_dense = 1

与 DeepSeek 的观察类似，我们也同样在训练中发现第一层 MoE 的 router 很难做到负载均衡，但不同的是第二层之后并没有发现什么大问题。为了更充分利用 MoE 优势，我们只保留第一层 dense，其余全用 MoE。这个操作对 prefill 几乎无影响，对 decode 每个 rank 大约增加几百 MB 访存，可以忽略不计。

2.4 n_group = 1（expert 无分组）

expert 分组的最大价值是当一个 rank 上存在多个 expert 时，可以让它们同组，在 device(GPU) level 让 MoE 计算更均衡。但在当前模型的参数规模下，我们不得不使用很大的 EP，每个 device 上只剩少量、甚至只有一个 expert，group level 的均衡则从GPU层面转换到了节点层面。而即便节点层面能够做到相对均衡，但每个节点内部遇到所有 token 都被 route 到当前 group 的同一个 expert上这种最坏情况下，MoE 计算耗时仍然不会理想。因此，EPLB方案里面的动态重排和冗余专家对于当前设定下的负载均衡问题相对来说要更为关键一些。而更自由的 router 方案能让 expert 的组合空间显著增大，从而进一步提升模型能力。

3. 小结

以上就是K2模型结构参数被设定为当前这个状态，来自推理侧的完整思维链了。综合以上四个相比DSv3的改动，我们能够得到一个在相同EP数量下，虽然总参数增大到1.5倍，但除去通信部分，理论的prefill和decode耗时都更小的推理方案。即使考虑与通信overlap等复杂因素，这个方案也不会比DSv3有显著的成本增加。可以自豪的说，虽然只有小小的4个参数的改动，但每一个决策的背后都有充足的理论分析和实验验证。也希望模型开源后，能有更多的推理厂商和框架共同帮我们验证前面分析的正确性。再次感谢所有小伙伴对Kimi-K2的关注！

[1] DeepSeek-V3 Technical Report https://arxiv.org/pdf/2412.19437 [2] One More Thing, DeepSeek-V3/R1 Inference System Overview. https://github.com/deepseek-ai/open-infra-index/blob/main/202502OpenSourceWeek/day_6_one_more_thing_deepseekV3R1_inference_system_overview.md

【补充】

前面的估计确实有个重要前提被遗漏了，就是GPU的型号。所有关于计算仿存的roofline都和实际具体每张卡的计算密度有关。文中的估计都是根据我们手头的现有卡来的。具体多少batch_size会打到compute这种问题还需要大家根据自己手头GPU的实际情况来重新估算，无法直接作为结论抄作业。

架构方面，可能会让一些想看“新东西”的读者失望，因为我们大体上是复用了DeepSeek-V3的架构，仅在个别地方有所改动，下图比较清晰地展示了改动之处：

但这里想要说的是，不管是跟DeepSeek-V3的“异”还是“同”，背后其实都有很多思考和实验在支撑着我们的选择。

异

首来说说“异”。MLA部分，我们没有像DeepSeek-V3那样double heads，

然后是MoE部分，我们将num_experts从256升到384，这部分有两个原因，一是为了补回没有double heads带来的损失，二是符合我们测出来的Sparsity的Scaling Law。后者简单来说，就是在一定范围内，固定激活的专家数量，增加总专家数带来的效果收益，大于增加的Infra成本。

还有就是num_experts=384，刚好凑出了1T总参数量这个吉祥的数字，代表着我们开始迈进 1t.ai（可以点击）时代！其实很多时候就是这些简单的美感，支撑我们渡过枯燥的实验过程。

此外，还有一些细微的改动，比如我们将Routed Experts的Scaling Factor从2.5改为了2.827，这个改动很小，大体上不会带来实质变化，但它是我们自己从理论出发推导出来的结果（参考这里），它遵循了我们一贯的原则——每一处继承或者调整，都需要足够的证据，不管是理论的还是实验的。

同

接着再来说说“同”。首先还是MLA，为什么要选择MLA，笔者已经用两篇文章《Transformer升级之路：21、MLA好在哪里?（上）》和《Transformer升级之路：21、MLA好在哪里?（下）》介绍了MLA的优越性，简单来说就是MLA真的很能打。

当然，实际情况还复杂得多，我们内部也一直致力于探索MLA更好的替代品，但是作为我们第一个开源的超大模型，我们还是选择了致敬DeepSeek，复用它的MLA设计。还有MoE部分，我们沿用了DeepSeek-V3的Shared Expert、大Sparsity和Loss-Free负载均衡，这些我们都是经过自己的实验，确认它能要不能提高MoE的训练稳定性，要不能提高MoE的学习效率，最终才确定使用的。

这里特别提一下Sparsity部分，在相当长一段时间内，我们都是在Sparsity=8（64选8）下做实验，直到有一天我们解决了Infra问题后，开始尝试更大的Sparsity，发现收益非常可观，于是开始探索Sparsity的Scaling Law，并且感叹似乎一切都在DeepSeek的预料中（DeepSeek-V3已经是256选8，32的Sparsity；K2则是384选8，48的Sparsity）。

所以，我们最终决定，K2还是先狠狠地致敬DeepSeek了，而不是刻意地标新立异。更多的想法和创意，我们把它们做得更扎实后，在K3、K4再见了。

优化

优化方面，很多读者都应该注意到，我们已经转移到Muon的大本营，然而Muon在超大规模上的实践完全是空白的，所以从Moonlight到K2，实际上也是一场Muon的Scaleup之路上的“大开荒”。

器

几个月前的Moonlight中，我们已经在16B的MoE上验证了Muon，确认了Weight Decay的必要性和Match Adam Update RMS的有效性。但单凭这两点，还不足以将Muon扩展到千亿甚至万亿参数的模型，因为新的拦路虎出现了——MaxLogit爆炸。

对此，我们的解决方法是MuonClip，它是Muon与QK-Clip的组合，技术细节大家可以看文章《QK-Clip：让Muon在Scaleup之路上更进一步》，这里我们尽量多聊一点文章以外的内容。

$$\text{MuonClip}\quad\left\{\quad\begin{aligned} &\boldsymbol{M}_t = \mu \boldsymbol{M}_{t−1} + \boldsymbol{G}_t \\[8pt] &\boldsymbol{O}_t = \mathop{\text{msign}}(\boldsymbol{M}_t) \underbrace{\times \sqrt{\max(n,m)}\times 0.2}_{\text{Match Adam Update RMS}} \\[8pt] &\boldsymbol{W}_t = \boldsymbol{W}_{t−1} − \eta_t (\boldsymbol{O}_t + \lambda \boldsymbol{W}_{t-1}) \\[8pt] &\left.\begin{aligned} &\text{if }S_{\max}^{(l,h)} > \tau: \\ &\qquad\text{if }\boldsymbol{W} \in \{\boldsymbol{W}_{qc}^{(l,h)}, \boldsymbol{W}_{kc}^{(l,h)}\}: \\ &\qquad\qquad\boldsymbol{W}_t \leftarrow \boldsymbol{W}_t \times \sqrt{\tau / S_{\max}^{(l,h)}} \\ &\qquad\text{elif }\boldsymbol{W} \in \{\boldsymbol{W}_{qr}^{(l,h)}\}: \\ &\qquad\qquad\boldsymbol{W}_t \leftarrow \boldsymbol{W}_t \times \tau / S_{\max}^{(l,h)} \end{aligned}\quad\right\} \text{QK-Clip} \end{aligned}\right.\\$$ \text{MuonClip}\quad\left\{\quad\begin{aligned} &\boldsymbol{M}_t = \mu \boldsymbol{M}_{t−1} + \boldsymbol{G}_t \\[8pt] &\boldsymbol{O}_t = \mathop{\text{msign}}(\boldsymbol{M}_t) \underbrace{\times \sqrt{\max(n,m)}\times 0.2}_{\text{Match Adam Update RMS}} \\[8pt] &\boldsymbol{W}_t = \boldsymbol{W}_{t−1} − \eta_t (\boldsymbol{O}_t + \lambda \boldsymbol{W}_{t-1}) \\[8pt] &\left.\begin{aligned} &\text{if }S_{\max}^{(l,h)} > \tau: \\ &\qquad\text{if }\boldsymbol{W} \in \{\boldsymbol{W}_{qc}^{(l,h)}, \boldsymbol{W}_{kc}^{(l,h)}\}: \\ &\qquad\qquad\boldsymbol{W}_t \leftarrow \boldsymbol{W}_t \times \sqrt{\tau / S_{\max}^{(l,h)}} \\ &\qquad\text{elif }\boldsymbol{W} \in \{\boldsymbol{W}_{qr}^{(l,h)}\}: \\ &\qquad\qquad\boldsymbol{W}_t \leftarrow \boldsymbol{W}_t \times \tau / S_{\max}^{(l,h)} \end{aligned}\quad\right\} \text{QK-Clip} \end{aligned}\right.\\

QK-Clip是应对MaxLogit爆炸非常直接有效的方法，而且从事后来看，它的思想也非常简单，但它的出现并没有那么显然易得，甚至可以说非常“突然”。

思

让我们回到K2启动训练的前几天，尽管大家都觉得K2的训练已经迫在眉睫，然而我们仍没找到能杜绝MaxLogit爆炸的方法，我们已经在这个问题上前前后后试了几个月，所有的方法在参数量进一步变大后均以失败告终。我们甚至已经打算尝试“Q、K的weight单独用回Adam优化器”这样的丑陋思路（因为Adam似乎没那么容易炸）。

另一边，笔者隐隐觉得参数裁剪特别是奇异值裁剪能够解决这个问题，但奇异值裁剪并不直观，多数同学不好接受，而且实现复杂度也略高，所以依然是待定状态。在一筹莫展之际，笔者突然醒悟般意识到“MaxLogit本身就是触发Clip的最本质信号”，然后就顺理成章提出了QK-Clip。这个方法相比奇异值裁剪直观太多了，大家很快都理解了它，并开始测试，终于K2的最后一环被打通，K2训练正式启动！期间还经历了从Total Clip到Per-Head Clip的变化，但这都是小插曲了。

QK-Clip的出现，对笔者以及整个团队的冲击都蛮大的。就MaxLogit这个问题而言，我相信各位读者随便都能提出很多可能的解决方案，然而事实就是前几个月我们严重困在了“可能”两个字中，因为对于一个问题，有太多“可能”的解决办法了，没有直接思考“保证能”解决问题的办法，是我们前几个月困境的本质所在。

所以，QK-Clip不仅解决了MaxLogit爆炸这个具体的问题，还给了我们很大启发，在面对一个类似的难题时，我们应该至少要问自己三个问题：

• 这个方法能保证解决这个问题吗？（我们此前的一系列无效尝试）
• 如果不能，有什么方法能保证解决这个问题吗？（QK-Clip）
• 这是解决这个问题的最小改动的方法吗？（Per-Head的QK-Clip）

另外，我们认为QK-Clip的思想，实际上是解决很多训练不稳定问题的“抗生素”，它可以推广成“哪里不稳Clip哪里”，只要我们的监控指标出现了异常，那么就可以考虑构建类似的Clip。这也许不是最优雅的解决方法，但肯定是一个非常直接有效的解决方法。

最后

最后的最后，K2的成功训练，是很多同学共同努力的结果，本文只谈及到了自己相关的非常小的部分，还有很多精彩的内容没有参与和谈及，敬请大家期待我们的Tech Report。

更新：已接入claude code，Agent能力确实牛！

这次开源诚意十足，而且看到擅长Agentic任务这里，我真心动了，高低得试试。

模型代号：kimi-k2-0711-preview，128K上下文长度，MoE架构，1T参数量，激活参数32B。暂不支持Vision，亮点是代码能力和agent能力。
K2定位为具备“通用Agent”能力的大模型，原生支持多工具调用和联网功能（太好了，国内终于有专门为AI Agent优化的模型了）。

目前kimi官网可以免费直接用：

而且模型已经完全开源，可以本地部署：

再看看分数，SWE仅次于Claude，数学能力强，Tool Use和Claude接近。在SWE-Bench等编程基准测试中取得了开源模型的SOTA(比Claude还是差点)，在数学竞赛题（如AIME）、逻辑谜题等高难度推理任务上，K2同样超越了主流开源模型的水平。尤其在需要多步推导的题目中，善于处理多步骤指令和复杂任务链。

官方博客里面还贴心给了一个案例（大家可以从这里看：https://moonshotai.github.io/Kimi-K2/）

但昨天才吃了Grok 4刷榜的亏，今天我高低得试试难的！

接下来我这里把Kimi k2接入Gemini Cli，看看真实的编程+Agentic能力。

未完待续，正在测试，一会更新。

更新：Gemini Cli接入有点复杂，发现Claude Code有个ANTHROPIC_BASE_URL可以直接换base url，我用这种方式把Kimi K2接入进去了。

工具调用丝滑（出现过一次工具调用格式不对），代码能力也不弱。考虑到Claude Code是Anthropic根据Claude定制的prompt，直接套在Kimi K2上能用，还是很惊讶了（之前试过好几个国内的模型，Kimi K2在这里是独一档）。

所以我们终于可以搞到低成本Claude Code了，开心（这个大夜没白熬啊）

未完待续（还在跑...我的kimi api账号新开的，中间请求频繁报429了）。


更新2：kimi倾向于给很长的todo list，然后中间常常会某次工具调用失败，导致无法一次性全部完成，需要人工矫正（这个还是明显不如Claude 4，当然claude code的prompt本就是为claude 4设计的有点不公平。且k2的工具调用稳定性已经明显超过了deepseek v3）
尝试了3次，均在某步出错，有一次甚至给我claude code干废了，原因不明：

00:04

claude code + kimi k2，某次出现的奇怪bug

这是1个提示词做出来的（中间失败过1次，然后错误日志贴个它，一次性修好了），成品：

00:09

claude code + kimi k2，一次性完成的文字rpg游戏

完成度不错了，有基本的移动/地图/属性/log/背包逻辑。


我后续打算针对k2这个模型，优化下claude code的prompt，然后继续试试。

（对了，由于claude code走了本地的代理服务，里面的prompt在日志全都暴露出来了。不是啥大事情，但还挺奇妙的，只能说这种本地安装的Agent的安全还是个大问题...）

未完待续。

1. 使用Muon优化器的Concerns

值得注意的是，moonshot发布的K2是端到端使用Muon ^[2][3] 训练的。Muon最早是Keller提出，在Speedrun中表现出色^[2] ，然后被moonshot盯上，进行了一些调整并scale up ^[3][4] 。

在moonshot早期的工作中^[3]，尽管moonshot在其工作中claim Muon有巨大的token efficiency优势并给出了一个Megatron-LM实现 ^[5]，然而Muon在后续相关工作和X上的讨论还是引发了一些concern，主要包括三个方面：

1. Muon需要full parameter matrix来进行NS计算。在modern LLM training infra的parallel setting下，很多人认为这个full parameter matrix的操作太expensive
2. Muon需要“几套”不同的调参机制，对模型的调参提出了更高的要求。对比下，看起来AdamW这种self-adaptive的优化器调参更简单稳定
3. Muon可能引发训练的不稳定性，比如Moonshot自己的work^[3] 其实就有提到，有潜在的attention max logit问题

事实上，随着K2^[6]的发布，一定程度上可以理解这些问题都不是大问题。这篇blog尝试“狡辩”一下原因。

2. Concern 1: Muon's Infra Scalability

最先讨论Muon full parameter matrix的操作是否expensive，弥补一下前期论文^[3]里没有把这里说清楚的一个小遗憾。为了说清楚，需要仔细介绍一下Zero 1 Sharding，然后理解一下Zero 1 Sharding的实现方式和国内外（为啥说国外？因为主要是国外的一些公司在X上进行challenge）training cluster的区别（本质就是国外公司太有钱了卡太多了），才能够详细说清楚为啥Moonshot的人觉得Muon's infra is scalable，但还是有人再challenge。

2.1 Zero 1 Sharding

首先需要介绍一下背景，因为modern LLM training往往是large model + large cluster，Zero1 optimzier实际上是训练中标配使用技术，无论是Megatron-LM，Deepspeed，FSDP都已经实现支持zero1。

Zero 1技术本质上就是把需要吃掉很多显存的optimizer states，比如AdamW的exp, exp_square, fp32_master_weights，在DP group之间进行sharding。在使用AdamW时，具体的Zero 1 Distributed Optimizer的lifecyle是：

1. 先进行DP间grads的reduce scatter。因为有sharding，所以是reduce scatter而不是all reduce，每个DP rank只需要自己负责维护的那一部分local params的gradients是准的就可以了
2. 再进行local params的AdamW update计算。这里AdamW的计算本质上是elementwise，所以也只需要计算local params的updates即可
3. 最后进行DP间的params-allgather。这里因为每个DP只负责一部分params的更新，其他的DP想要拿到准确的更新完的参数，需要再进行一次allgather

注意1和3看着通信很expensive，但是实际上可以和model fwdbwd进行overlap（很成熟的技术，所有框架都会做），所以其实不用担心。而step2，因为AdamW是elementwise计算且DP越大，每一个DP rank的计算量越小，所以其实也很scalable。综上来看，就会发现这套Distributed Optimizer技术对于AdamW来说非常友好。事实上，AdamW optimizer time的开销通常在整个global step的1%以下，和fwdbwd一比基本忽略不计。

然而Muon在step 2里有比较大的challenge，因为不再是elementwise的计算。在这里Muon requires full parameter matrix来进行NS的计算，这就势必要引入额外的通信和更大的计算量（整个矩阵跑NS）。如果想要Muon尽量scalable，也就是需要step2的通信开销尽量小（因为几乎不可隐藏），引入的额外的计算尽量小（单个小矩阵跑NS很快，所以尽量避免每个DP rank需要在太大的矩阵和太多矩阵上进行NS）

2.2 Moonshot解决方案

基于Moonshot的开源工作^[5]，这里推测Moonshot是基于Megatron-LM的某个版本进行开发并且之后一直进行维护。对于Megatron-LM而言，他早期的Zero 1 optimizer ^[7]实现如下（我们称之为flat-param concat zero1）：

可以看到，他的做法是把所有的optimizer states全部flatten起来，然后concat在一起，然后按DP平均分配切。这样的分配方式是显存最优的，因为没有任何optimizer state的duplicates。而且，这样切对于Muon是十分有利，因为大多数的local params都是完备的，直接可以进行NS运算。只有concat在DP边缘的params会导致被切在两个DP上而不完整，需要进行一些操作。

具体而言，以DP0和DP1共同处理Param 1为例，如果我们在这里开脑洞雕花的话，有几种搞法：

• 无脑gather法。DP 0和DP 1分别进行gather，两个人都拿全params，两个rank都进行full NS，算完之后只更新自己local param的那一部分，剩下的丢掉。step1 和step3的grad reduce scatter和params all gather保持不变
• 避免重算法。每个DP i把自己边缘的params往DP i-1传，让DP i-1负责算边缘params的计算，算完之后再拿过来在当前rank维护的那一部分进行更新。这样就避免了任何额外的计算，通信量实际上也比无脑gather法要好。不过对于一些极端情况，如params1打穿了3个dp rank，就需要一些额外的heuristic
• 精密排布法。排布distributed optimizer的时候，直接不让DP边缘切割行为发生，这样没有任何额外通信和额外计算。代价是显存分配不再balance，且最优分配变成一个背包问题。不balance的显存分配，显然infra的同学是受不了的，因为会导致估不准训练期间的显存，进而影响并行分配策略。

事实上，moonshot采用的是无脑gather法，因为写起来最简单，且不用应对任何特殊情况。注意虽然听起来无脑，但是引入的额外计算其实不多！！！因为真正发生在DP边缘的参数并不多（DP * 2个），理论上来说只有这部分的参数需要额外通信以及duplicate的NS计算。其他的params比如图上的param0，param2，因为本来就是完整的，并不需要任何额外通信和计算！

可以看到empirically，这些通信和计算的实际性能都会受DP数以及模型里的最大矩阵大小的影响。考虑到modern MoE arch（感谢deepseek v3），实际上一个模型里不会有特别大的矩阵，因为都是fine-grained experts（word embedding和lm head是AdamW控制的，不需要Muon）。因此，长远来看，Muon的scalability前途一片光明，稳中向好。事实上因为无脑法的开销本就不大，因此雕别的方案的收益也不大，进而蒋kernel也没啥继续雕的动力了（不过印象里YouJiacheng游老师应该在speedrun上雕过一些类似的方案？）。

2.3 Others' concern

然而实际上，在一些外国公司的研究里，Muon scalability是偏向唱衰^{[8][9][10][11]} ，并且moonshot的方法^[5]被反复拷打。显然不可能别人都是傻逼，但是基于2.2的分析，以及实际上moonshot把K2 large scale训出来了，所以moonshot应该也不是傻逼...

我个人感觉这里造成conflicts的主要原因是因为两边对于Zero 1实现的不同，进而导致了对于Step 2这个额外通信量和计算量的估计产生了很大的偏差。国外的主流方法称之为dim 0 sharding Zero1。比如 国外的主流并行框架PyTorch FSDP2实际上的Zero 1实现方式是这样的 ^[12]：

而更新版本的Megatron-LM ^[13]里引入了bucket的概念。这个概念究其本质和params dim 0 sharding造成的影响应该是类似的：

这些更新实际上对moonshot之前的Muon实现是一个“毁灭性”打击，因为这种Zero 1实现，实际上让每一个参数都被DP进行了sharding！那无论是无脑gather法，避免重算法，还是精妙排布法，这些基于flat params contcat zero1全毁了，因为每一个参数都需要通信和额外重算，进而导致了海量的额外通信量和计算量，Muon也就变得不可接受了。

2.4 长期解决方案

国外很多公司肯定不是傻逼，早期的并行设计实际上都是使用flat params concat zero1 ^[14]。后期为了一些别的concern（主要是国外公司卡太多了，flat params不利于做grad reduce scatter和params all gather的overlap），而改成了dim0 params sharding Zero 1。在必须进行dim0 params sharding的context下，moonshot的方法确实不scalable。但这不表示Muon本质上没法变得scalable，这里必然会有新的解决方案出现。事实上听说好像似乎大概可能也行已经有人在做了（手动狗头）

3. Concern 2: Muon needs more hyperparameters

另外一个常见complain就是说Muon会有好几套hyperparameters，相比AdamW是一个显著劣势：

• 需要额外的调参efforts。需要额外的调参表示需要额外的心智负担才能找到最好的模型
• 不是对AdamW的fair comparisons。如果对AdamW也进行分块调参，可能会有更好的效果

我个人认为这种concern实际上是对Muon优化器的数学性质没有准确的理解。要理解Muon，实际上要从Standard Parametrization (SP)和Maximal Parametrization (MuP)两个角度进行理解，才能理解为啥会有多套参数需要进行调整。此外，Muon是为了matrix设计的^[2]，因此非matrix参数word embedding，lm head， rmsnorm gamma都是用AdamW进行更新的。

3.1 Standard Parametrization + Muon

先来看SP下的Muon，在moonshot早期（24年12月左右）开始研究/复现（抄）Keller的Muon时^[15]，它当时长这样（没有weight decay，没有游老师加入的各种类似于zero1的各种工程优化）：

可以看到这个时候没有那么令人发指的那么多套参数，就是AdamW一套，Muon一套。这里其实Muon的update RMS是和AdamW差特别多的，在moonshot的工作中^[3]，苏隐给出了推导：

可以看到，AdamW的update rms empirically 是0.2~0.4左右，而Muon是远远要小的。如果不把Muon的update rms调大（最简单的方式就是搞一套专用的lr），Muon是更新不动的，因此反而是不公平的比较方式。

在SP setting下，不想搞两套调参的话，可以直接使用moonshot的工作^[3]，通过matching update rms的方式，开箱即用，直接使用一套AdamW的超参即可。AdamW的超参咋调的工作不要太多（e.g. stepfun law），随便找一个抄过来，然后用moonshot的方式迁移到Muon大概率都有不错的improved loss token efficiency。事实上，moonshot的工作的主要贡献就是这里，让大家在SP setting下可以无脑迁移到Muon。

我个人浅显的理解，这里其实等于是在矩阵frobenius norm constrain下的最速优化，实际上控制了update的rms，和AdamW类似。符合SP的要求，但不是最优的。对于Muon，理论上最优的方式是接下来要讨论的spectral norm constrain下的最速优化。

3.2 MuP Parametrization + Muon

Muon最精彩的用法不是SP，而是与MuP结合。有一系列开源的工作做了非常精彩的介绍！^[16][17][18]。简单来说，Muon几乎是为MuP量身定做的优化器，不像使用MuP + AdamW时那样引入很多假设和基于variance的假设，Muon天然控制spectral norm（因为NS数学上对max singular values进行了裁剪，而max singular values by definition is spectral norm），进而非常适合高阶MuP^[17]需要的sepctral norm控制！

查看Keller对Muon的改进历史，除了游老师等高手进行infra上的优化，主要就是男神Jeremy对MuP思想的引入（Jeremy是MuP的作者之一，也是Muon blog的作者之一，可以说double男神了）。在引入MuP的类似思想后，对Embedding，LM Head，Hidden Matrix都会有各自进行控制的逻辑^[19]，虽然令人发指，但是细究起来是为了适配MuP是合理的（事实上AdamW里适配MuP也要对不同的模块做lr adjustment）。特别是这里对Muon update rms的调整，不看那个max(1, x)部分，只看里面的 $\sqrt{d_{out}/d_{in}}$\sqrt{d_{out}/d_{in}} 部分，这就是苏隐高阶MuP里的推导^[17]一模一样！（不过不知道为啥要加这个max(1,x)的操作，加了max之后实际上变回frobenius norm了？）

4. Concern 3: Muon Training Instabilities

实际上large scale训练Muon的公司没几个，Muon训练的instabilities实际上是moonshot自抛自扣^[3][6]。这里的instabilities主要是weight decay和max attention logit问题。weight decay好理解，max attention logit问题涉及到了这次blog里说的muonclip方法^[6]。

Max attention logit问题通常可以使用qknorm来解决，但是Moonshot在K2里使用了MLA（不得不说deepseek太狠了，model arch基本一试一个不吱声），估计是测起来效果确实好，也没必要强行创新而不使用这种好的技术。MLA在comppression上加了norm，但是为了inference效率，实际上q和k heads不会被materialize出来，这样也就没有办法做qk head norm。因此moonshot也就另辟蹊径，搞了muonclip（实际上也有一些别的人对qknorm有效果上的concern^[20]）

我个人觉得muonclip是非常优雅的！在苏隐的高阶MuP blog中 ^[17]，我们可以学到spectral norm比frobenius norm小：

而sepctral norm直接和max logit大小相关，即 $||xW||_2 <= ||x||_2 ||W||_2$||xW||_2 <= ||x||_2 ||W||_2 相关（这里W是矩阵，所以是spectral norm），所以最直接的搞法就是控制spectral norm。但是spectral norm比较难算，因此直接利用spectral norm和frobenius norm的不等式关系，直接在frobenius norm上做clip，那 $||xW||_2$||xW||_2 就被控制了！

不过后来找机会和苏隐聊了一下，苏隐表示没有想这么多，而且我的理解不一定对（人麻了），他的思路就是直接对根本问题进行操作。苏隐说之后几天会有blog放出来，到时候可以关注一波。

5. Conclusion

感觉K2应该会是一个很强力的模型，期待更多相关人士的测评。此外moonshot之前非常擅长VL和RL，可以期待经过一段时间后训练之后，基于K2的thinking & vision understanding模型有机会大放异彩！

同时moonshot作为一家拥有苏隐，蒋kernel，飞来阁打手张宇等诸多高手的公司，感觉很有前途！此外moonshot既能把fancy的新技术例如Muon落地，也能大方承认和使用竞对开发的先进技术，感觉格局还是很大的！

参考

1. ^飞来阁的故事： https://zhuanlan.zhihu.com/p/1915601328211759191
2. ^^abcKeller Muon Blog:  https://kellerjordan.github.io/posts/muon/
3. ^^abcdefghMoonshot Muon:  https://arxiv.org/abs/2502.16982
4. ^Why using Muon:  https://kexue.fm/archives/10739
5. ^^abcMegatron-LM PR:  https://github.com/NVIDIA/Megatron-LM/pull/1428
6. ^^abcMoonshot K2: https://github.com/MoonshotAI/Kimi-K2
7. ^Megatron-LM Zero 1:  https://github.com/NVIDIA/Megatron-LM/blob/main/docs/source/images/distrib_optimizer/sharding_scheme.png
8. ^Keller fighting on X:  https://x.com/kellerjordan0/status/1893868235381961140
9. ^Essential AI拷打Moonshot： https://www.essential.ai/blog/infra
10. ^Dion拷打Moonshot： https://arxiv.org/pdf/2504.05295
11. ^韩国小哥Seo拷打Moonshot： https://x.com/SeunghyunSEO7/status/1943731232119964027
12. ^Pytorch FSDP2 Sharding： https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md
13. ^Megatron-LM bucket： https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/distributed/distributed_data_parallel.py#L58
14. ^PyTorch FSDP1 Flat Params： https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_flat_param.py
15. ^Keller早期的Muon实现： https://github.com/KellerJordan/Muon/blob/7f9342f50bb418d14a52ec89449e7bc93bebca95/muon.py
16. ^Jeremry's Deriving Muon： https://jeremybernste.in/writing/deriving-muon
17. ^^abcd高阶MuP推导： https://kexue.fm/archives/10795
18. ^Some people amazed by Muon + MuP on X： https://x.com/JingyuanLiu123/status/1931223767449309657
19. ^keller最新muon实现： https://github.com/KellerJordan/Muon/blob/master/muon.py#L157
20. ^qknorm is a band-aid:  https://x.com/giffmana/status/1943731151497027962

DeepSeek彻底改变格局后，Kimi的转向其实挺明确，而且之前Kimi 一系列工作都在指向这个模型。

几个关键点：（1）DeepSeek后它的撞车论文，预示它转向开源预训练，我们当时就说它最终是要和DeepSeek撞更大的车的。（2）muon，非常关键的工作。（3）最近的agent产品，背后开发团队分享时，一个点最吸引人：“agent就是模型”。

和Kimi这个团队的几次接触，感觉还是有一口气，要争一争。这次的效果就更加关键，接下来看看后面几天社区的发酵和评价了。

引用一段我在2月的硅星人的文章里写的话：

月之暗面的应对思路也就变得引人注目。其中开源是必须要做的一步。而看起来月之暗面的选择，是想要真正匹配DeepSeek的开源思路——目前在DeepSeek后出现的诸多开源大部分像是应激反应，它们还是用此前Llama时期的开源思路来跟随。事实上DeepSeek的开源与此前的已经不同，不再是Llama似的防守扰乱闭源对手式的开源，而是一种能带来明确收益的竞争策略。
月之暗面最近传出内部“把SOTA（state-of-the-art）结果做为目标”，看起来是最接近这种新的开源模式的策略，要开的是最强的模型，最强的架构方法，这样反而会得到它一直渴望的在应用侧的影响力。
根据两家的论文，MoBA已经用在月之暗面的模型和产品上，NSA也一样，它甚至让外界可以对DeepSeek接下来的模型有更明确预期。于是接下来的看点就是，月之暗面和DeepSeek用MoBA和 NSA各自训练出的下一代模型，要不要再撞车一次，并且还是以开源的方式——这可能也是月之暗面在等待着的那个节点。

另外，希望下次可以中国时间白天发。像DeepSeek那个开源周那样。让硅谷彻夜无眠去吧。

用了一下kimi的anthropic接口桥接的claude code，感觉很棒，希望kimi把anthropic的模型内搜索也做了