万字长文精读Kimi K2技术要点

大模型向量数据库机器学习
目录
  • 引言
  • 摘要
  • 简介
  • 预训练
  • MuonClip:通过权重裁剪实现稳定训练
  • 预训练数据:通过复述提升Token效用
  • 模型架构
  • 训练基础设施
  • 训练方案
  • 后训练阶段
  • 有监督微调
  • 强化学习
  • RL基础设施
  • 补充信息
  • 工具调用的Token模板
引言

前文简要介绍了Kimi K2的技术要点,本文将详细地介绍Kimi K2的各技术要点。

摘要

Kimi K2是一个拥有320亿(32B)激活参数和1万亿(1T)总参数的混合专家(MoE)大语言模型。K2的一个重要贡献是提出了MuonClip优化器,该优化器通过新颖的QK-clip技术改进了Muon,在享有Muon先进token效率的同时解决了训练不稳定性问题。基于MuonClip,K2在15.5万亿个token上进行了预训练,实现了零损失尖峰(loss spike)。这意味着整个训练过程非常稳定,没有出现模型性能突然急剧下降的情况。在后训练阶段,K2经历了多阶段后训练过程,其中突出的是大规模智能体数据合成流程和联合强化学习(RL)阶段,模型通过与真实和合成环境的交互来提升其能力。

Kimi K2在开源非思考模型中实现了最先进的性能(当初发布之际,并非指当下,当然即使当下也很能打),在智能体能力方面表现突出。值得注意的是,K2在Tau2-Bench上获得66.1分,在ACEBench(英文)上获得76.5分,在SWE-Bench Verified上获得65.8分,在SWE-Bench Multilingual上获得47.3分——在非思考设置下超越了大多数开源和闭源基线。它还在编码、数学和推理任务中展现出强大能力,在LiveCodeBench v6上获得53.7分,在AIME 2025上获得49.5分,在GPQA-Diamond上获得75.1分,在OJBench上获得27.1分,所有这些都没有使用扩展思考。这些结果将Kimi K2定位为迄今为止(K2发布之际)最强大的开源大型语言模型之一,特别是在软件工程和智能体任务方面。官方发布了基础和后训练模型检查点模型,以促进智能体智能的未来研究和应用。

论文: https://arxiv.org/abs/2507.20534

GitHub: https://github.com/MoonshotAI/Kimi-K2

更多AI智能体相关欢迎关注公众号"小窗幽记机器学习":

简介

大语言模型(LLMs)的发展正在经历一场深刻的范式转变,朝着智能体智能(Agentic Intelligence)的方向发展——即模型在复杂动态环境中自主感知、规划、推理和行动的能力。这一转变标志着从静态模仿学习向主动学习的转型,模型通过交互进行学习,获得超越其训练分布的新技能,并通过经验调整行为。相信这种方法能够使AI agent突破静态人类生成数据的限制,通过自身的探索(exploration)和利用(exploitation)获得超人类的能力。因此,智能体智能正迅速成为下一代基础模型的决定性能力,在工具使用、软件开发和现实世界自主性方面具有广泛影响。

实现智能体智能在预训练和后训练阶段都面临着挑战。预训练必须在有限的高质量数据约束下赋予模型广泛的通用先验知识,这使得token效率(每个token的学习信号)成为关键的扩展因素。后训练必须将这些先验知识转化为可执行的行为,然而多步推理、长期规划和工具使用等智能体能力在自然数据中很少见,且扩展成本高昂。如何可扩展合成结构化、高质量智能体轨迹 ,并结合融入偏好和自我批判的通用强化学习(RL)技术,对于弥合这一差距至关重要。

Kimi K2的贡献涵盖了预训练和后训练两个前沿领域:

  • 提出了MuonClip,这是一种新型优化器,将token高效的Muon算法与称为QK-Clip的稳定性增强机制相结合,提出MuonClip,成功地在15.5万亿个token上预训练了Kimi K2,没有出现任何损失尖峰。
  • 引入了大规模智能体数据合成pipeline,通过模拟和真实世界环境 系统地生成工具使用演示。该系统构建多样化的工具、agent、任务和轨迹,以大规模创建高保真、可验证正确的智能体交互。
  • 设计了一个通用强化学习框架,将可验证奖励(RLVR)与自我批判评分奖励机制相结合。模型不仅从外部定义的任务中学习,还从评估自身输出中学习 ,将对齐从静态扩展到开放领域。

Kimi K2在广泛的智能体和前沿基准测试中展现出强大的性能。它在Tau2-bench上获得66.1分,在ACEBench(英文)上获得76.5分,在SWE-bench Verified上获得65.8分,在SWE-bench Multilingual上获得47.3分,在非思考评估设置下超越了大多数开源和闭源基线,缩小了与Claude 4 Opus和Sonnet的差距。在编码、数学和更广泛的STEM领域,Kimi K2在LiveCodeBench v6上获得53.7分,在OJBench上获得27.1分,在AIME 2025上获得49.5分,在GPQA-Diamond上获得75.1分,进一步突显了其在通用任务中的能力。在LMSYS Arena排行榜(2025年7月17日)上,基于超过3,000个用户投票,Kimi K2排名第一的开源模型,总体排名第五。

picture.image

2 预训练

Kimi K2的基础模型是一个万亿参数的混合专家(MoE)transformer模型,在15.5万亿高质量token上进行预训练。鉴于高质量人类数据的可用性日益有限,Kimi团队的研究员认为token效率正在成为大型语言模型扩展中的关键因素。为解决这一问题,专门设计了一套用于最大化token效率的预训练技术。具体而言,采用了token高效的Muon优化器,并通过引入QK-Clip来缓解其训练不稳定性。此外,还结合合成数据,以进一步从可用的高质量token中挖掘智能。模型架构采用超稀疏MoE与多头潜在注意力(MLA),类似于DeepSeek-V3,底层基础设施旨在优化训练效率和研究效率。

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

使用token高效的Muon优化器训练Kimi K2,结合了权重衰减和一致更新RMS缩放。在之前的工作Moonlight中的实验表明,在相同的计算预算和模型大小时,使用相同的训练数据量进行训练,Muon大幅优于AdamW,使其成为提高大型语言模型训练中token效率的有效选择。

扩展Muon而引发的训练不稳定性 尽管效率很高,但扩大Muon训练规模揭示了一个挑战:由于注意力logits爆炸导致的训练不稳定性,在实验中这个问题在Muon中出现得更频繁,而在AdamW中较少。现有的缓解策略不足。例如,logit soft-cap 直接裁剪注意力logits,但查询和键之间的点积在应用上限之前仍可能过度增长。另一方面,查询-键归一化(QK-Norm)不适用于多头潜在注意力(MLA),因为其键矩阵在推理期间不会完全实例化。

用QK-Clip驯服Muon 为解决这个问题,提出了一种新的权重裁剪机制QK-Clip来显式约束注意力logits。QK-Clip通过在更新后重新缩放查询和键投影权重来限制注意力logits的增长。

设transformer层的输入表示为X。对于每个注意力头h,其查询、键和值投影计算如下:

其中

是模型参数,attention输出如下:

定义最大logit,即每个头的标量,作为该批次B中softmax的最大输入:

其中

是训练样本X中不同token的索引。

QK-Clip的核心思想是当

(注意力头的最大 logit)超过目标阈值 τ (这里设置为100) 时重新缩放

。重要的是,这个操作不会改变当前步骤中的前向/反向计算——仅使用最大logit作为指导信号来确定控制权重增长的强度。

一个朴素的实现是同时裁剪所有头:

其中

,α是平衡参数,通常设置为0.5,对查询和键应用相等的缩放。

然而,观察到实际上只有一小部分头表现出爆炸性logits。为了最小化对模型训练的干预,确定每个头的缩放因子

,并选择应用每头QK-Clip。这种裁剪对于常规多头注意力(MHA)很简单。对于MLA,仅在非共享注意力头组件上应用裁剪:

(head-specific components): 除以

  • (head-specific rotary): 除以

,

  • (shared rotary): 保持原样,以避免影响其他部分。

MuonClip:新优化器 将Muon与权重衰减、一致RMS匹配和QK-Clip集成到单个优化器中,称之为MuonClip(见算法1)。

picture.image

从几个缩放实验中展示了MuonClip的有效性。首先,使用原始Muon训练一个中等规模的9B激活和53B总参数的MoE模型。如图2(左)所示,观察到最大注意力logits迅速超过1000的量级,这说明注意力logits爆炸在Muon训练到这个规模时已经很明显。这个水平的最大logits通常会导致训练期间的不稳定性,包括显著的损失尖峰和偶尔的发散。

picture.image

图2:左:在中等规模训练运行期间,注意力logits迅速超过1000,这可能导致潜在的数值不稳定性甚至训练发散。右:使用MuonClip和τ = 100的Kimi K2在整个训练运行中的最大logits。最大logits迅速增加到100的上限值,并且仅在大约30%的训练步骤后衰减到稳定范围,展示了QK-Clip的有效调节效果。

接下来,证明QK-Clip不会降低模型性能,并确认MuonClip优化器保留了Muon的优化特性,而不会对损失轨迹产生不利影响。实验设计和发现的详细讨论在原始论文的附录D中,感兴趣的小伙伴可以查阅原始论文。

最后,使用MuonClip(τ = 100)训练Kimi K2这个大规模MoE模型,并在整个训练运行中监控最大注意力logits(图2(右))。最初,由于QK-Clip,logits被限制在100。在训练过程中,最大logits逐渐衰减到典型的操作范围,而无需对τ进行任何调整。重要的是,如图3所示,训练损失保持平滑稳定,没有可观察到的尖峰,验证了MuonClip在大规模语言模型训练中为注意力动态提供了稳健且可扩展的控制。

picture.image

Figure 3:Kimi K2 的每步训练损失曲线(未进行平滑或子采样)。整个训练过程中曲线没有出现尖峰。PS:为了清晰起见,省略了训练的最初阶段。

2.2 预训练数据:通过复述提升Token效用

预训练中的Token效率指的是每消耗一个Token在训练中带来的性能提升。提升Token效用,即每个Token所贡献的有效学习信号,有助于增强每个Token对模型更新的影响,从而直接提升Token效率。这一点在高质量Token供应有限、需要最大化利用时尤为重要。提升Token效用的一种简单方法是对同一Token进行多次暴露,但这容易导致过拟合和泛化能力下降。

Kimi K2在预训练数据方面相较于Kimi K1.5的一个关键进步,是引入了合成数据生成策略以提升Token效用。具体而言,采用了精心设计的复述流程,在不显著引发过拟合的前提下,扩大高质量Token的数量。本报告介绍了两种面向特定领域的复述技术,分别针对知识和数学领域,实现了可控的数据增强。

1、知识数据复述
在自然、知识密集型文本上的预训练存在权衡:单个epoch难以实现知识的全面吸收,而多epoch重复则收益递减且过拟合风险增加。为提升高质量知识Token的效用,提出了一套合成复述框架,主要包括以下关键环节:

  • 风格与视角多样的提示词 :为提升语言多样性并保持事实准确性,应用了一系列精心设计的提示词,引导大语言模型以不同风格和视角对原文进行忠实复述。
  • 分块自回归生成 :为保持长文档的整体连贯性并避免信息丢失,采用了基于分块的自回归重写策略。文本被分割为若干片段,分别复述后再拼接为完整段落,有效缓解了大模型常见的输出长度限制。该流程如图4所示。
  • 一致性校验 :为确保原文与复述内容的一致性,进行一致性检查,比较每段复述文本与原文的语义对齐情况,作为训练前的初步质量控制。

picture.image

Figure 4:自动回归分块复述流程用于处理较长的输入片段。该流程将输入内容拆分为较小的片段,同时保留上下文信息,依次对每个片段进行重写,最后将所有重写后的片段拼接成完整的重写段落。

通过在SimpleQA上测试数据复述与多epoch重复的准确率,比较了两者的效果。以K2的早期检查点为例,评估了三种训练策略:

  • (1)原始数据集重复10个epoch;
  • (2)数据复述一次后重复10个epoch;
  • (3)数据复述10次后仅训练1个epoch。表1显示,这些策略下准确率持续提升,验证了基于复述的数据增强方法的有效性。该方法已扩展至其他大规模知识语料库,并取得了类似的积极效果,每个语料库最多复述两次。

picture.image

2、数学数据复述
为提升数学推理能力,将高质量数学文档重写为“学习笔记”风格,方法参考SwallowMath。此外,通过将其他语言的高质量数学资料翻译为英文,进一步提升了数据多样性。尽管在数据集的复述子集上初步实验结果令人鼓舞,但将合成数据作为持续扩展策略仍处于探索阶段。主要挑战包括在不损失事实准确性的前提下,如何将该方法泛化到多样化的领域,同时最小化幻觉和避免有害内容,并确保可扩展至大规模数据集。

预训练数据总体情况
Kimi K2的预训练语料库包含15.5万亿个精心筛选的高质量Token,涵盖四大主要领域:网页文本、代码、数学和知识。大部分数据处理流程 遵循Kimi K1.5的方法。针对每个领域,均进行了严格的正确性和质量验证,并设计了有针对性的数据实验,以确保所选数据集具备高度多样性和有效性。

2.3 模型架构

Kimi K2 是拥有1.04万亿参数的混合专家(MoE)Transformer模型,激活参数为320亿。该架构设计与 DeepSeek-V3 类似,采用多头潜在注意力(MLA)作为注意力机制,模型隐藏维度为7168,MoE专家隐藏维度为2048。缩放定律分析显示,持续提升稀疏度能够显著提升模型性能 ,因此专家数量由 DeepSeek-V3 的256个提升至384个。这意味着在处理任何给定输入时,激活参数量占比下降了。为降低推理时的计算开销,注意力头数由 DeepSeek-V3 的128个减少至64个。表2详细对比了 Kimi K2 与 DeepSeek-V3 的架构参数。

picture.image

稀疏度缩放律(Sparsity Scaling Law)
针对混合专家(MoE)模型家族,基于 Muon 制定了专门的稀疏度缩放律。稀疏度定义为专家总数与激活专家数的比值。通过精心设计的小规模实验发现,在激活参数数量固定(即FLOPs恒定)的情况下,增加专家总数(即提升稀疏度)能够持续降低训练损失和验证损失,从而提升整体模型性能 (见图5)。具体而言,在计算最优的稀疏度缩放律下,若验证损失为1.5,稀疏度为48时,FLOPs分别比稀疏度为8、16、32时降低了1.69倍、1.39倍和1.15倍。尽管提升稀疏度带来更优性能,但也会增加基础设施复杂度。为平衡模型性能与成本,Kimi K2 采用稀疏度48,每次前向传播激活384个专家中的8个。

picture.image

Figure 5:稀疏性缩放定律。增加稀疏性会提升模型性能。将激活的专家数量固定为8,共享专家数量固定为1,并改变专家的总数,从而得到具有不同稀疏性水平的模型。

注意力头数
DeepSeek-V3 将注意力头数设置为模型层数的约两倍,以更好地利用内存带宽并提升计算效率。然而,随着上下文长度增加,注意力头数翻倍会显著增加推理开销,导致长序列处理效率下降。这在需要高效长上下文处理的智能体应用中成为主要瓶颈。例如,在序列长度为128k时,注意力头数从64增加到128(专家总数固定为384),推理FLOPs增加了83%。为评估该设计影响,进行了对比实验,分别测试注意力头数等于层数和为层数两倍的配置,并在不同训练FLOPs下进行对比。在等token训练条件下,注意力头数翻倍带来的验证损失提升有限(提升幅度为0.5%至1.2%),见图6。鉴于稀疏度48已具备较强性能,注意力头数翻倍带来的边际收益不足以抵消推理成本。因此,最终选择64个注意力头

picture.image

Figure 6:对于注意力头数量等于层数的模型及其注意力头数量加倍的对应模型的缩放曲线。将注意力头数量加倍,验证损失大约降低0.5%到1.2%。

2.4 训练基础设施

2.4.1 计算集群
Kimi K2 的训练在配备 NVIDIA H800 GPU 的集群上进行。每个 H800 节点配备2TB内存和8块GPU,节点内通过 NVLink 和 NVSwitch 互联。不同节点间采用8×400 Gbps RoCE互连进行通信。

2.4.2 模型扩展的并行策略
大规模语言模型训练常常需要适应动态的资源可用性。与其针对特定资源量优化单一并行策略,不如采用灵活策略,使 Kimi K2 能在任意32的倍数节点数量下进行训练。该策略结合了16路流水线并行(PP)及虚拟阶段 、16路专家并行(EP)和 ZeRO-1 数据并行。

在此配置下,模型参数以BF16格式存储,梯度累积缓冲区为FP32格式,约需6TB GPU内存,分布在256块GPU的模型并行组中。优化器状态的存放方式取决于训练配置。当训练节点总数较大时,优化器状态分布存储,单设备内存占用降至可忽略水平;当节点数较少(如32个)时,可将部分优化器状态卸载至CPU。

这种方法允许在小规模和大规模实验中复用相同的并行配置,每块GPU可持有约30GB内存用于存储所有状态。其余GPU内存用于激活存储,详见2.4.3节。如此一致的设计有助于提升研究效率,简化系统并显著加快实验迭代速度。

EP通信与交错1F1B的重叠
通过增加预热微批次数,可以在标准交错1F1B调度下,将EP的全互连通信与计算进行重叠。相比之下,DualPipe会使参数和梯度所需的内存翻倍,需要提升并行度来补偿。增加流水线并行(PP)会引入更多空闲周期,而增加专家并行(EP)则会带来更高的开销,如下所述。这些额外成本对于训练超过1万亿参数的大模型来说过于高昂,因此未采用DualPipe方案。

然而,交错1F1B会将模型划分为更多阶段,带来不可忽视的PP通信开销。为降低这一成本,将每个微批次反向传播中的权重-梯度计算与对应的PP通信解耦,并并行执行。这样,除预热阶段外,所有PP通信都能实现有效重叠。

更小的EP规模
为确保1F1B阶段计算与通信的完全重叠,K2中注意力头数由DeepSeek-V3的128降至64,导致注意力计算时间缩短,因此需尽量减少EP操作时间。为此,采用最小可行的EP并行策略,即EP=16。较小的EP分组也放宽了专家均衡约束,无需进一步调优即可实现接近最优的速度。

2.4.3 激活内存优化
在为参数、梯度缓冲区和优化器状态预留空间后,每块设备上剩余的GPU内存不足以容纳完整的MoE激活。为确保激活内存在约束范围内,尤其是在1F1B预热阶段初始流水线阶段激活量最大时,采用了以下技术:

  • 选择性重计算
    对计算开销较低但占用内存较大的阶段(如LayerNorm、SwiGLU和MLA上投影)进行重计算。此外,训练过程中还会对MoE下投影进行重计算,以进一步减少激活内存。虽然重计算为可选操作,但能有效维持GPU内存充足,防止训练初期因专家分布不均导致的崩溃。
  • 对不敏感激活采用FP8存储
    MoE上投影和SwiGLU的输入以FP8-E4M3格式(1×128 tile,FP32缩放)压缩存储。小规模实验表明,这不会带来可测量的损失增加。由于初步研究中观察到潜在的性能下降风险,因此未在计算过程中使用FP8。
  • 激活数据CPU卸载
    其余所有激活数据均卸载至CPU内存。专用拷贝引擎负责激活数据的流式卸载和加载,并与计算及通信内核重叠。在1F1B阶段,会在卸载前一微批次的前向激活的同时,预取下一微批次的反向激活。预热和冷却阶段也采用类似处理,整体流程如图7所示。尽管卸载可能因PCIe带宽拥堵对EP通信产生轻微影响,但测试结果显示EP通信仍能实现完全重叠。

2.5 训练方案

模型预训练时,采用4096 token的上下文窗口,使用MuonClip优化器(算法1)和WSD(Warmup-Stable-Decay)学习率调度(出自MiniCPM),共处理了15.5万亿token。前10万亿token在500步预热后,采用2e-4的恒定学习率训练;随后5.5万亿token采用余弦衰减,从2e-4降至2e-5。权重衰减始终设为0.1,全局批量大小为6700万token。整体训练曲线见前面的Figure 3。

在预训练后期,进行了退火阶段和长上下文激活阶段。批量大小保持在6700万token,学习率从2e-5衰减至7e-6。在此阶段,模型以4k序列长度训练了4000亿token,随后以32k序列长度训练了60亿token。为将上下文窗口扩展至128k,采用了YaRN方法。

3 后训练阶段

3.1 有监督微调

后训练阶段采用了Muon优化器,官方推荐在K2微调中使用。此前研究已表明,经过Muon预训练的检查点在Muon微调下能获得最佳性能。

为此,构建了一个涵盖多领域的大规模指令微调数据集,核心原则是最大化提示多样性并确保高响应质量。针对不同任务领域,开发了一套数据生成流程,结合了人工标注、提示工程和验证机制。采用K1.5 及其他自研领域专家模型生成各类任务的候选回复,随后通过大模型或人工评审进行自动化质量评估与筛选。对于智能体数据,设计了数据合成流程,通过多步交互式推理,训练模型的工具使用能力。

3.1.1 大规模智能体工具使用数据合成
现代大模型智能体的一项关键能力,是能自主使用陌生工具、与外部环境交互,并通过推理、执行和纠错不断优化行为。工具使用能力对于解决需要动态与真实系统交互的复杂多步任务至关重要。近期的ACEBench 和τ-bench 等基准测试强调了全面评估工具使用能力的重要性,而ToolLLM 和ACEBench 等框架则展示了教授模型高效使用数千种工具的潜力。

然而,大规模训练此类能力面临巨大挑战:真实环境虽能提供丰富且真实的交互信号,但因成本、复杂性、隐私和可访问性等限制,难以大规模构建。近期关于合成数据生成的研究(如AgentInstruct 、Self-Instruct 、StableToolBench 、ZeroSearch)已在无需依赖真实交互的情况下,展现出大规模数据生成的可行性。在这些进展基础上,并借鉴ACEBench的综合数据合成框架,开发了一套可大规模模拟真实工具使用场景的数据生成流程,能够生成数万条多样且高质量的训练样本。

该数据合成流程分为三步,如图8所示:

  • 工具规格生成:首先从真实工具和大模型合成工具中构建了一个大型工具规格库;
  • 智能体与任务生成:针对每组从工具库中采样的工具,生成相应的智能体及任务;
  • 轨迹生成:为每个智能体和任务生成轨迹,智能体通过调用工具完成任务。

picture.image

Figure 8:工具使用的数据合成流程。(a)工具规格来自现实世界的工具和大语言模型(LLMs);智能体和任务则由工具库生成。(b)多智能体流程用于生成并筛选带有工具调用的轨迹。

1、领域演化与工具生成
工具库的构建采用了两种互补方式。首先,直接从GitHub仓库获取了3000多个真实MCP(Model Context Protocol)工具,利用现有高质量工具规格。其次,通过分层领域生成流程系统性地演化合成工具:以关键类别(如金融交易、软件应用、机器人控制)为起点,在每个类别下进一步细分多个具体应用领域,并为每个领域合成专用工具,明确接口、描述和操作语义。该流程共生成了2万多种合成工具。图9通过t-SNE嵌入展示了工具库的多样性,表明MCP和合成工具共同覆盖了工具空间的互补区域。

picture.image

Figure 9:工具嵌入的t-SNE可视化。(a) 真实MCP工具根据其原始来源类别自然聚类。(b) 合成工具被组织在预定义的领域类别中,系统性覆盖工具空间。两者共同确保了不同工具功能的全面覆盖。

2、智能体多样化
通过合成多样的系统提示,并为其配置不同工具组合,生成了数千种不同的智能体。这些智能体具备不同能力、专长和行为模式,确保了潜在应用场景的广泛覆盖。

3、基于评分标准的任务生成
针对每种智能体配置,生成了从简单到复杂的各类任务。每个任务都配有明确的评分标准,规定了成功判据、预期工具使用模式和评估检查点。此方法确保了智能体表现评估的一致性和客观性。

4、多轮轨迹生成
通过以下组件模拟真实的工具使用场景:

  • 用户模拟:由大模型生成的用户角色,具备不同沟通风格和偏好,与智能体进行多轮对话,形成自然的交互模式。
  • 工具执行环境:高级工具模拟器(功能等同于世界模型)负责执行工具调用并提供真实反馈。模拟器在每次工具执行后维护并更新状态,支持复杂的多步交互及持久性效果,并引入可控的随机性,生成包括成功、部分失败和边界情况在内的多样结果。

5、质量评估与筛选
每条轨迹均由大模型评审,依据任务评分标准进行评估。仅保留符合成功标准的轨迹用于训练,确保数据质量,同时允许任务完成策略的自然多样性。

6、真实执行环境的混合方案
尽管模拟环境具备可扩展性,但其真实性有限。为此,在关键场景(如编程和软件工程任务)中,辅以真实执行沙箱。这些沙箱可执行实际代码,与真实开发环境交互,并通过测试用例通过率等客观指标提供真实反馈。通过结合模拟多样性与真实执行的真实性,显著提升了智能体的实用能力。

通过这种结合可扩展模拟与有针对性的真实执行的混合流程,生成了兼具覆盖面和真实性的多样高质量工具使用示范。大规模自动化的合成数据生成与真实执行环境的结合,通过质量筛选流程有效实现了大规模拒绝采样。这些高质量合成数据在有监督微调中,已在多种真实应用场景下显著提升了模型的工具使用能力。

3.2 强化学习

强化学习(RL)被认为在token效率和泛化能力方面优于监督微调(SFT)。基于K1.5的工作,在K2中继续扩展RL的任务多样性和训练FLOPs。为此,开发了一个类似Gym的可扩展框架,支持在各种场景下进行RL训练。通过大量具有可验证奖励的任务扩展了该框架。对于依赖主观偏好的任务,如创意写作和开放式问答,引入了自我批评奖励机制 ,让模型通过成对比较来评判自己的输出。这种方法使各个领域的任务都能从RL范式中受益。

3.2.1 可验证奖励Gym

1、数学、STEM和逻辑任务 对于数学、STEM和逻辑推理领域,RL数据准备遵循两个关键原则:多样化覆盖和适度难度。

  • 多样化覆盖 。对于数学和STEM任务,通过结合专家标注、内部问答提取流程和开放数据集收集高质量的问答对。在收集过程中,利用标签系统有意增加覆盖不足领域的覆盖率。对于逻辑任务,数据集包含多种格式,包括结构化数据任务(如多跳表格推理、跨表聚合)和逻辑谜题(如24点游戏、数独、谜语、密码算术和摩尔斯电码解码)。
  • 适度难度 。RL提示集既不应太简单也不应太难,两者都可能产生很少的信号并降低学习效率。使用SFT模型的pass@k准确率评估每个问题的难度,只选择难度适中的问题。

2、复杂指令遵循 有效的指令遵循不仅需要理解显式约束,还需要处理隐式要求、边缘情况,并在扩展对话中保持一致性。通过结合自动验证和对抗检测的混合验证框架,以及可扩展的课程生成流程来解决这些挑战。方法采用双路径系统以确保精确性和鲁棒性:

  • 混合规则验证 。实施两种验证机制:
  • (1) 通过代码解释器对具有可验证输出的指令(如长度、风格约束)进行确定性评估
  • (2) 对需要细致理解约束的指令使用LLM作为评判者进行评估。为解决模型可能声称完成指令但实际未遵守的潜在对抗行为,加入了额外的hack检查层,专门检测此类欺骗性声明。
  • 多源指令生成 。为构建训练数据,采用三种不同的生成策略以确保全面覆盖:
  • (1) 由数据团队开发的专家制作的复杂条件提示和评分标准
  • (2) 受AutoIF 启发的agentic指令增强
  • (3) 专门用于生成探测特定失败模式或边缘情况的额外指令的微调模型。这种多管齐下的方法确保了指令覆盖的广度和深度。

3、真实性 真实性对于在多轮工具使用、自主生成推理链以及开放环境交互等场景下运行的智能体模型至关重要。受 FACTS Grounding 评估框架的启发,训练了一个句子级的真实性判别模型,用于自动化验证。该判别模型能够有效检测在上下文中没有支持证据却做出事实性陈述的句子。它作为奖励模型,用于提升整体的真实性表现。

4、编码与软件工程 为增强解决竞赛级编程问题的能力,从开源数据集和合成来源收集问题及其评判器。为确保合成数据的多样性和奖励信号的正确性,纳入了从预训练数据中检索的高质量人工编写的单元测试。

对于软件工程任务,从GitHub收集大量拉取请求和问题,构建包含用户提示/问题和可执行单元测试的软件开发环境。该环境建立在强大的沙箱基础设施上,由Kubernetes提供可扩展性和安全性支持。它支持超过10,000个并发沙箱实例,性能稳定,非常适合竞赛编程和软件工程任务。

5、安全性 增强安全性的工作始于人工策划的种子提示集,手工制作以涵盖暴力、欺诈和歧视等主要风险类别。

为模拟复杂的越狱尝试(如角色扮演、文学叙事和学术话语),采用了具有三个关键组件的自动提示演化流程:

  • 攻击模型:迭代生成旨在从目标LLM引出不安全响应的对抗性提示。
  • 目标模型:对这些提示产生响应,模拟潜在漏洞。
  • 判断模型:评估交互以确定对抗性提示是否成功绕过安全机制。

使用特定任务的评分标准评估每次交互,使判断模型能够提供二元成功/失败 标签。

3.2.2 超越验证:自我批评评分奖励

为将模型对齐扩展到可验证奖励之外的任务,引入了一个基于自我批评反馈的通用强化学习框架。该方法旨在通过将从可验证场景中学到的能力扩展到更广泛的主观任务,使LLM与细致的人类偏好对齐,包括有用性、创造性、推理深度、事实性和安全性。该框架使用自我批评评分奖励机制运行,模型评估自己的输出以生成偏好信号。为将K2引导成为合格的评判者,策划了开源和内部偏好数据集的混合,并在SFT阶段初始化其批评能力。

1、自我批评策略优化 在学习循环的第一个核心过程中,K2 actor为涵盖广泛用例的通用提示生成响应。K2 critic然后通过对照评分标准组合进行成对评估来对所有结果进行排名,该组合包括核心评分标准(附录F.1)——代表Kimi珍视的AI助手基本价值观,规定性评分标准(附录F.2)——旨在消除奖励黑客行为,以及数据团队为特定指令上下文制作的人工标注评分标准。尽管某些评分标准可以指定为强制性的,K2仍保留根据其内部先验对它们进行权衡的灵活性。这种能力实现了与其不断演化的策略行为的动态持续对齐,确保模型的响应在适应特定指令的同时与其核心身份保持一致。

2、闭环critic改进和对齐 在RL训练期间,使用可验证信号改进critic模型。从可验证奖励提示生成的on-policy rollouts用于持续更新critic,这是将RLVR的客观性能信号直接提炼到其评估模型中的关键步骤。这个迁移学习过程将其更主观的判断建立在可验证数据的基础上,使可验证任务的性能提升,从而增强critic对缺乏明确奖励信号的复杂任务的判断。这个闭环过程确保critic与策略的演化同步持续重新校准其评估标准。通过将主观评估建立在可验证数据的基础上,该框架实现了与复杂、不可验证的人类目标的稳健且可扩展的对齐。

因此,这种整体对齐在广泛的领域产生了全面的性能改进,包括用户意图理解、创意写作、复杂推理和细致的语言理解。

3.2.3 RL算法

采用K1.5 中引入的策略优化算法作为K2的基础。对于每个问题x,从先前策略πold中采样K个响应{y1, ..., yk},并针对以下目标优化模型

π

θ

其中

1是采样响应的平均奖励,

是促进稳定学习的正则化参数。与SFT一样,采用Muon优化器来最小化此目标。随着在K2中将RL训练扩展到更广泛的任务范围,主要挑战是在所有领域实现一致的性能改进。为解决这个问题,对RL算法引入了几项补充。

1、预算控制 已广泛观察到RL通常会导致模型生成的响应长度大幅增加。虽然更长的响应可以使模型利用额外的测试时计算来提高复杂推理任务的性能,但在非推理领域,收益往往不能证明其推理成本的合理性。为鼓励模型合理分配推理预算,在整个RL训练过程中强制执行每个样本的最大token预算,其中预算根据任务类型确定。超过此token预算的响应将被截断并分配惩罚,这激励模型在指定限制内生成解决方案。从经验上看,这种方法显著提高了模型的token效率,鼓励在所有领域生成简洁而有效的解决方案。

2、PTX损失 为防止在联合RL训练期间可能忘记有价值的高质量数据,策划了一个包含精心挑选的高质量样本的数据集,并通过辅助PTX损失将其整合到RL目标中。这种策略不仅利用了高质量数据的优势,还减轻了对训练体系中明确存在的有限任务集过拟合的风险。这种增强大大提高了模型在更广泛领域的泛化能力。

3、温度衰减 对于创意写作和复杂推理等任务,发现在训练初期通过高采样温度促进探索至关重要。高温度允许模型生成多样化和创新的响应,从而促进发现有效策略并降低过早收敛到次优解决方案的风险。然而,在训练后期或评估期间保持高温度可能有害,因为它会引入过多的随机性并损害模型输出的可靠性和一致性。为解决这个问题,采用温度衰减计划,在整个训练过程中从探索转向利用。这种策略确保模型在最有益时利用探索,同时最终收敛到稳定和高质量的输出。

3.3 RL基础设施

3.3.1 协同定位架构

与K1.5 类似,在同步RL训练中采用了混合协同定位架构,其中训练引擎和推理引擎位于相同的工作节点上。当一个引擎正在工作时,另一个引擎会释放或卸载其GPU资源以进行配合。在RL训练的每次迭代中,中央控制器首先调用推理引擎生成新的训练数据。然后通知训练引擎在新数据上进行训练,并将更新后的参数发送给推理引擎用于下一次迭代。

每个引擎都针对吞吐量进行了大量优化。此外,随着模型规模扩展到K2的大小,引擎切换和故障恢复的延迟变得十分显著。以下将介绍在这些方面的系统设计考虑。

3.3.2 高效的引擎切换

在推出阶段,训练引擎的参数被卸载到DRAM中。因此,启动训练引擎只需简单的H2D传输步骤。然而,启动推理引擎是一个更大的挑战,因为它必须从具有不同分片范式的训练引擎获取更新的参数。

考虑到K2的规模和涉及的大量设备,使用网络文件系统进行参数重新分片和广播是不切实际的。保持低开销所需的总带宽达到每秒数PB。为了应对这一挑战,开发了一个分布式检查点引擎,与训练节点协同定位来管理参数状态。执行参数更新时,每个检查点引擎工作节点从训练引擎获取参数的本地副本,然后在所有检查点引擎工作节点之间广播完整的参数集。随后,推理引擎仅从检查点引擎检索其所需的参数分片。此过程如图10所示。为了在1T模型上实现这一点,更新以流水线方式逐参数执行,最大限度地减少内存占用(见原文附录G)。

选择在整个集群中广播完整的参数集,而不考虑每个推理工作节点上的特定分片方案。虽然这比理论上的最优方法传输了几倍的数据,但它提供了一个更简单的系统设计,对训练和推理引擎的侵入性更小。选择牺牲这种轻微的开销来完全解耦训练引擎和推理引擎,大大简化了维护和测试。

值得注意的是,由于减少了同步开销和提高了网络带宽利用率,这种方法优于按需传输方法。该系统可以在30秒内完成Kimi K2的完整参数更新,对于典型的RL训练迭代来说,这个时间可以忽略不计。

3.3.3 高效的系统启动

由于大规模训练容易出现系统故障,对于像Kimi K2这样大的模型,优化启动时间至关重要。为了启动训练引擎,让每个训练工作节点有选择地从磁盘读取部分或不读取参数,并向其对等节点广播必要的参数。设计目标是确保所有工作节点共同只读取一次检查点,最大限度地减少昂贵的磁盘IO。

由于推理引擎是独立的副本,希望避免在它们之间引入额外的同步屏障。因此,选择重用检查点引擎进行启动:让检查点引擎从磁盘共同读取检查点,类似于训练引擎的启动方式。然后使用前一节介绍的方法更新未初始化的推理引擎的状态。通过利用专用的检查点引擎,系统还能够抵御单点故障,因为推理副本可以在不与其他副本通信的情况下重新启动。

3.3.4 Agent推出

RL基础设施支持长期、多轮agent任务的训练。在推出期间,这些任务带来了独特的挑战,例如复杂的环境交互和延长的推出持续时间。这里介绍一些优化措施来缓解这些问题。

由于环境的多样性,某些交互可能会阻塞在等待环境反馈上(例如,虚拟机或代码解释器),导致GPU空闲。采用两种策略来最大化GPU利用率:

  • 将重型环境部署为可以更容易扩展的专用服务;
  • 采用大量并发推出来分摊某些昂贵交互引起的延迟。

agent推出中的另一个挑战是单个推出轨迹可能极长。为了防止长尾轨迹阻塞整个推出过程,采用了部分推出技术。这种策略允许暂停长尾未完成的任务,并在下一次RL迭代中恢复。

为了提高研究效率,还设计了一个受OpenAI Gym框架启发的统一接口,以简化新环境的集成。希望未来能将RL基础设施扩展到更多样化的交互环境。

补充信息

工具调用的Token模板

工具调用的token结构包含三个组成部分:

  • 工具声明消息:定义可用工具列表和参数模式;
  • 助手消息中的工具调用部分:编码模型调用工具的请求;
  • 工具结果消息:封装被调用工具的执行结果。

工具声明消息的原始token格式如下:

picture.image

蓝色高亮标记代表特殊token,绿色部分(用括号引用)是工具声明内容。使用TypeScript来表达工具声明内容,因为TypeScript是一种简洁的语言,具有全面的类型系统,能够用简短的文本表达工具参数的类型和约束。以下代码1展示了两个简单工具的示例,采用与OpenAI聊天完成API兼容的JSON格式,作为对比,用TypeScript定义的相同工具(列在代码2中)要简短得多。为了提高兼容性,部分训练数据也使用JSON作为工具声明语言,这样第三方框架无需额外开发即可支持工具调用方案。

代码1:使用JSON在OpenAI兼容API中定义工具:

  
[  
    {  
        " type ": " function ",  
        " function ": {  
            " name ": " get\_weather ",  
            " description ": "Get weather for a location and date ",  
            " parameters ": {  
                " type ": " object ",  
                " properties ": {  
                    " location ": {  
                        " type ": " string ",  
                        " description ": " City and country e.g. Beijing , China "  
                    },  
                    " date ": {  
                        " type ": " string ",  
                        " description ": " Date to query , format in ‘%Y -%m -%d’"  
                    }  
                },  
                " required ": [  
                    " location "  
                ]  
            }  
        }  
    },  
    {  
        " type ": " function ",  
        " function ": {  
            " name ": " Calculator ",  
            " description ": " Simple calculator ",  
            " parameters ": {  
                " properties ": {  
                    " expr ": {  
                        " type ": " string ",  
                        " description ": " Arithmetic expression in javascript "  
                    }  
                },  
                " type ": " object "  
            }  
        }  
    }  
]  

代码 2: Tool definition in TypeScript

  
namespace functions {  
// Get weather for a location and date  
type get\_weather = (\_: {  
    // City and country e.g. Beijing , China  
    location : string ,  
    // Date to query , format in ‘%Y -%m -%d’  
    date ?: string  
}) => any;  
    // Simple calculator  
    type Calculator = (\_: {  
    // Arithmetic expression in javascript  
    expr ?: string  
}) => any;  
}  

模型响应消息中工具调用部分的token模板如下所示:

picture.image

如模板所示,通过在单个响应轮次中放置多个工具调用来支持并行工具调用。每个工具调用都有唯一的调用ID,格式为functions.{tool-name}:{counter},其中tool-name是工具名称,counter是对话中所有工具调用从0开始的自增计数器。

在推理过程中,模型偶尔会生成意外的token,导致解析工具调用时出现格式错误。为解决这个问题,开发了一个名为enforcer的约束解码模块,灵感来自lm-format-enforcer6。当生成<tool_call_section_begin|> token时,它确保即将到来的工具相关token遵循预定义模板,并且JSON参数字符串遵循声明的模式。

工具结果消息只是一个用工具调用ID和相应结果编码的文本消息。

picture.image

0
0
0
0
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 XR 技术的探索与实践
火山引擎开发者社区技术大讲堂第二期邀请到了火山引擎 XR 技术负责人和火山引擎创作 CV 技术负责人,为大家分享字节跳动积累的前沿视觉技术及内外部的应用实践,揭秘现代炫酷的视觉效果背后的技术实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论