Kimi K2：MuonClip 与万亿参数 Agent 模型

交互

Kimi K2 是 Moonshot AI 在 2025 年 7 月发布的开源 MoE 模型：总参数 1.04T，激活 32.6B，用 15.5T token 预训练，全程没有一次 loss spike。它真正新的东西有三块：

1. 全貌：K2 有哪些数字

K2 的骨架和 DeepSeek-V3 几乎一样，差别集中在几个数字上：

项目DeepSeek-V3Kimi K2
层数6161
总参数671B1.04T
激活参数37B32.6B
专家总数（routed）256384
每 token 激活专家88
共享专家11
注意力头数12864
Dense 层数31
专家分组（expert grouping）有无

另外几个常用数字：隐藏维度 7168，每个专家的 FFN 中间维度 2048，注意力用 MLA；预训练 15.5T token、上下文 4K，之后退火并做长上下文激活，用 YaRN 扩到 128K；优化器 MuonClip，QK-Clip 阈值 \(\tau=100\)。

MLA（见第 7 节）的具体尺寸是：q_lora_rank=1536、kv_lora_rank=512、qk_nope_head_dim=128、qk_rope_head_dim=64、v_head_dim=128，路由打分函数是 sigmoid（scoring_func: sigmoid，topk_method: noaux_tc，即 DeepSeek-V3 式的无辅助损失负载均衡）。

2. 从 Adam 到 Muon：正交化动量

2.1 先回忆 SGD 动量和 Adam

训练就是反复做 \(W \leftarrow W - \eta\,\Delta\)。不同优化器的区别只在“更新方向 \(\Delta\) 怎么从梯度 \(G\) 算出来”：

这两者都把参数当成一堆独立的标量。但神经网络里绝大多数参数是矩阵 \(W\in\mathbb R^{n\times m}\)，它的作用是一个线性算子 \(x\mapsto Wx\)。对矩阵来说，更自然的问题是：这一步更新会把输出 \(Wx\) 改变多少？

2.2 谱范数下的最速下降

Muon 背后的观点（Bernstein & Newhouse 的“在范数约束下做最速下降”）是：对矩阵参数，应该约束更新的谱范数（最大奇异值）\(\|\Delta\|_2\)，因为它正好是“对任意单位输入，输出最多变化多少”。于是问题是

$$\Delta^\star = \arg\max_{\|\Delta\|_2 \le 1} \langle G, \Delta\rangle .$$

设 \(G = U\Sigma V^\top\) 是奇异值分解。\(\langle G,\Delta\rangle = \operatorname{tr}(\Sigma\, U^\top \Delta V)\)，而 \(U^\top\Delta V\) 的对角元绝对值都不超过它的谱范数 1，所以最大值 \(\sum_i \sigma_i\) 在 \(U^\top \Delta V = I\) 时取到：

$$\Delta^\star = U V^\top =: \operatorname{msign}(G).$$

这就是“正交化”：把梯度的所有奇异值都换成 1，只保留它的奇异向量（方向）。它常被记作 \(\operatorname{msign}\)（matrix sign），类比标量的 sign：Adam 在逐元素意义下取符号，Muon 在矩阵意义下取符号。等价的写法是 \(\operatorname{msign}(G) = (GG^\top)^{-1/2}G\)。

Muon 的完整更新是：

$$M_t = \mu M_{t-1} + G_t,\qquad O_t = \operatorname{msign}(M_t),\qquad W_t = W_{t-1} - \eta\, O_t .$$

动量先在原始梯度上累积，再整体正交化。动量系数通常取 \(\mu=0.95\)，Moonlight 和 K2 都用这个值。

2.3 Newton–Schulz 迭代：不做 SVD 也能正交化

每步对每个矩阵做 SVD 太贵，而且不适合 GPU。Muon 用一个只含矩阵乘法的迭代逼近 \(UV^\top\)。先归一化 \(X_0 = M/\|M\|_F\)，保证所有奇异值都在 \((0,1]\)；然后迭代

$$X_{k} = a X_{k-1} + b\,(X_{k-1}X_{k-1}^\top) X_{k-1} + c\,(X_{k-1}X_{k-1}^\top)^2 X_{k-1}.$$

关键观察：如果 \(X = U\Sigma V^\top\)，那么 \((XX^\top)X = U\Sigma^3V^\top\)，\((XX^\top)^2X = U\Sigma^5 V^\top\)。所以这个迭代不改变奇异向量，只对每个奇异值独立地施加奇多项式

$$\sigma \mapsto f(\sigma) = a\sigma + b\sigma^3 + c\sigma^5 .$$

只要 \(f\) 把 \((0,1]\) 里的值都推向 1 附近，迭代几次后 \(X_k \approx U V^\top\)。Muon 用的系数是 \(a=3.4445,\ b=-4.7750,\ c=2.0315\)。这里的设计取舍是：\(f'(0)=a\approx 3.44\) 很大，小奇异值每步被放大三倍多，几步就能拉上来；代价是 \(f\) 并不精确收敛到 1，而是让奇异值在 1 附近振荡。按这组系数，初始奇异值只要不小于 0.01，迭代 5 步后都落在大约 \([0.68,\ 1.13]\) 里。对优化来说这已经够了：Moonlight 的实验里迭代 10 步比 5 步正交化得更准，但训练效果没有更好，所以用 5 步。

开销：对 \(n\times m\)（\(n\le m\)）的矩阵，一步迭代大约是 \(XX^\top\)（\(2n^2m\)）、平方（\(2n^3\)）和两次乘回 \(X\)（\(4n^2 m\)）。5 步加起来比前向反向小得多；实测优化器带来的额外延迟一般只占前向反向的 1%–3%。

def newton_schulz5(G, steps=5, eps=1e-7):
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.bfloat16() / (G.norm() + eps)
    transposed = X.size(0) > X.size(1)
    if transposed: X = X.T          # 让 X X^T 是较小的那个方阵
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X
    return X.T if transposed else X
Newton–Schulz 的 5 步（系数 3.4445, −4.7750, 2.0315）。右边是一个 8 个奇异值按 0.6 的幂衰减的动量矩阵（示意），左边是它把单位圆在最大、最小两个奇异方向张成的平面里映成的椭圆。每一步都把小奇异值放大 3 倍多、把大的往回压，椭圆一步步变圆，奇异向量（椭圆的朝向）始终不变。

2.4 为什么正交化有帮助

直观解释：梯度矩阵往往是“低秩、病态”的——少数几个奇异方向占了绝大部分能量。SGD 和 Adam 的更新会主要沿这几个方向走，其它方向几乎不动。正交化之后所有奇异方向的步长相同，那些弱但有用的方向也能被学到。换句话说，正交化让更新矩阵“各向同性”，权重不会只沿少数主导方向学习。

Moonlight 给了两类证据：

把“原始梯度、Adam 式符号更新、Newton–Schulz 正交化”放在两种尺度上对比。上半部分是 2×2 矩阵作用在单位圆上的像：病态的梯度把圆压成细长椭圆，正交化后接近一个圆（奇异值都接近 1）。下半部分是一个 32×16 的随机矩阵，奇异值按幂律衰减，柱状图显示三种更新的奇异值（各自除以最大值）。所有数都在页面上实时计算：奇异值由 Jacobi 方法求 \(G^\top G\) 的特征值得到。

交互演示 · Muon 与 Adam 的更新谱

单位圆在三种更新下的像（2×2）

梯度 G/‖G‖Fsign(G)（Adam 极限）NS 正交化

迭代多项式 f 作用 k 次后的映射 σ → fk(σ)

32×16 矩阵的奇异值（各自归一化到最大值为 1，降序）

原始梯度sign(G)NS 正交化
2×2：NS 后的两个奇异值–
SVD 熵：G / sign / NS–
NS 后最小/最大奇异值–

矩阵为随机生成（示意）。sign(G) 的谱虽然也比 G 平，但它不保持 G 的奇异向量，且仍有明显的衰减；NS 只改奇异值、不改奇异向量。2×2 情形里 sign(G) 常常退化成一条线段：对接近秩 1 的矩阵 \(G\approx\sigma uv^\top\)，\(\operatorname{sign}(G_{ij})=\operatorname{sign}(u_i)\operatorname{sign}(v_j)\) 仍是秩 1，逐元素取符号救不回弱方向。

3. Moonlight 的两处修改：权重衰减与更新 RMS 对齐

原始 Muon 在小规模语言模型上很强，但扩到大模型、多 token 时收益会消失。Moonlight 找到两个原因并修掉，K2 原样继承了这两处修改。

3.1 加回权重衰减

原始 Muon 没有权重衰减。Moonlight 观察到训练久了权重和层输出的 RMS 一直增长，超出 bf16 的高精度区间。他们加入 AdamW 式的解耦权重衰减：

$$W_t = W_{t-1} - \eta_t\,(O_t + \lambda W_{t-1}).$$

在一个 800M 模型、100B token（约 5 倍算力最优 token 数）的过训练实验里，原始 Muon 早期收敛更快，但后期被带权重衰减的 Muon 反超，后者也优于 AdamW。为什么 Muon 特别需要它，可以从更新的“形状”理解：Muon 每步的更新是满秩、每个方向步长相同的，不像 Adam 那样大多数方向几乎不动，所以权重范数更容易持续增长；权重衰减提供一个平衡点。这个“奇异值容易增长”的性质，也正是注意力 logit 爆炸的根源。

3.2 更新 RMS 随矩阵形状变化

Adam 的更新逐元素大约是 \(\pm1\) 量级，所以它的更新 RMS（均方根）不依赖矩阵形状，实际在 0.2–0.4 左右。Muon 不是这样。引理：形状为 \([A,B]\) 的满秩矩阵，Muon 更新 \(O = UV^\top\) 的 RMS 是 \(\sqrt{1/\max(A,B)}\)。

证明很短。设 \(n=\max(A,B)\ge m=\min(A,B)\)，满秩时 \(O = \sum_{k=1}^{m} u_k v_k^\top\)，\(\{u_k\}\)、\(\{v_k\}\) 各自正交归一。则

$$\|O\|_F^2 = \operatorname{tr}(O^\top O) = \operatorname{tr}\Big(\sum_k v_k v_k^\top\Big) = m,\qquad \operatorname{RMS}(O) = \sqrt{\frac{\|O\|_F^2}{nm}} = \sqrt{\frac{1}{n}} .$$

这带来两个问题：大矩阵（比如 dense MLP，\(\max(A,B)\) 很大）的更新太小，学不动；小矩阵（比如把 MLA/GQA 的每个 KV 头当成独立参数）更新太大，训练不稳。修法是把每个矩阵的更新乘上 \(\sqrt{\max(A,B)}\) 抵消形状影响，再乘 0.2 对齐 AdamW 的典型更新 RMS：

$$W_t = W_{t-1} - \eta_t\big(0.2\cdot O_t\cdot\sqrt{\max(A,B)} + \lambda W_{t-1}\big).$$

这样做的另一个好处是 Muon 可以直接复用为 AdamW 调好的学习率和权重衰减——这对实践很重要，因为非矩阵参数（embedding、LM head、RMSNorm 的 gamma）仍然用 AdamW，两套优化器共享一组超参。Moonlight 做过一组扫描：把 Muon 更新 RMS 设成 0.05/0.1/0.2/0.4/0.8，2k 步时 0.2 与 0.4 最好且相近（训练 loss 3.198 / 3.199），AdamW 基线是 3.512。

Moonlight 还比较了三种控制更新 RMS 的方法：统一乘 \(0.2\sqrt{H}\)（基线）、把更新直接归一到 RMS=0.2（Update Norm）、按形状乘 \(0.2\sqrt{\max(A,B)}\)（Adjusted LR）。后两者都比基线好且彼此接近，Adjusted LR 计算更便宜，于是被采用；K2 的优化器里用的就是它：\(O_t = \text{Newton-Schulz}(M_t)\cdot\sqrt{\max(n,m)}\cdot 0.2\)。

3.3 分布式 Muon 与 SFT 阶段

把 Muon 用到大规模训练，还要解决分布式的问题。AdamW 的更新逐元素计算，所以 ZeRO-1 可以把优化器状态任意切片。Muon 需要完整的矩阵才能做 Newton–Schulz。Moonlight 的 Distributed Muon 在 ZeRO-1 基础上多两步：在数据并行组内把本地分片的（带动量的）梯度 gather 成完整矩阵，算完整更新，只保留自己那一片去更新参数。内存上 Muon 只有一个动量缓冲，是 AdamW 的一半；通信量是 Distributed AdamW 的 \((1, 1.25]\) 倍（多出来的 gather 用 bf16）。

Moonlight 还做了预训练与 SFT 优化器的交叉实验：Muon 预训练 + Muon SFT 最好；如果预训练用 AdamW，SFT 换成 Muon 没有明显优势。K2 据此在 SFT 和 RL 阶段也用 Muon，并建议用户微调 K2 时用 Muon。

4. 注意力 logit 爆炸

4.1 现象

Muon 在 K2 的规模上遇到了一个新的失稳模式。先定义要观察的量：对一个注意力头 \(h\)，\(Q^h = XW_q^h,\ K^h = XW_k^h\)，softmax 的输入是 \(\frac{1}{\sqrt d}Q^hK^{h\top}\)。每个头在一个 batch 里的最大 logit 是

$$S^h_{\max} = \frac{1}{\sqrt d}\max_{X\in B}\ \max_{i,j}\ Q^h_i K^{h\top}_j .$$

K2 团队在一个中等规模（激活 9B、总参 53B 的 MoE）用原始 Muon 训练时，看到最大 logit 很快超过 1000。logit 到这个量级，softmax 几乎变成 one-hot，梯度要么消失要么剧烈，通常会伴随 loss spike，偶尔发散。他们在实验里发现这个问题在 Muon 下比 AdamW 下更频繁。Moonlight 当时其实已经观察到了苗头：训练早期部分层的最大 logit 超过 100，只是 logit 超过 100 的比例只有约 \(10^{-4}\)，且后来自己降下去了，所以没出事。到 K2 的规模就不再安全。

4.2 为什么由权重的谱范数决定

logit 有一个简单的上界：

$$|q_i\cdot k_j| \le \|q_i\|\,\|k_j\| \le \|x_i\|\,\|x_j\|\,\|W_q\|_2\,\|W_k\|_2 .$$

注意力层的输入经过 RMSNorm，\(\|x\|\) 是有界的，所以 logit 能无限增长只能是 \(W_q\) 或 \(W_k\) 的谱范数在增长。更准确地说，logit 是双线性型 \(x_i^\top W_q W_k^\top x_j\)，乘积 \(W_qW_k^\top\) 把两个谱范数乘在一起，任何一个奇异值的增长都被“平方地”放大。

4.3 为什么 Muon 更容易出事

K2 团队对此给出的是一个假设而非证明：Muon 的更新所有奇异值相等、有效秩满；Adam 的更新谱偏斜、有效秩低。把权重写成 \(W_{t-1} = \sum_i\sigma_i u_iv_i^\top\)，更新写成 \(\Delta W_t = \sum_j\bar\sigma\,\bar u_j\bar v_j^\top\)。权重和更新的有效秩都更高时，某个更新方向 \(\bar u_j\bar v_j^\top\) 与权重已有的奇异方向 \(u_iv_i^\top\) 对齐的概率更大，一旦对齐，对应奇异值就被逐步加性地推高。再经过上面的“平方放大”，就变成 logit 爆炸。Moonlight 的 SVD 熵结果（Muon 训练的权重有效秩更高）被引用为支持证据。

4.4 已有办法为什么不够

5. QK-Clip 的推导与 MuonClip

5.1 核心想法：事后缩放权重

QK-Clip 不改前向、不改反向，而是在优化器更新之后，看这一步前向里观察到的 \(S^h_{\max}\)，如果超过阈值 \(\tau\)，就把这个头的 \(W_q\)、\(W_k\) 按比例缩小。最大 logit 只是一个“该压多狠”的信号。

推导只用到 logit 关于 \(W_q\)、\(W_k\) 是双线性的。若

$$W_q \leftarrow \gamma^{\alpha}W_q,\qquad W_k\leftarrow \gamma^{1-\alpha}W_k,$$

那么对同样的输入，

$$q_i\cdot k_j \ \to\ \gamma^{\alpha}\gamma^{1-\alpha}\,q_i\cdot k_j = \gamma\, q_i\cdot k_j .$$

取 \(\gamma = \min(1, \tau/S_{\max})\)，缩放后（在同一批数据上）最大 logit 恰好变成 \(\min(S_{\max},\tau)\)。\(\alpha\) 控制 query 与 key 各承担多少，通常取 0.5，两边各乘 \(\sqrt\gamma\)。下一步的数据不同、权重又更新了一步，所以这不是严格的硬上界，而是持续把权重拉回“在当前数据上 logit 不超过 \(\tau\)”的区域。

5.2 为什么要按头裁剪

朴素做法是全局用一个 \(\gamma=\min(1,\tau/\max_h S^h_{\max})\) 缩放所有头。但实际上只有少数头会爆炸，全局缩放会把正常的头也压下去，相当于凭空给它们加了强正则。所以 K2 对每个头单独算 \(\gamma_h = \min(1,\tau/S^h_{\max})\)，只动超标的头。演示里可以切换两种方式，看正常头被“误伤”多少。

5.3 MLA 里怎么裁

MLA 的每个头，query 和 key 都由两部分拼成：不带位置编码的“内容”部分（C）和带 RoPE 的部分（R）：

$$q^h = [\,q^{C,h};\ q^{R,h}\,],\qquad k^h = [\,k^{C,h};\ k^{R}\,],\qquad q^h\cdot k^h = q^{C,h}\cdot k^{C,h} + q^{R,h}\cdot k^{R}.$$

要注意 \(k^R\) 是所有头共享的（它由输入直接投影得到，再加 RoPE，只有一份 64 维）。如果缩放它，会同时影响所有头，违背“只动超标的头”。所以 K2 的方案是：

两项都乘了 \(\gamma_h\)，所以整个 logit 恰好缩放 \(\gamma_h\)，与 5.1 的推导一致。落实到参数上，就是缩放 query 上投影里该头对应的 \(W_{qc}^h\)、\(W_{qr}^h\)，以及 key 上投影（\(W^{UK}\)）里该头对应的 \(W_{kc}^h\)。这些都是头专属的权重切片，缩放互不影响。

5.4 MuonClip 算法

把 Muon、权重衰减、RMS 对齐和 QK-Clip 合在一起，就是 K2 的 MuonClip 优化器：

for each step t:
    # 1. Muon（含 Moonlight 的两处修改）
    for each matrix W (n×m):
        M_t = μ·M_{t-1} + G_t
        O_t = NewtonSchulz(M_t) · sqrt(max(n, m)) · 0.2      # 对齐 AdamW 的更新 RMS
        W_t = W_{t-1} − η·(O_t + λ·W_{t-1})                  # 解耦权重衰减
    # 2. QK-Clip（逐头）
    for each attention head h in every attention layer:
        S = S_max^h            # 前向时顺手记录，不需要额外计算
        if S > τ:
            γ = τ / S
            W_qc^h ← W_qc^h · sqrt(γ)
            W_kc^h ← W_kc^h · sqrt(γ)
            W_qr^h ← W_qr^h · γ

实现上唯一的额外开销是前向时为每个头记录一个 max（在 FlashAttention 类 kernel 里可以顺带算出），以及偶尔对少量权重切片做一次标量乘法。

x c^Q（1536） c^KV（512） q^C,h　×√γ_h q^R,h　×γ_h k^C,h　×√γ_h k^R（共享）不动 logit · γ_h= qC·kC + qR·kR
QK-Clip 在 MLA 中的作用位置：头专属的三个分量被缩放，共享的 RoPE key 不动，整个 logit 恰好乘 γh。括号里是 K2 的维度。
逐头 QK-Clip 的 5 个训练步（玩具数字）。每步优化器更新让权重谱范数变大，头 3 的最大 logit 越过 τ；更新之后只把这个头的 \(W_q\)、\(W_k\) 各乘 \(\sqrt{\gamma}\)，\(\gamma=\tau/S_{\max}\)，它的 logit 正好回到 τ，其他头不受影响。右边灰色虚线是不裁剪时的走势。

5.5 实验证据

一个玩具动力学就能复现这些现象：每个头的 \(W_q\)、\(W_k\) 谱范数按“对齐带来的乘性增长 − 权重衰减”演化，增长率随训练衰减（模拟“权重和更新方向对齐”在早期最强）；少数头增长率特别大。logit \(S^h = \kappa_h\sigma_q\sigma_k\)。打开 QK-Clip 时，每步更新后按 5.1 的规则缩放。

交互演示 · QK-Clip 模拟器（玩具模型）
无裁剪：所有头中的最大 logitQK-Clip：最大 logit阈值 τ
无裁剪的峰值 Smax–
触发过裁剪的头占比–
最后一次触发在–
正常头 logit 被额外压缩–

数据为示意。logit 的对数增长率 = 2·(rh·e−t/T − λ)，rh 对少数“爆炸头”取大值。“正常头被额外压缩”= 未爆炸头在裁剪方案下的平均 logit 相对无裁剪时下降的比例；逐头方案下它应为 0，全局方案下明显大于 0。

6. Token 效率与数据改写

6.1 token 效率是新的瓶颈

高质量的人类数据越来越有限，token 效率（每个 token 带来的性能提升）成了缩放中的关键系数。提升它有两条路：优化器（Muon 每个 token 学得更多）和数据（让每个 token 携带更多有效信号，即提高 token utility）。最朴素的“多学几遍”（多 epoch 重复）会过拟合、泛化变差。

6.2 知识类数据的改写

知识密集的文本有个两难：只过一遍学不全，重复多遍收益递减还容易过拟合。K2 的改写框架有三个部件：

块 1 块 2 块 3 改写(块1) 改写(块2 | 前文) 改写(块3 | 前文) 拼接 → 完整改写文+ 忠实度校验
分块自回归改写：每块在前文的条件下改写，最后拼回。

改写到底有没有用？K2 团队在一个早期 checkpoint 上用 SimpleQA 准确率比较了三种方案：

改写次数epoch 数SimpleQA 准确率
0（原始 wiki 文本）1023.76
11027.39
10128.94

同样是“看十遍”，十个不同的改写版本各看一遍最好。实际大规模使用时，每个语料最多改写两次。

6.3 数学数据

数学文档按 SwallowMath 的方法改写成“学习笔记”风格；另外把其他语言的高质量数学材料翻译成英文来增加多样性。合成数据作为持续缩放的手段还远没有成熟，难点是推广到更多领域时保证事实正确、减少幻觉和毒性、以及规模化。

6.4 总体语料

15.5T token，四个领域：网页、代码、数学、知识。大部分处理流程沿用 Kimi K1.5，每个领域都做了正确性、质量验证和针对性的数据实验。

7. 架构：MLA、384 专家与稀疏度缩放律

7.1 MLA 回顾（按 K2 的尺寸）

长上下文推理的主要显存开销是 KV cache。MLA（Multi-head Latent Attention，DeepSeek-V2 引入）的思路是把每个 token 的 K/V 压缩到一个低维隐向量，只缓存这个隐向量。K2 的尺寸如下：

按这些尺寸算，每层注意力参数约 1.01 亿，61 层合计约 6.2B。

7.2 MoE 配置

除第一层是 dense FFN（中间维 18432）外，其余 60 层都是 MoE：384 个路由专家 + 1 个共享专家，每 token 选 8 个路由专家，专家 FFN 中间维 2048（SwiGLU，每个专家 \(3\times 7168\times 2048\approx 44\)M 参数）。和 DeepSeek-V3 不同，K2 没有专家分组（V3 先选组再在组内选专家，用于限制跨节点通信）。

7.3 稀疏度缩放律

专家该设多少个？先定义稀疏度为“总专家数 / 激活专家数”。K2 团队在 Muon 下做了小规模的受控实验：固定激活专家数 8、共享专家 1（即固定激活参数和 FLOPs），只改变总专家数。结论是稀疏度越高，训练和验证 loss 都越低。在算力最优的缩放律下，要达到验证 loss 1.5，稀疏度 48 相比稀疏度 8、16、32 分别节省 1.69×、1.39×、1.15× 的 FLOPs。

收益在递减（从 32 到 48 只省 15%），而更高的稀疏度意味着更多专家参数要存、要通信，基础设施更复杂。K2 取了 48：384/8。

7.4 注意力头数：64 而不是 128

DeepSeek-V3 把头数设成层数的两倍左右（128），目的是更好地利用显存带宽。但上下文变长后，注意力的计算随头数线性增长：在 128K 序列长度下、总专家数固定 384 时，头数从 64 加到 128 会让推理 FLOPs 增加 83%。

这个 83% 可以粗算出来。生成第 \(L\) 个位置的 token 时，每个头要和前面 \(L\) 个 key 做点积、再对 \(L\) 个 value 加权求和；按每头 128 维算（忽略 64 维的 RoPE 分量），两步各是 \(2\cdot L\cdot 128\) FLOPs。于是每层注意力是 \(4\cdot H\cdot L\cdot 128\)，代入 \(H=64\)、\(L=131072\)、61 层：

$$4\times 64\times 131072\times 128\times 61 \approx 262\ \text{GFLOPs}.$$

其余所有矩阵乘按激活参数算，每 token 约 \(2N_{\text{act}}\approx 65\) GFLOPs。头数加倍时，注意力这部分翻倍（+262），注意力投影的参数也多出约 5B（+10 GFLOPs），总增幅约 \((262+10)/(65+262)\approx 83\%\)。也就是说，在 128K 上下文下，注意力计算已经是推理的大头，头数直接决定成本。

收益这一边，受控实验（头数 = 层数 vs 头数 = 2× 层数，同 token 数）显示加倍头数只让验证 loss 降低 0.5%–1.2%。在 Agent 场景长上下文很关键，这点收益不值得，于是 K2 用 64 头。

计算器按 MLA 和 MoE 的尺寸估算参数量和每 token 的 FLOPs。K2 预设下它给出的总参数是 1.026T，比 1.04T 少约 17B，差的是一层 MTP（multi-token prediction，预训练时额外预测后面 token 的模块，结构和一个 MoE 层相同，约 17B 参数）。K2 训练时带这一层，但公开权重不含它（num_nextn_predict_layers=0），加上它就是 1.043T。

交互演示 · MoE 稀疏度计算器
总参数（不含 MTP）–
激活参数–
稀疏度 E/k–
BF16 权重–
每 token 参数 FLOPs ≈ 2Nact–
每 token 注意力 FLOPs（该长度）–
头数加倍后总 FLOPs 增加–

参数公式：MLA 每层 = d·1536 + 1536·H·192 + d·576 + 512·H·256 + H·128·d；每个专家 3·d·2048；路由器 E·d；embedding 与 LM head 各 163840·d。注意力 FLOPs 按每层 4·H·L·128 粗算（忽略 RoPE 分量），L 是当前位置之前的上下文长度。柱状图是 K2 稀疏度实验的结果（达到验证 loss 1.5 时，稀疏度 48 相对其它稀疏度节省的 FLOPs 倍数），不随滑块变化。

在 K2 预设、128K 上下文下，“头数加倍后总 FLOPs 增加”一栏就是上面算的 83%。把上下文缩短，这个比例迅速下降：注意力 FLOPs 和 \(L\) 成正比，到 4K 时增幅只剩约 25%，其中大半来自多出来的投影参数。这正是 K2 在“为长上下文 agent 设计”的前提下选 64 头的原因。

8. 训练基础设施与训练配方

8.1 集群

K2 在 NVIDIA H800 集群上训练，每节点 8 卡、2 TB 内存，节点内 NVLink/NVSwitch，节点间 8×400 Gbps RoCE。

8.2 并行策略：能在任意 32 的倍数个节点上训练

训练资源是动态变化的，所以他们没有为某个固定规模优化，而是设计了一个能在“任意 32 的倍数个节点”上跑的方案：

PP×EP = 256 卡 = 32 节点，构成一个模型并行组。BF16 参数加 FP32 梯度累积缓冲约 6 TB（\(1.04\text{T}\times(2+4)\ \text{bytes}\)），分布在 256 卡上。优化器状态的放法取决于规模：节点多时由 ZeRO-1 摊得很薄；只有 32 节点时把一部分卸载到 CPU。这样每卡用于各种状态的显存约 30 GB，其余给激活。同一套并行配置可以用于小规模和大规模实验，这对研究效率很重要。

8.3 EP 通信与 interleaved 1F1B 的重叠

MoE 的 all-to-all（把 token 发给专家、再收回）很慢，需要和计算重叠。K2 的做法是在标准 interleaved 1F1B 调度上增加 warmup 阶段的 micro-batch 数，让一个 micro-batch 的 EP 通信与另一个的计算重叠。他们没有用 DeepSeek-V3 的 DualPipe：DualPipe 要存两份参数和梯度，显存翻倍，只能靠更大的 PP（更多 bubble）或 EP（开销更高）来补，对万亿模型代价太大。

interleaved 1F1B 的问题是 stage 更多，PP 通信更多。缓解办法是把每个 micro-batch 反向里的权重梯度计算拆出来（类似 zero-bubble 的 B/W 拆分），和对应的 PP 通信并行执行。这样除 warmup 阶段外，所有 PP 通信都能被重叠。

8.4 更小的 EP

K2 只有 64 个注意力头，注意力计算时间比 V3 短，留给 EP 通信重叠的时间也短。所以要尽量缩短 EP 操作：选最小可行的 EP=16。小 EP 组的另一个好处是专家负载均衡的约束更宽松，不需要额外调就能接近最优速度。

8.5 激活显存的三板斧

扣掉参数、梯度缓冲和优化器状态后，剩下的显存放不下完整的 MoE 激活，尤其是 1F1B warmup 阶段积累激活最多的前几个 stage。三种手段：

8.6 训练配方

9. SFT 与大规模 Agent 数据合成

9.1 SFT 的总原则

SFT 用 Muon。指令数据的两条原则是最大化提示多样性和保证回答质量。不同任务领域有各自的数据流水线，混合使用人工标注、提示工程和验证；候选回答由 K1.5 和内部的领域专家模型生成，再由 LLM 或人工评审打分过滤。

9.2 工具调用数据的合成流水线

真实环境能提供真实的交互信号，但成本、复杂度、隐私和可达性都限制了规模。K2 借鉴 ACEBench 的合成框架，建了一条模拟真实工具使用场景的流水线，生成数万条多样、高质量的训练样本。分三个阶段：

工具库3000+ 真实 MCP20000+ 合成工具 Agent系统提示 + 工具组合数千个 任务 + rubric成功标准、预期工具使用模式、检查点 多轮轨迹用户模拟器工具模拟器（有状态） LLM 评审按 rubric 过滤 → 只留成功轨迹编码/SWE 场景改用真实沙箱 + 单测 SFT 数据（相当于大规模拒绝采样）
工具调用数据合成流水线。

整体上，大规模生成 + 质量过滤本质上就是一种大规模拒绝采样（rejection sampling）。

9.3 工具调用的 token 模板

K2 的工具调用格式有几个值得记下的设计：

10. 强化学习：可验证奖励、自我批评与三个技巧

SFT 之后是 RL。K2 团队的判断是 RL 比 SFT 有更好的 token 效率和泛化，所以在 K1.5 的基础上扩大了任务多样性和训练 FLOPs，搭了一个类 Gym 的可扩展框架。

10.1 可验证奖励的“Gym”

10.2 超越验证：自我批评 rubric 奖励

创意写作、开放问答这类主观任务没有可验证的答案。K2 的办法是让模型给自己的输出打分（Self-Critique Rubric Reward）：

核心 rubric 有三条：清晰与相关（简洁且完整回应意图，非必要不长列表，需要做选择时给出唯一明确答案）；对话流畅与投入；客观落地的互动（避免元评论和对用户的奉承）。规定性 rubric 两条：不准以夸赞开头（“好问题！”）；不准写解释“这个回答为什么好”的句子。这套 rubric 也有局限：它可能偏好显得自信、果断的回答，在本该表达不确定性或多视角的场景里，模型可能过度确定。

10.3 RL 目标函数

沿用 K1.5 的策略优化算法。对每个问题 \(x\)，从旧策略 \(\pi_{\text{old}}\) 采样 \(K\) 个回答，最小化

$$L_{\text{RL}}(\theta) = \mathbb E_{x}\left[\frac1K\sum_{i=1}^K\Big(r(x,y_i)-\bar r(x) - \tau\log\frac{\pi_\theta(y_i|x)}{\pi_{\text{old}}(y_i|x)}\Big)^2\right],\qquad \bar r(x)=\frac1K\sum_i r(x,y_i).$$

怎么理解这个平方损失？在 \(\theta=\theta_{\text{old}}\) 处 log 比值为 0，梯度是

$$\nabla_\theta L = -\frac{2\tau}{K}\sum_i \big(r_i-\bar r\big)\,\nabla_\theta\log\pi_\theta(y_i|x),$$

所以第一步梯度下降等价于带均值基线的策略梯度：比平均好的回答概率上升，差的下降。随着 \(\pi_\theta\) 偏离 \(\pi_{\text{old}}\)，损失要求 \(\tau\log(\pi_\theta/\pi_{\text{old}})\) 去“拟合”优势 \(r_i-\bar r\)，而不是无限推高，这起到类似 KL 正则的稳定作用（\(\tau\) 越大更新越保守）。这个目标也用 Muon 优化。

10.4 三个技巧

11. RL 基础设施

11.1 共置架构与引擎切换

和 K1.5 一样，K2 的同步 RL 用共置（colocated）架构：训练引擎和推理引擎在同一批机器上，一个工作时另一个释放或卸载 GPU 资源。每轮由中心控制器先调推理引擎生成数据，再让训练引擎训练，然后把新参数发给推理引擎。

rollout 期间训练引擎的参数卸载到 DRAM，恢复时只需一次 H2D 拷贝。难的是推理引擎：它要从训练引擎拿到新参数，而两者分片方式不同。到 K2 规模，用网络文件系统做重分片和广播不可行（需要每秒数 PB 的总带宽）。他们做了一个和训练节点共置的分布式 checkpoint engine：每个 worker 从训练引擎拿本地参数副本，然后在所有 checkpoint engine worker 间广播完整参数集，推理引擎只取自己需要的那一片。为控制显存，更新逐参数、流水线式进行。

“广播全部参数”比理论最优的“只传需要的”多传好几倍数据，但它让训练和推理引擎完全解耦、更易维护测试；而且因为同步开销小、网络带宽利用率高，实际上反而更快。K2 的完整参数更新不到 30 秒。代码开源在 MoonshotAI/checkpoint-engine。

具体的流水线是这样的：每卡三个等大缓冲区——一个 H2D 缓冲、两个 IPC 缓冲（共享给推理引擎直接访问）。理想是三阶段流水（H2D、广播、推理引擎重载）；但在 H800 上并发 H2D 与广播会占满共享的 PCIe，三阶段退化成串行。所以实际用两阶段：先所有设备做一次同步 H2D，然后广播和重载并行。规模大时每卡分片小，整套参数一次就能放进 H2D 缓冲，开销消失。

11.2 启动与故障恢复

训练引擎启动时，每个 worker 只从磁盘读部分或不读参数，再广播给同伴，保证所有 worker 合起来只读一次 checkpoint。推理引擎是独立副本，不想引入额外同步屏障，所以复用 checkpoint engine：由它集体读盘，再用上面的方法更新未初始化的推理引擎。这也让系统对单点故障更鲁棒：单个推理副本可以独立重启。

11.3 Agentic rollout

12. 评测结果

所有对比都在非思考模式下进行（基线也用各自的非思考配置），输出上限 8192 token（SWE-bench Verified Agentless 为 16384）。下表是其中一部分。

基准K2-InstructDeepSeek-V3-0324Qwen3-235B-A22BClaude Sonnet 4Claude Opus 4GPT-4.1
LiveCodeBench v653.746.937.048.547.444.7
OJBench27.124.011.315.319.619.5
SWE-bench Verified（agentic 单次）65.838.834.472.7*72.5*54.6
SWE-bench Verified（多次尝试）71.6——80.2*79.4*—
SWE-bench Multilingual47.325.820.951.0—31.5
τ²-Bench（三域平均）66.148.837.3———
ACEBench (En)76.572.770.576.275.680.1
AIME 2025 (Avg@64)49.546.724.7*33.1*33.9*37.0
GPQA-Diamond (Avg@8)75.168.4*62.9*70.0*74.9*66.3
MMLU-Redux92.790.589.2*93.694.292.4
IFEval (Prompt Strict)89.881.183.2*87.687.488.0
SimpleQA31.027.713.215.922.842.3

* 为对应模型官方公布的数字。

基座模型（Kimi-K2-Base）方面，MMLU 87.79、MMLU-Pro 69.17、SimpleQA 35.25、EvalPlus 80.33、MATH 70.22、C-Eval 92.50，多数高于 DeepSeek-V3-Base、Llama 4 Maverick 和 Qwen2.5-72B 基座；GPQA-Diamond（48.11 vs V3 的 50.51）和 CMATH 略低于 V3。开放式评测上，截至 2025-07-17，K2-Instruct 在 LMSYS Arena 文本榜上是开源第一、总榜第五（3000 多票）。

K2 团队也列出了局限：遇到很难的推理或工具定义不清时会生成过多 token，导致截断或工具调用不完整；不必要地开启工具会让某些任务表现下降；一次性构建完整软件项目不如放在 agentic 编码框架里用。

13. 小结

把 K2 拆开看，最值得带走的是三点：

K2 之后，Moonshot 的注意力路线转向 KDA 线性注意力与 MLA 的混合，即 Kimi Linear，它也是后来 K3 的骨架。

参考与延伸阅读

想看更多实验细节，推荐直接读 K2 报告的 §2.1 和附录 E（MuonClip），以及 Moonlight 的 §2–3。