2 预训练

2 预训练

本文为论文《Kimi K2：面向智能体能力的开源大模型》「中文全文翻译」板块的对应章节，与「论文原文」板块同一章节一一对应，层级保持一致。

Kimi K2 的基础模型是一个拥有万亿参数的混合专家（MoE）Transformer [72] 模型，在 15.5 万亿高质量 token 上进行预训练。鉴于高质量人类数据的可用性日益受限，我们认为 token 效率正成为大语言模型扩展中的关键系数。为此，我们引入了一套专门用于最大化 token 效率的预训练技术。具体而言，我们采用了 token 高效的 Muon 优化器 [33, 46]，并通过引入 QK-Clip 来缓解其训练不稳定性。此外，我们还结合了合成数据生成，以进一步从现有高质量 token 中挖掘智能。模型架构遵循与 DeepSeek-V3 [10] 类似的超稀疏 MoE 与多头潜在注意力（MLA）设计，该设计源自经验性扩展定律分析。底层基础设施的构建同时优化了训练效率和研究效率。

2.1 MuonClip：通过权重裁剪实现稳定训练

我们使用 token 高效的 Muon 优化器 [33] 训练 Kimi K2，并结合了权重衰减和一致的更新 RMS 缩放 [46]。我们在先前工作 Moonlight [46] 中的实验表明，在相同的计算预算和模型规模下——因此训练数据量也相同——Muon 的表现显著优于 AdamW [36, 48]，使其成为提升大语言模型训练中 token 效率的有效选择。

Muon 扩展时的训练不稳定性

尽管 Muon 效率很高，但在扩展 Muon 训练规模时会遇到一个挑战：由于注意力 logit 爆炸导致的训练不稳定性。在我们的实验中，这一问题在使用 Muon 时出现得更频繁，而在使用 AdamW 时较少出现。现有的缓解策略并不足够。例如，logit soft-cap [69] 直接裁剪注意力 logit，但在应用上限之前，查询和键之间的点积仍可能过度增长。另一方面，Query-Key Normalization（QK-Norm）[11, 81] 不适用于多头潜在注意力（MLA），因为其 Key 矩阵在推理过程中并未完全实例化。

使用 QK-Clip 驯服 Muon

为了解决这一问题，我们提出了一种新颖的权重裁剪机制 QK-Clip，以显式约束注意力 logit。QK-Clip 的工作原理是在更新后对查询和键投影权重进行重缩放，从而限制注意力 logit 的增长。

设 Transformer 层的输入表示为 $\mathbf{X}$。对于每个注意力头 $h$，其查询、键和值投影计算如下：

| | $\mathbf{Q}^{h}=\mathbf{X}\mathbf{W}{q}^{h},\quad\mathbf{K}^{h}=\mathbf{X}\mathbf{W}{k}^{h},\quad\mathbf{V}^{h}=\mathbf{X}\mathbf{W}_{v}^{h}.$ | |

其中 $\mathbf{W}{q},\mathbf{W}{k},\mathbf{W}_{v}$ 是模型参数。注意力输出为：

| | $\mathbf{O}^{h}=\operatorname{softmax}\left(\frac{1}{\sqrt{d}}\mathbf{Q}^{h}\mathbf{K}^{h\top}\right)\mathbf{V}^{h}.$ | |

我们将最大 logit（每个头一个标量）定义为该批次 $B$ 中 softmax 的最大输入：

| | $S_{\max}^{h}=\frac{1}{\sqrt{d}}\max_{\mathbf{X}\in B}\max_{i,j}\mathbf{Q}{i}^{h}\mathbf{K}{j}^{h\top}$ | |

其中 $i,j$ 是训练样本 $\mathbf{X}$ 中不同 token 的索引。

QK-Clip 的核心思想是：每当 $S_{\max}^{h}$ 超过目标阈值 $\tau$ 时，对 $\mathbf{W}{k},\mathbf{W}{q}$ 进行重缩放。重要的是，该操作不会改变当前步骤的前向/反向计算——我们仅将最大 logit 作为指导信号，用于确定控制权重增长的强度。

一种朴素的实现会同时裁剪所有头：

| | $\mathbf{W}{q}^{h}\leftarrow\gamma^{\alpha}\mathbf{W}{q}^{h}\qquad\mathbf{W}{k}^{h}\leftarrow\gamma^{1-\alpha}\mathbf{W}{k}^{h}$ | |

其中 $\gamma=\min(1,\tau/S_{\max})$，$S_{\max}=\max_{h}S_{\max}^{h}$，$\alpha$ 是平衡参数，通常设为 $0.5$，即对查询和键施加相等的缩放。

然而，我们在实践中观察到，只有一小部分头会出现 logit 爆炸。为了尽量减少对模型训练的干预，我们为每个头确定一个缩放因子 $\gamma_{h}=\min(1,\tau/S_{\max}^{h})$，并选择应用逐头 QK-Clip。对于常规的多头注意力（MHA），这种裁剪非常直接。对于 MLA，我们仅对未共享的注意力头组件应用裁剪：

$\textbf{q}^{C}$ 和 $\textbf{k}^{C}$（头特定组件）：各自乘以 $\sqrt{\gamma_{h}}$

$\textbf{q}^{R}$（头特定旋转）：乘以 $\gamma_{h}$，

$\textbf{k}^{R}$（共享旋转）：保持不变，以避免对其他头产生影响。

算法 1：MuonClip 优化器

图 2：左图：在一次中等规模训练运行中，注意力 logit 迅速超过 1000，这可能导致潜在的数值不稳定甚至训练发散。右图：在整个训练过程中，使用 MuonClip 且 $\tau$ = 100 时 Kimi K2 的最大 logit。最大 logit 迅速上升到上限值 100，并且仅在大约 30% 的训练步数之后才衰减到稳定范围，展示了 QK-Clip 的有效调控作用。

MuonClip：新的优化器

我们将 Muon 与权重衰减、一致的 RMS 匹配以及 QK-Clip 集成到一个优化器中，我们将其称为 MuonClip（见算法 1）。

我们通过若干扩展实验展示了 MuonClip 的有效性。首先，我们使用原始 Muon 训练了一个中等规模的 MoE 模型（90 亿激活参数、530 亿总参数）。如图 2（左）所示，我们观察到最大注意力 logit 迅速超过 1000 的量级，表明在此规模的 Muon 训练中，注意力 logit 爆炸已经十分明显。该水平的最大 logit 通常会导致训练不稳定，包括显著的损失尖峰和偶发的发散。

接下来，我们证明 QK-Clip 不会降低模型性能，并确认 MuonClip 优化器在不对损失轨迹产生不利影响的情况下保留了 Muon 的优化特性。关于实验设计和发现的详细讨论见附录 D。

最后，我们使用 $\tau=100$ 的 MuonClip 训练 Kimi K2（一个大规模 MoE 模型），并在整个训练过程中监控最大注意力 logit（图 2（右））。初始阶段，由于 QK-Clip 的作用，logit 被限制在 100。在训练过程中，最大 logit 逐渐衰减到典型的工作范围，而无需对 $\tau$ 进行任何调整。重要的是，训练损失保持平滑稳定，未观察到任何尖峰，如图 3 所示，这验证了 MuonClip 在大规模语言模型训练中能够对注意力动态提供稳健且可扩展的控制。

图 3：Kimi K2 的逐步训练损失曲线，未进行平滑或子采样。在整个训练过程中未出现任何尖峰。请注意，为清晰起见，我们省略了训练最开始阶段。

2.2 预训练数据：通过改写提升 Token 效用

预训练中的 token 效率指的是在训练过程中每消耗一个 token 所获得的性能提升。提高 token 效用——即每个 token 所贡献的有效学习信号——可以增强每个 token 对模型更新的影响，从而直接提升 token 效率。当高质量 token 的供给有限且必须被最大化利用时，这一点尤为重要。提高 token 效用的一种朴素方法是对相同 token 进行重复暴露，但这可能导致过拟合并降低泛化能力。

Kimi K2 预训练数据相较于 Kimi K1.5 的一个关键进步，是引入了一种合成数据生成策略来提升 token 效用。具体而言，我们采用了精心设计的改写流水线，在不引发显著过拟合的情况下扩充高质量 token 的数量。在本报告中，我们描述了两种分别针对知识和数学领域的专用改写技术，它们实现了这种受控的数据增强。

知识数据改写

在自然的、知识密集型文本上进行预训练面临一种权衡：单个 epoch 不足以实现全面的知识吸收，而多 epoch 重复则会带来收益递减并增加过拟合风险。为了提高高质量知识 token 的效用，我们提出了一个合成改写框架，包含以下关键组件：

风格和视角多样化的提示：为了在保持事实完整性的同时增强语言多样性，我们应用了一系列精心设计的提示。这些提示引导大语言模型以不同风格和不同视角生成原始文本的忠实改写版本。

分块自回归生成：为了保持长文档的全局连贯性并避免信息丢失，我们采用了基于分块的自回归改写策略。文本被划分为若干片段，分别进行改写，然后重新拼接成完整段落。该方法缓解了 LLM 通常存在的隐式输出长度限制。该流水线的概览如图 4 所示。

保真度验证：为了确保原始内容与改写内容之间的一致性，我们执行保真度检查，比较每个改写段落与其来源在语义上的一致性。这是训练前的初步质量控制步骤。

我们通过在 SimpleQA 上测试相应准确率，比较了数据改写与多 epoch 重复的效果。我们使用 K2 的一个早期检查点进行实验，并评估了三种训练策略：（1）将原始数据集重复 10 个 epoch，（2）将数据改写一次并重复 10 个 epoch，以及（3）将数据改写 10 次并进行单次训练。如表 1 所示，准确率在这些策略中持续提升，证明了我们基于改写的增强方法的有效性。我们将该方法扩展到其他大规模知识语料库，并观察到了类似的积极结果，每个语料库最多改写两次。

表 1：三种改写-epoch 配置下的 SimpleQA 准确率

图 4：针对长输入文本的自回归分块改写流水线。输入被拆分为保留上下文的较小块，按顺序改写，然后拼接为完整的改写段落。

数学数据改写

为了增强数学推理能力，我们按照 SwallowMath [15] 中提出的方法，将高质量数学文档改写为“学习笔记”风格。此外，我们还通过将其他语言的高质量数学材料翻译成英文来增加数据多样性。

尽管在我们数据集的改写子集上进行的初步实验显示出令人鼓舞的结果，但将合成数据作为持续扩展策略的使用仍是一个活跃的研究领域。关键挑战包括：在不损害事实准确性的前提下将该方法推广到多样化的源领域、尽量减少幻觉和非预期毒性，以及确保对大规模数据集的可扩展性。

预训练数据总体情况

Kimi K2 的预训练语料库包含 15.5 万亿 token 的精选高质量数据，涵盖四个主要领域：网页文本、代码、数学和知识。大多数数据处理流水线遵循 Kimi K1.5 [35] 中概述的方法。对于每个领域，我们都进行了严格的正确性和质量验证，并设计了有针对性的数据实验，以确保精选数据集同时具备高度多样性和有效性。

2.3 模型架构

Kimi K2 是一个拥有 1.04 万亿参数的混合专家（MoE）Transformer 模型，具有 320 亿激活参数。该架构遵循与 DeepSeek-V3 [10] 类似的设计，采用多头潜在注意力（MLA）[44] 作为注意力机制，模型隐藏维度为 7168，MoE 专家隐藏维度为 2048。我们的扩展定律分析表明，持续增加稀疏性可带来显著的性能提升，这促使我们将专家数量增加到 384 个，相比之下 DeepSeek-V3 为 256 个。为了降低推理时的计算开销，我们将注意力头数量减少到 64 个，而 DeepSeek-V3 为 128 个。表 2 展示了 Kimi K2 与 DeepSeek-V3 在架构参数上的详细比较。

表 2：Kimi K2 与 DeepSeek-V3 的架构比较

稀疏性扩展定律

我们使用 Muon 为混合专家（MoE）模型族开发了一个稀疏性扩展定律。稀疏性定义为专家总数与激活专家数之比。通过精心控制的小规模实验，我们观察到——在固定激活参数数量（即固定 FLOPs）的条件下——增加专家总数（即增加稀疏性）会持续降低训练和验证损失，从而提升整体模型性能（图 6）。具体而言，在计算最优的稀疏性扩展定律下，要达到相同的 1.5 验证损失，稀疏性 48 相比稀疏性 8、16 和 32 分别可减少 1.69×、1.39× 和 1.15× 的 FLOPs。尽管增加稀疏性可以带来更好的性能，但这一收益伴随着更高的基础设施复杂性。为了在模型性能与成本之间取得平衡，我们在 Kimi K2 中采用稀疏性 48，即每次前向传播激活 384 个专家中的 8 个。

图 5：稀疏性扩展定律。增加稀疏性可提升模型性能。我们将激活专家数固定为 8，共享专家数固定为 1，并改变专家总数，从而得到具有不同稀疏性水平的模型。

图 6：注意力头数量等于层数的模型及其注意力头数量加倍的对应模型的扩展曲线。将注意力头数量加倍可使验证损失降低约 $0.5\%$ 至 $1.2\%$。

注意力头数量

DeepSeek-V3 [10] 将注意力头数量设置为大约模型层数的两倍，以更好地利用内存带宽并提高计算效率。然而，随着上下文长度增加，将注意力头数量加倍会带来显著的推理开销，降低长序列下的效率。这在智能体应用中成为一个主要限制，因为这些应用需要高效的长上下文处理能力。例如，在序列长度为 128k 的情况下，在专家总数固定为 384 的前提下，将注意力头数量从 64 增加到 128 会导致推理 FLOPs 增加 83%。为了评估这一设计的影响，我们进行了受控实验，在不同训练 FLOPs 条件下比较注意力头数量等于层数的配置与头数加倍的配置。在等 token 训练条件下，我们观察到在不同计算预算下，将注意力头加倍仅带来验证损失的适度改善（范围为 0.5% 至 1.2%）（图 6）。考虑到稀疏性 48 已经提供了强劲性能，将注意力头数量加倍所带来的边际收益不足以抵消推理成本。因此，我们选择使用 64 个注意力头。

2.4 训练基础设施

2.4.1 计算集群

Kimi K2 在配备 NVIDIA H800 GPU 的集群上进行训练。H800 集群中的每个节点包含 2 TB RAM 和 8 个 GPU，节点内通过 NVLink 和 NVSwitch 连接。在不同节点之间，使用 $\text{8}!\times!\text{400}~\text{Gbps}$ RoCE 互连来促进通信。

2.4.2 面向模型扩展的并行策略

大语言模型的训练通常在动态变化的资源可用性条件下进行。我们没有针对仅适用于特定资源数量的单一并行策略进行优化，而是追求一种灵活的策略，使 Kimi K2 能够在任意数量为 32 的倍数的节点上进行训练。我们的策略结合了 16 路流水线并行（PP）与虚拟阶段 [28, 53, 38, 57, 47, 21]、16 路专家并行（EP）[39] 以及 ZeRO-1 数据并行 [60]。

在此设置下，以 BF16 存储模型参数并以 FP32 存储其梯度累积缓冲区大约需要 6 TB 的 GPU 内存，分布在由 256 个 GPU 组成的模型并行组上。优化器状态的放置取决于训练配置。当训练节点总数较大时，优化器状态会被分片，使其在每个设备上的内存占用降至可忽略的水平。当训练节点总数较小（例如 32 个）时，我们可以将部分优化器状态卸载到 CPU。

这种方法使我们能够在小规模和大规模实验中复用相同的并行配置，同时让每个 GPU 为所有状态保留大约 30 GB 的 GPU 内存。其余 GPU 内存用于激活值，如第 2.4.3 节所述。这种一致的设计对于研究效率非常重要，因为它简化了系统并大幅加快了实验迭代速度。

EP 通信与交错 1F1B 的重叠

通过增加预热微批次数量，我们可以在标准交错 1F1B 调度 [21, 53] 下将 EP 的 all-to-all 通信与计算重叠。相比之下，DualPipe [10] 将参数和梯度所需的内存翻倍，因此需要增加并行度来补偿。增加 PP 会引入更多气泡，而增加 EP（如下所述）会带来更高开销。对于训练超过 1 万亿参数的大模型，这些额外成本过高，因此我们选择不使用 DualPipe。

然而，交错 1F1B 将模型划分为更多阶段，引入了不可忽视的 PP 通信开销。为了缓解这一成本，我们将权重梯度计算与每个微批次的反向传播解耦，并与相应的 PP 通信并行执行。因此，除预热阶段外，所有 PP 通信都可以被有效重叠。

更小的 EP 规模

为了确保在 1F1B 阶段实现完全的计算-通信重叠，K2 中注意力计算时间的减少（相比 DeepSeek-V3 的 128 个注意力头，K2 为 64 个）要求尽量减少 EP 操作的时间。这是通过采用最小可行的 EP 并行策略（即 EP = 16）来实现的。使用更小的 EP 组还放宽了专家均衡约束，使得无需进一步调优即可达到接近最优的速度。

2.4.3 激活值缩减

在为参数、梯度缓冲区和优化器状态预留空间之后，每个设备上剩余的 GPU 内存不足以容纳完整的 MoE 激活值。为了确保激活值内存符合约束，特别是对于在 1F1B 预热阶段累积最大激活值的初始流水线阶段，我们采用了以下技术。

选择性重计算

对计算开销低但占用内存大的阶段应用重计算，包括 LayerNorm、SwiGLU 和 MLA 上投影 [10]。此外，在训练过程中对 MoE 下投影进行重计算，以进一步减少激活值内存。虽然这是可选的，但该重计算能够保持足够的 GPU 内存，防止在训练早期阶段因专家不均衡而导致崩溃。

对不敏感激活值进行 FP8 存储

MoE 上投影和 SwiGLU 的输入在 $1\times 128$ 的 tile 中压缩为 FP8-E4M3，并使用 FP32 缩放因子。小规模实验未显示可测量的损失增加。由于我们在初步研究中观察到可能存在性能下降的风险，因此未在计算中应用 FP8。

图 7：在不同 PP 阶段中重叠的计算、通信和卸载。

激活值 CPU 卸载

所有剩余的激活值都被卸载到 CPU RAM。由一个拷贝引擎负责流式处理卸载和加载，与计算和通信内核重叠。在 1F1B 阶段，我们在卸载前一个微批次的前向激活值的同时，预取下一个微批次的反向激活值。预热和冷却阶段的处理方式类似，总体模式如图 7 所示。尽管卸载可能会因 PCIe 流量拥塞而对 EP 通信产生轻微影响，但我们的测试表明 EP 通信仍能被完全重叠。

2.5 训练方案

我们使用 MuonClip 优化器（算法 1）和 WSD 学习率调度 [25]，在 4,096 token 的上下文窗口下对模型进行了预训练，共处理了 15.5T token。前 10T token 在 500 步预热后以 2e-4 的恒定学习率进行训练，随后 5.5T token 的学习率从 2e-4 余弦衰减到 2e-5。权重衰减在整个过程中设为 0.1，全局批大小保持在 67M token。整体训练曲线如图 3 所示。

在预训练接近尾声时，我们进行了退火阶段，随后是长上下文激活阶段。批大小保持在 67M token 不变，而学习率从 2e-5 衰减到 7e-6。在此阶段，模型先在 4k 序列长度上训练 4000 亿 token，然后在 32k 序列长度上再训练 600 亿 token。为了将上下文窗口扩展到 128k，我们采用了 YaRN 方法 [55]。