首页/Kimi 发布首个万亿参数开源模型 K2 模型，哪些信息值得关注？/toothacher17/
7个回答
分享
查看原文

Kimi 发布首个万亿参数开源模型 K2 模型，哪些信息值得关注？

holiday111
96个点赞 👍

利益相关：不请自来，曾经moonshot第一蒋吹（可能要和

@Andrew Lu
 竞争），飞来阁[1]
@苏剑林
 粉丝，趁K2发布来硬蹭一波。

1. 使用Muon优化器的Concerns

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

Moonshot论文中的Muon具有巨大的优势

在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）：

Megatron-LM里的早期Zero 1 Distributed Optimizer实现

可以看到，他的做法是把所有的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]：

FSDP V2 Sharding on Params Dim 0

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

新版本Megatron-LM的DDP里引入了bucket

这些更新实际上对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的各种工程优化）：

Keller Jordan的早期版本Muon

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

苏隐的Update RMS推导

可以看到，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}} 部分，这就是苏隐高阶MuP里的推导[17]一模一样！（不过不知道为啥要加这个max(1,x)的操作，加了max之后实际上变回frobenius norm了？）

Keller的muon update对hidden矩阵进行adjust LR的方法


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]。

Moonshot自抛自扣

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]）

moonshot's muonclip

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

spectral norm比frobenius norm小

而sepctral norm直接和max logit大小相关，即 ||xW||_2 <= ||x||_2 ||W||_2 相关（这里W是矩阵，所以是spectral norm），所以最直接的搞法就是控制spectral norm。但是spectral norm比较难算，因此直接利用spectral norm和frobenius norm的不等式关系，直接在frobenius norm上做clip，那 ||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
toothacher17
自由评论 (0)
分享
Copyright © 2022 GreatFire.org