摘要
通过下一个词预测进行语言模型预训练已被证明在扩展计算能力方面是有效的,但其发展受限于可用训练数据的数量。强化学习 (RL) 的扩展为人工智能的持续进步开辟了一条新的道路,其承诺是大型语言模型 (LLM) 可以通过学习使用奖励进行探索来扩展其训练数据。然而,先前发表的工作并未产生具有竞争力的结果。鉴于此,我们报告了我们的最新多模态 LLM——Kimi k1.5 的训练实践,包括其 RL 训练技术、多模态数据配方以及基础设施优化。长上下文扩展和改进的策略优化方法是我们的方法的关键要素,它建立了一个简单而有效的 RL 框架,而无需依赖更复杂的技术,如蒙特卡洛树搜索、价值函数和过程奖励模型。值得注意的是,我们的系统在多个基准测试和多种模态上均取得了最先进的推理性能——例如,在 AIME 上达到 77.5 分,在 MATH 500 上达到 96.2 分,在 Codeforces 上达到 94 百分位,在 MathVista 上达到 74.9 分,与 OpenAI 的 o1 持平。此外,我们提出了有效的长到短方法,该方法使用长链式思维 (CoT) 技术来改进短 CoT 模型,从而实现了最先进的短 CoT 推理结果——例如,在 AIME 上达到 60.8 分,在 MATH500 上达到 94.6 分,在 LiveCodeBench 上达到 47.3 分——比现有的短 CoT 模型(如 GPT-4o 和 Claude Sonnet 3.5)高出很多(高达 +550%)。Kimi k1.5 将在 kimi.ai 上提供。
1. 引言
在缩放定律的背景下,通过下一个词预测进行语言模型预训练已被研究,其中模型参数和数据大小按比例缩放会导致智能的持续提升 (Kaplan 等人,2020 年;Hoffmann 等人,2022 年)。然而,这种方法受限于可用的高质量训练数据的数量 (Villalobos 等人,2024 年;Muennighoff 等人,2023 年)。在本报告中,我们介绍了 Kimi k1.5 的训练配方,这是我们最新使用强化学习 (RL) 训练的多模态 LLM。目标是探索持续扩展的可能新方向。使用 RL 与 LLM,模型学习使用奖励进行探索,因此不限于预先存在的静态数据集。
K1.5 的设计和训练有几个关键要素:
- 长上下文扩展:我们将 RL 的上下文窗口扩展到 128k,并观察到随着上下文长度的增加,性能持续提升。我们方法背后的一个关键思想是使用部分回溯来提高训练效率——即通过重用之前轨迹的大部分来对新的轨迹进行采样,避免从头开始重新生成新轨迹的成本。我们的观察确定上下文长度是 RL 与 LLM 持续扩展的关键维度。
- 改进的策略优化:我们推导了长 CoT RL 的公式,并采用在线镜像下降的变体进行鲁棒的策略优化。该算法通过我们有效的采样策略、长度惩罚和数据配方的优化而得到进一步改进。
- 简单框架:长上下文扩展与改进的策略优化方法相结合,建立了一个用于 LLM 学习的简单 RL 框架。由于我们能够扩展上下文长度,学习的 CoT 表现出规划、反思和纠正的特性。上下文长度的增加具有增加搜索步骤数量的效果。因此,我们证明了即使不依赖更复杂的技术,如蒙特卡洛树搜索、价值函数和过程奖励模型,也可以实现强大的性能。
- 多模态性:我们的模型在文本和视觉数据上进行联合训练,具有对这两种模态进行联合推理的能力。此外,我们提出了有效的长到短方法,该方法使用长 CoT 技术来改进短 CoT 模型。具体来说,我们的方法包括在长 CoT 激活时应用长度惩罚和模型合并。
我们的长 CoT 版本在多个基准测试和模态上均取得了最先进的推理性能——例如,在 AIME 上达到 77.5 分,在 MATH 500 上达到 96.2 分,在 Codeforces 上达到 94 百分位,在 MathVista 上达到 74.9 分,与 OpenAI 的 o1 持平。我们的模型还实现了最先进的短 CoT 推理结果——例如,在 AIME 上达到 60.8 分,在 MATH500 上达到 94.6 分,在 LiveCodeBench 上达到 47.3 分——比现有的短 CoT 模型(如 GPT-4o 和 Claude Sonnet 3.5)高出很多(高达 +550%)。结果如图 1 和图 2 所示。
2. 方法:使用 LLM 进行强化学习
Kimi k1.5 的开发包括几个阶段:预训练、普通的监督微调 (SFT)、长 CoT 监督微调以及强化学习 (RL)。本报告侧重于 RL,从 RL 提示集整理(第 2.1 节)和长 CoT 监督微调(第 2.2 节)的概述开始,然后深入讨论第 2.3 节中的 RL 训练策略。有关预训练和普通监督微调的更多详细信息,请参见第 2.5 节。
2.1 RL 提示集整理
通过初步实验,我们发现 RL 提示集的质量和多样性对于确保强化学习的有效性至关重要。一个构造良好的提示集不仅能引导模型走向强大的推理,还能降低奖励黑客攻击和过度拟合表面模式的风险。具体来说,高质量 RL 提示集具有以下三个关键属性:
- 多样化覆盖范围:提示应涵盖广泛的学科领域,例如 STEM、编码和一般推理,以提高模型的适应能力,并确保在不同领域具有广泛的适用性。
- 难度平衡:提示集应包括分布均匀的简单、中等和困难问题,以促进逐步学习并防止过度拟合特定复杂度水平。
- 准确的可评估性:提示应允许验证者进行客观可靠的评估,确保模型性能是根据正确的推理而不是表面模式或随机猜测来衡量的。
为了在提示集中实现多样化覆盖,我们采用自动过滤器来选择需要丰富推理且易于评估的问题。我们的数据集包括来自不同领域的问题,例如 STEM 领域、竞赛和一般推理任务,包含了纯文本和图像-文本问答数据。此外,我们开发了一个标记系统,用于按领域和学科对提示进行分类,确保不同学科领域之间的平衡代表 (M. Li 等人,2023 年;W. Liu 等人,2023 年)。
我们采用基于模型的方法,利用模型自身的容量来自适应地评估每个提示的难度。具体来说,对于每个提示,SFT 模型使用相对较高的采样温度生成十个答案。然后计算通过率,并将其用作提示难度的代理——通过率越低,难度越高。这种方法使难度评估与模型的内在能力保持一致,使其对 RL 训练非常有效。通过利用这种方法,我们可以预过滤大多数琐碎的情况,并在 RL 训练期间轻松探索不同的采样策略。
为了避免潜在的奖励黑客攻击 (Everitt 等人,2021 年;Pan 等人,2022 年),我们需要确保每个提示的推理过程和最终答案都可以准确验证。经验观察表明,一些复杂的推理问题可能有相对简单且易于猜测的答案,导致错误的肯定验证——即模型通过不正确的推理过程得出正确答案。为了解决这个问题,我们排除了容易出现此类错误的题目,例如选择题、对错题和证明题。此外,对于一般的问答任务,我们提出了一种简单但有效的方法来识别并删除容易黑客攻击的提示。具体来说,我们提示模型在不进行任何 CoT 推理步骤的情况下猜测潜在答案。如果模型在 N 次尝试内预测出正确答案,则认为该提示太容易黑客攻击并被删除。我们发现设置 N=8 可以删除大多数容易黑客攻击的提示。开发更先进的验证模型仍然是未来研究的一个开放方向。
2.2 长 CoT 监督微调
有了经过提炼的 RL 提示集,我们采用提示工程来构建一个小型但高质量的长 CoT 预热数据集,其中包含经过准确验证的文本和图像输入的推理路径。这种方法类似于拒绝抽样 (RS),但侧重于通过提示工程生成长 CoT 推理路径。由此产生的预热数据集旨在封装对类人推理至关重要的关键认知过程,例如规划,即模型在执行之前系统地概述步骤;评估,包括对中间步骤的批判性评估;反思,使模型能够重新考虑和完善其方法;以及探索,鼓励考虑替代解决方案。通过对这个预热数据集进行轻量级的 SFT,我们有效地引导模型将推理策略内化。因此,经过微调的长 CoT 模型在生成更详细、更逻辑连贯的响应方面表现出更强的能力,从而提高了其在各种推理任务中的表现。
2.3 强化学习
2.3.1 问题设定
2.3.2 策略优化
我们采用在线策略镜像下降的变体作为我们的训练算法 (Abbasi-Yadkori 等人,2019 年;Mei 等人,2019 年;Tomar 等人,2020 年)。该算法执行迭代操作。在第 ii 次迭代中,我们使用当前模型 πθiπθi 作为参考模型,并优化以下相对熵正则化策略优化问题:
对于那些熟悉策略梯度方法的人来说,这个梯度类似于使用采样奖励的平均值作为基线的策略梯度 (Kool 等人,2019 年;Ahmadian 等人,2024 年)。主要区别在于,响应是从 πθiπθi 而不是策略上采样的,并且应用了 l2l2 正则化。因此,我们可以将这视为通常的策略梯度算法到非策略情况下的自然扩展 (Nachum 等人,2017 年)。我们从 DD 中采样一批问题,并将参数更新为 θi+1θi+1,随后它作为下一次迭代的参考策略。由于每次迭代都考虑了不同的优化问题,因为参考策略的变化,我们也在每次迭代开始时重置优化器。
在我们的训练系统中排除了价值网络,这在以前的研究中也得到了利用 (Ahmadian 等人,2024 年)。虽然这种设计选择显著提高了训练效率,但我们也假设在经典 RL 中用于信用分配的常规价值函数的使用可能不适合我们的环境。考虑一个场景,其中模型生成了部分 CoT (z1,z2,…,zt)(z1,z2,…,zt),并且有两种可能的下一个推理步骤:zt+1zt+1 和 zt+1′zt+1′。假设 zt+1zt+1 直接导致正确答案,而 zt+1′zt+1′ 包含一些错误。如果可以获得价值函数,它会表明 zt+1zt+1 保留了比 zt+1′zt+1′ 更高的价值。根据标准的信用分配原则,选择 zt+1′zt+1′ 将受到惩罚,因为它相对于当前策略具有负优势。然而,探索 zt+1′zt+1′ 对于训练模型生成长 CoT 非常有价值。通过使用从长 CoT 派生的最终答案的正当性作为奖励信号,只要它成功地恢复并达到正确答案,模型就可以学习从采取 zt+1′zt+1′ 的模式。这个例子的关键要点是我们应该鼓励模型探索不同的推理路径,以提高其解决复杂问题的能力。这种探索性方法产生了丰富的经验,支持关键规划技能的发展。我们的主要目标不仅仅局限于在训练问题上获得高精度,而是侧重于为模型配备有效的解决问题策略,最终提高其在测试问题上的表现。
2.3.3 长度惩罚
从本质上讲,我们提倡更短的响应,并惩罚更长的正确响应,同时明确惩罚错误的长响应。然后将此基于长度的奖励与原始奖励一起添加,并带有权重参数。
在我们的初步实验中,长度惩罚可能会在初始阶段减慢训练速度。为了缓解这个问题,我们建议在训练过程中逐步引入长度惩罚。具体来说,我们采用标准的策略优化而不使用长度惩罚,然后在训练的其他部分采用恒定的长度惩罚。
2.3.4 采样策略
尽管 RL 算法本身具有相对较好的采样特性(更困难的问题提供更大的梯度),但它们的训练效率有限。因此,一些定义明确的先验采样方法可以带来更大的性能提升。我们利用多种信号来进一步改进采样策略。首先,我们收集的 RL 训练数据自然带有不同的难度标签。例如,数学竞赛问题比小学数学问题更难。其次,由于 RL 训练过程多次采样同一个问题,我们还可以跟踪每个问题的成功率作为难度指标。我们提出了两种采样方法来利用这些先验来提高训练效率。
课程采样:我们首先从训练较简单的任务开始,然后逐渐过渡到更具有挑战性的任务。由于初始 RL 模型性能有限,在非常困难的问题上花费有限的计算预算通常会产生很少的正确样本,从而导致训练效率低下。同时,我们收集的数据自然包括等级和难度标签,因此基于难度的采样是一种直观有效的方法来提高训练效率。
优先采样:除了课程采样之外,我们使用优先采样策略来关注模型表现不佳的问题。我们跟踪每个问题 ii 的成功率 sisi,并根据 $1-s_{i}$ 对问题进行采样,以便成功率较低的问题获得更高的采样概率。这将模型的努力引向其最薄弱的地方,从而加快学习速度并提高整体性能。
2.3.5 训练配方的更多细节
编码测试用例生成:由于网络上的许多编码问题没有测试用例,我们设计了一种自动生成测试用例的方法,作为训练模型 RL 的奖励。我们的重点主要集中在不需要特殊裁判的问题上。我们还假设这些问题的真实解决方案是可用的,这样我们就可以利用解决方案来生成更高质量的测试用例。
我们利用广泛认可的测试用例生成库 CYaRon1 来增强我们的方法。我们使用我们的基础 Kimi k1.5 来根据问题陈述生成测试用例。CYaRon 的使用说明和问题描述作为生成器的输入。对于每个问题,我们首先使用生成器生成 50 个测试用例,并且还为每个测试用例随机抽取 10 个真实提交。我们将测试用例与提交内容进行对比。如果至少有 7 个提交内容产生匹配结果,则认为测试用例有效。经过这一轮过滤后,我们得到一组选定的测试用例。如果至少有 9 个提交内容通过整个选定的测试用例集,则将问题和其相关选定的测试用例添加到我们的训练集中。
在统计数据方面,从 1000 个在线竞赛问题的样本中,大约 614 个不需要特殊裁判。我们开发了 463 个测试用例生成器,产生了至少 40 个有效测试用例,导致在我们的训练集中包含 323 个问题。
1.经典 RM:借鉴 InstructGPT (Ouyang 等人,2022 年) 方法,我们实现了一个基于价值头的奖励模型,并收集了大约 800k 数据点进行微调。该模型最终将“问题”、“参考答案”和“响应”作为输入,并输出一个单一标量,指示响应是否正确。
2.链式思维 RM:最近的研究 (Ankner 等人,2024 年;McAleese 等人,2024 年)表明,通过链式思维 (CoT) 推理增强的奖励模型可以显著优于经典方法,特别是在需要细致正确标准的任务上——例如数学。因此,我们收集了大约 800k CoT 标记示例的同等大数据集来微调 Kimi 模型。在与经典 RM 相同的输入基础上,链式思维方法显式生成逐步推理过程,然后在 JSON 格式中提供最终的正确性判断,从而实现更强大、更易解释的奖励信号。
在我们的手动抽查中,经典 RM 达到了大约 84.4 的准确性,而链式思维 RM 达到了 98.5 的准确性。在 RL 训练过程中,我们采用链式思维 RM 来确保更正确的反馈。
视觉数据:为了提高模型在现实世界图像推理方面的能力,并实现视觉输入与大语言模型 (LLM) 之间更有效的对齐,我们的视觉强化学习 (Vision RL) 数据主要来自三个不同的类别:现实世界数据、合成视觉推理数据以及文本渲染数据。
1.现实世界数据包括各种年级、不同科学问题,这些问题需要图形理解和推理,位置猜测任务需要视觉感知和推理,以及涉及理解复杂图表的数据分析等。这些数据集提高了模型在实际场景中进行视觉推理的能力。
2.合成视觉推理数据是人工生成的,包括程序化生成的图像和场景,旨在提高特定的视觉推理技能,例如理解空间关系、几何图案和对象交互。这些合成数据集为测试模型的视觉推理能力提供了一个受控环境,并提供了无穷无尽的训练示例。
3.文本渲染数据是通过将文本内容转换为视觉格式而创建的,使模型在处理不同模态的基于文本的查询时能够保持一致性。通过将文本文档、代码片段和结构化数据转换为图像,我们确保模型无论输入是纯文本还是文本渲染为图像(如截图或照片)都能提供一致的响应。这也帮助提高了模型处理文本密集型图像的能力。
每种数据对于构建一个全面的视觉语言模型至关重要,该模型能够有效地管理广泛的现实世界应用,同时确保在不同输入模态之间保持一致的性能。
2.4 长到短:短 CoT 模型的上下文压缩
尽管长 CoT 模型实现了强大的性能,但它消耗的测试时标记比标准的短 CoT LLM 要多。然而,可以将长 CoT 模型的思维先验转移到短 CoT 模型中,以便即使在有限的测试时标记预算下也能提高性能。我们提出了几种解决长到短问题的方法,包括模型合并 (Yang 等人,2024 年)、最短拒绝抽样、DPO (Rafailov 等人,2024 年) 和长到短 RL。以下是这些方法的详细描述:
模型合并:模型合并已被证明有助于保持泛化能力。我们还发现了它在合并长 CoT 模型和短 CoT 模型时提高标记效率的有效性。这种方法结合了一个长 CoT 模型和一个较短的模型,从而获得一个新的模型而无需训练。具体来说,我们通过简单地平均它们的权重来合并两个模型。
最短拒绝抽样:我们观察到,我们的模型对同一问题生成具有很大长度变化的响应。基于这一点,我们设计了最短拒绝抽样方法。这种方法对同一个问题采样 nn 次(在我们的实验中,n=8n=8),并选择最短的正确响应用于监督微调。
DPO:与最短拒绝抽样类似,我们利用长 CoT 模型生成多个响应样本。选择了最短的正确解决方案作为正样本,而较长的响应被视为负样本,包括错误的较长响应和正确的较长响应(比选定的正样本长 1.5 倍)。这些正负样本对形成了用于 DPO 训练的对偏好数据。
长到短 RL:在标准 RL 训练阶段之后,我们选择一个在性能和标记效率之间提供最佳平衡的模型作为基础模型,并进行单独的长到短 RL 训练阶段。在第二阶段,我们应用第 2.3.3 节中介绍的长度惩罚,并显著降低最大回溯长度,以进一步惩罚超出所需长度的响应,尽管可能是正确的。
2.5 其他训练细节
2.5.1 预训练
Kimi k1.5 基础模型是在一个多样化的、高质量的、多模态语料库上进行训练的。语言数据涵盖五个领域:英语、中文、代码、数学推理和知识。多模态数据,包括字幕、图像-文本交错、OCR、知识以及问答数据集,使我们的模型能够获得视觉-语言能力。严格的质控确保了整体预训练数据集的相关性、多样性和平衡性。我们的预训练分为三个阶段:(1) 视觉-语言预训练,在建立强大的语言基础之后,逐步进行多模态整合;(2) 冷却阶段,使用精选和合成数据巩固能力,特别是对于推理和知识型任务;(3) 长上下文激活,将序列处理扩展到 131,072 个标记。有关我们预训练工作的更多详细信息,请参见附录 B。
2.5.2 普通监督微调
我们创建了涵盖多个领域的普通 SFT 语料库。对于非推理任务,包括问答、写作和文本处理,我们首先通过人工标注创建种子数据集。这个种子数据集用于训练种子模型。随后,我们收集了多种提示,并使用种子模型为每个提示生成多个响应。然后,标注者对这些响应进行排名,并完善排名最高的响应以产生最终版本。对于推理任务,例如数学和编码问题,其中基于规则和基于奖励模型的验证比人工判断更准确、更有效,我们利用拒绝抽样来扩展 SFT 数据集。
我们的普通 SFT 数据集包含大约 1 百万个文本示例。具体来说,500k 示例用于一般问答,200k 用于编码,200k 用于数学和科学,5k 用于创意写作,以及 20k 用于长上下文任务,例如摘要、文档问答、翻译和写作。此外,我们构建了 1 百万个文本-视觉示例,涵盖各种类别,包括图表解释、OCR、图像基础对话、视觉编码、视觉推理以及带有视觉辅助的数学/科学问题。
我们首先在 32k 标记的序列长度上训练模型 1 个 epoch,然后在 128k 标记的序列长度上再训练 1 个 epoch。在第一阶段 (32k),学习率从 2 × 10−5 to 2 × 10−6,然后在第二阶段(128k)重新升温到,然后在第二阶段(128k)重新升温到 1 × 10−5,最终下降到 1 × 10−6。为了提高训练效率,我们将多个训练示例打包到每个训练序列中。
2.6 RL 基础设施
2.6.1 用于 LLM 的大规模强化学习训练系统
在人工智能领域,强化学习 (RL) 已成为大型语言模型 (LLM) 的关键训练方法 (Ouyang 等人,2022 年;Jaech 等人,2024 年),其灵感来自其在掌握复杂游戏(如围棋、星际争霸 II 和 Dota 2)方面的成功,通过 AlphaGo (Silver 等人,2017 年)、AlphaStar (Vinyals 等人,2019 年) 和 OpenAI Dota Five (Berner 等人,2019 年) 等系统。秉承这一传统,Kimi k1.5 系统采用迭代同步 RL 框架,经过精心设计,通过持续学习和适应来增强模型的推理能力。该系统的一个关键创新是引入部分回溯技术,旨在优化处理复杂推理轨迹。
如图 3a 所示,RL 训练系统通过迭代同步方法运行,每次迭代包括回溯阶段和训练阶段。在回溯阶段,回溯工作者由中央主控协调,与模型进行交互生成回溯轨迹,生成对各种输入的响应序列。然后将这些轨迹存储在重放缓冲区中,这确保了训练数据的多元化和无偏性,通过破坏时间相关性。在随后的训练阶段,培训工作者访问这些经验来更新模型的权重。这个循环过程使模型能够不断从其行动中学习,随着时间的推移调整其策略以提高性能。
中央主控作为中央指挥者,管理数据流和沟通回溯工作者、培训工作者、奖励模型评估以及重放缓冲区之间的通信。它确保系统和谐运作,平衡负载并促进高效的数据处理。
培训工作者访问这些回溯轨迹,无论是在单次迭代中完成还是分布在多次迭代中,以计算梯度更新,从而完善模型的参数并提高其性能。这个过程由奖励模型监督,该模型评估模型输出的质量,并提供基本反馈以指导训练过程。奖励模型的评估对于确定模型策略的有效性并引导模型走向最佳性能至关重要。
此外,该系统还包括一个代码执行服务,该服务专门用于处理代码相关问题,并且是奖励模型不可或缺的一部分。该服务在实际的编码场景中评估模型的输出,确保模型的学习与现实世界的编程挑战保持一致。通过将模型的解决方案与实际代码执行进行验证,这个反馈循环对于完善模型的策略并提高其在代码相关任务中的性能至关重要。
2.6.2 部分回溯用于长 CoT RL
我们工作的主要思想之一是扩展长上下文 RL 训练。部分回溯是一种关键技术,它有效地解决了处理长 CoT 特征的问题,通过管理长和短轨迹的回溯。该技术建立了一个固定的输出标记预算,限制了每个回溯轨迹的长度。如果在回溯阶段轨迹超过标记限制,则未完成的部分被保存到重放缓冲区,并在下一次迭代中继续。它确保没有单个冗长的轨迹独占系统的资源。此外,由于回溯工作者异步运行,当一些工作者忙于长轨迹时,其他工作者可以独立处理新的、短的回溯任务。异步操作通过确保所有回溯工作者都积极参与训练过程,从而优化了系统的整体性能。
如图 3b 所示,部分回溯系统的工作原理是将长响应分解为跨迭代的片段(从 iter n-m 到 iter n)。重放缓冲区充当中央存储机制,维护这些响应片段,其中只有当前迭代(iter n)需要策略内计算。先前的片段(iter n-m 到 n − 1n−1)可以从缓冲区中高效地重用,无需重复回溯。这种分段方法显著降低了计算开销:系统不是一次性回溯整个响应,而是逐步处理和存储片段,从而能够在保持快速迭代时间的同时生成更长的响应。在训练过程中,可以从损失计算中排除某些片段,以进一步优化学习过程,使整个系统既高效又可扩展。
部分回溯的实现还提供了重复检测。该系统识别生成内容中的重复序列,并提前终止它们,减少不必要的计算,同时保持输出质量。检测到的重复可以被分配额外的惩罚,有效阻止提示集中生成冗余内容。
2.6.3 训练和推理的混合部署
RL 训练过程包括以下阶段:
- 训练阶段:一开始,Megatron (Shoeybi 等人,2020 年) 和 vLLM (Kwon 等人,2023 年) 在单独的容器中执行,由一个名为 checkpoint-engine (第 2.6.3 节) 的垫片进程封装。Megatron 开始训练程序。训练完成后,Megatron 卸载 GPU 内存,并准备将当前权重转移到 vLLM。
- 推理阶段:Megatron 卸载后,vLLM 从虚拟模型权重开始,并使用从 Megatron 传输的最新权重通过 Mooncake (Qin 等人,2024 年) 更新它们。回溯完成后,checkpoint-engine 停止所有 vLLM 进程。
- 后续训练阶段:一旦分配给 vLLM 的内存被释放,Megatron 加载内存并开始另一轮训练。
我们发现现有工作难以同时支持以下所有特性:
- 复杂的并行策略:Megatron 可能与 vLLM 有不同的并行策略。Megatron 中分布在几个节点上的训练权重可能很难与 vLLM 共享。
- 最小化 GPU 资源闲置:对于策略内 RL,最近的工作,如 SGLang (L. Zheng 等人,2024 年) 和 vLLM 可能会在训练过程中保留一些 GPU,这反过来可能导致训练 GPU 闲置。如果在训练和推理之间共享相同的设备,效率会更高。
- 动态扩展能力:在某些情况下,通过增加推理节点的数量,同时保持训练过程不变,可以显著加速。我们的系统能够在需要时有效地利用闲置 GPU 节点。
如图 4 所示,我们在此基础上实现了这种混合部署框架(第 2.6.3 节),实现了从训练到推理阶段不到一分钟,反之亦然,大约十秒钟。
混合部署策略:我们提出了一个用于训练和推理任务的混合部署策略,该策略利用 Kubernetes Sidecar 容器共享所有可用的 GPU,将两种工作负载并置在一个 pod 中。这种策略的主要优点是:
- 它促进了高效的资源共享和管理,防止训练节点在训练节点和推理节点分别部署时闲置等待推理节点。
- 利用不同的部署图像,训练和推理可以各自独立迭代以提高性能。
- 该架构不限于 vLLM,其他框架也可以方便地集成。
Checkpoint Engine:Checkpoint Engine 负责管理 vLLM 进程的整个生命周期,公开 HTTP API 以触发对 vLLM 的各种操作。为了整体一致性和可靠性,我们使用 etcd 服务管理的全局元数据系统来广播操作和状态。由于 CUDA 图、NCCL 缓冲区和 NVIDIA 驱动程序的原因,vLLM 卸载主要可能难以完全释放 GPU 内存。为了最大限度地减少对 vLLM 的修改,我们终止并重新启动它以提高 GPU 利用率和故障容忍度。
Megatron 中的工作器将拥有的检查点转换为共享内存中的 Hugging Face 格式。这种转换还考虑了管道并行性和专家并行性,以便在这些检查点中只保留张量并行性。共享内存中的检查点随后被分割成碎片并注册到全局元数据系统中。我们采用 Mooncake 将检查点传输到对等节点之间的 RDMA。对 vLLM 进行一些修改以加载权重文件并执行张量并行转换。
2.6.4 代码沙盒
我们开发了沙盒作为执行用户提交代码的安全环境,优化了代码执行和代码基准评估。通过动态切换容器图像,沙盒通过 MultiPL-E (Cassano, Gouwar, D. Nguyen, S. Nguyen, 等人,2023 年)、DMOJ Judge Server 2、Lean、Jupyter Notebook 和其他图像支持不同的用例。
对于编码任务中的 RL,沙盒通过提供一致且可重复的评估机制,确保训练数据判断的可靠性。其反馈系统支持多阶段评估,例如代码执行反馈和仓库级编辑,同时保持统一的上下文,以确保跨编程语言的公平公正的基准比较。
我们部署该服务在 Kubernetes 上以提高可扩展性和弹性,通过 HTTP 端点公开它以进行外部集成。Kubernetes 的自动重启和滚动更新等功能确保了可用性和故障容忍。
为了优化性能并支持 RL 环境,我们整合了几种技术到代码执行服务中,以提高效率、速度和可靠性。这些包括:
- 使用 Crun:我们使用 crun 作为容器运行时,而不是 Docker,显著减少了容器启动时间。
- Cgroup 重用:我们预先创建 cgroup 以供容器使用,这在高并发场景中至关重要,因为为每个容器创建和销毁 cgroup 可能会成为瓶颈。
- 磁盘使用优化:使用带有上层挂载为 tmpfs 的叠加文件系统来控制磁盘写入,提供固定大小的高速存储空间。这种方法有利于临时工作负载。
这些优化提高了代码执行中 RL 的效率,为评估 RL 生成的代码提供了一个一致且可靠的环境,对于迭代训练和模型改进至关重要。
3. 实验
3.1 评估方法
由于 Kimi k1.5 是一个多模态模型,我们在不同模态的各种基准测试上进行了全面评估。详细的评估设置可以在附录 C 中找到。我们的基准测试主要涵盖以下三个类别:
3.1.1 文本基准测试
- MMLU (Hendrycks 等人,2020 年):涵盖 STEM、人文学科、社会科学等 57 个学科,难度从小学水平到高级专业水平不等,测试了世界知识和问题解决能力。
- IF-Eval (J. Zhou 等人,2023 年):用于评估大型语言模型遵循可验证指令能力的基准测试,包含 500 多个提示,例如“写一篇超过 800 字的文章”等。由于版本更新,表 3 中报告的 IF-Eval 分数来自一个中间模型。我们将根据最终模型更新分数。
- CLUEWSC (L. Xu 等人,2020 年):CLUE 基准测试中的共指消解任务,要求模型确定句子中的代词和名词短语是否共指,数据来自中文小说。
- C-EVAL (Y. Huang 等人,2023 年):一个全面的中文评估套件,用于评估基础模型的高级知识和推理能力。它包括 52 个学科、四个难度级别的 13,948 个多项选择题。
3.1.2 推理基准测试
- HumanEval-Mul:MultiPL-E (Cassano, Gouwar, D. Nguyen, S. D. Nguyen 等人,2022 年) 的一个子集。MultiPL-E 将 HumanEval 基准测试和 MBPP 基准测试扩展到涵盖多种编程范式和流行度的 18 种语言。我们选择了 8 种主流编程语言(Python、Java、C++、C#、JavaScript、TypeScript、PHP 和 Bash)的 HumanEval 翻译版本。
- LiveCodeBench (Jain 等人,2024 年):一个全面且无污染的基准测试,用于评估编码任务中的大型语言模型 (LLM)。它具有实时更新以防止数据污染、在多个编码场景中进行整体评估、高质量的问题和测试以及平衡的问题难度。我们使用 2408-2411 版本 (release v4) 的问题测试短 CoT 模型,使用 2412-2502 版本 (release v5) 的问题测试长 CoT 模型。
- AIME 2024:包含 2024 年 AIME 的竞赛问题。AIME 是一个著名的、仅限受邀参加的高中数学竞赛,评估高级数学技能,并要求扎实的基础和高水平的逻辑思维。
- MATH-500 (Lightman 等人,2023 年):一个包含 500 个问题的综合数学基准测试,涵盖代数、微积分、概率等各个数学主题,测试计算能力和数学推理能力。分数越高,表明数学问题解决能力越强。
- Codeforces:一个著名的在线判题平台,是评估长 CoT 编码模型的常用测试平台。为了在 Div2 和 Div3 竞赛中获得更高的排名,我们对 k1.5 长 CoT 模型生成的代码片段采用多数投票,并使用同一模型生成的测试用例。Codeforce ELO 评级的百分位数是从 OpenAI Day12 演讲中提取的。
3.1.3 图像基准测试
- MMMU (Yue, Ni 等人,2024 年):精心策划的 11.5K 多模态问题集合,来源包括大学考试、测验和教科书。这些问题涵盖六大主要学术领域:艺术与设计、商业、科学、健康与医学、人文与社会科学以及技术与工程。
- MATH-Vision (MATH-V) (K. Wang 等人,2024 年):一个精心策划的 3,040 个高质量数学问题集合,具有视觉上下文,来源包括真实的数学竞赛。它涵盖 16 个不同的数学学科,并分为 5 个难度级别。这个数据集提供了全面多样的挑战,是评估 LMM 数学推理能力的理想选择。
- MathVista (Lu 等人,2023 年):一个将来自各种数学和视觉任务的挑战整合在一起的基准测试,要求参与者展示细粒度的、深刻的视觉理解以及组合推理才能成功完成任务。
3.2 主要结果
3.2.1 Kimi k1.5 长 CoT 模型
Kimi k1.5 长 CoT 模型的性能如表 2 所示。通过长 CoT 监督微调(详见第 2.2 节)和视觉-文本联合强化学习(详见第 2.3 节),模型的长程推理能力得到了显著提升。测试时的计算扩展进一步增强了其性能,使模型能够在多种模态上取得最先进的成果。我们的评估表明,模型在长上下文中的推理、理解和解构信息的能力有了显著提高,代表了多模态 AI 能力的进步。
3.2.2 Kimi k1.5 短 CoT 模型
Kimi k1.5 短 CoT 模型的性能如表 3 所示。该模型整合了几种技术,包括传统的监督微调(详见第 2.5.2 节)、强化学习(详见第 2.3 节)和长到短蒸馏(详见第 2.4 节)。结果表明,k1.5 短 CoT 模型的性能与领先的开源和专有模型相比具有竞争力或更优。这些包括文本、视觉和推理挑战,在自然语言理解、数学、编码和逻辑推理方面表现出显著的优势。
3.3 长上下文扩展
我们采用一个中等规模的模型来研究 RL 与 LLM 的扩展特性。图 5 展示了在数学提示集上训练的模型变体在训练迭代过程中训练准确性和响应长度的演变。随着训练的进行,我们观察到响应长度和性能准确性同时增加。值得注意的是,更具挑战性的基准测试表现出更陡峭的响应长度增加,表明模型学会了为复杂问题生成更精细的解决方案。图 6 表明,模型的输出上下文长度与其问题解决能力之间存在很强的相关性。我们最终的 k1.5 运行将上下文长度扩展到 128k,并观察到在硬推理基准测试上持续改进。
3.4 长到短
我们将提出的长到短 RL 算法与第 2.4 节中介绍的四种方法(DPO、最短拒绝抽样和模型合并)进行了比较,重点关注长到短问题的标记效率 (X. Chen 等人,2024 年),特别是获得的长期模型如何使短期模型受益。在图 7 中,k1.5-long 代表我们选择用于长到短训练的长 CoT 模型。k1.5-short w/ rl 指的是使用长到短 RL 训练获得的短期模型。k1.5-short w/ dpo 表示通过 DPO 训练提高标记效率的短期模型。k1.5-short w/ merge 代表合并后的模型,而 k1.5-short w/ merge + rs 表示对合并后的模型应用最短拒绝抽样获得的短期模型。k1.5-shortest 代表我们在长到短训练期间获得的最短模型。如图 7 所示,与 DPO 和模型合并等其他方法相比,提出的长到短 RL 算法表现出最高的标记效率。值得注意的是,k1.5 系列的所有模型(用橙色标出)与其他模型(用蓝色标出)相比表现出更好的标记效率。例如,k1.5-short w/ rl 在 AIME2024 上实现了 60.8 的 Pass@1 分数(8 次运行的平均值),而平均仅使用 3,272 个标记。同样,k1.5-shortest 在 MATH500 上获得了 88.2 的 Pass@1 分数,同时消耗的标记数量与其他短期模型大致相同。
3.5 消融研究
模型大小和上下文长度的扩展:我们的主要贡献是将 RL 应用于增强模型生成长 CoT 的能力,从而提高其推理能力。一个自然而然的问题出现了:这种方法与简单地增加模型大小相比如何?为了证明我们方法的有效性,我们使用相同的数据集训练了两个不同大小的模型,并记录了 RL 训练期间所有检查点的评估结果和平均推理长度。这些结果如图 8 所示。值得注意的是,尽管较大的模型最初优于较小的模型,但较小的模型可以通过使用通过 RL 优化的更长 CoT 来实现可比的性能。然而,较大的模型通常比小模型具有更好的标记效率。这也表明,如果目标是尽可能高的性能,那么扩展较大模型的上下文长度具有更高的上限,并且更节省标记。然而,如果测试时的计算能力有限,训练具有较大上下文长度的小模型可能是可行的解决方案。
使用负梯度的效果:我们研究了在我们设置中使用 ReST (Gulcehre 等人,2023 年) 作为策略优化算法的有效性。ReST 与包括我们的方法在内的其他基于 RL 的方法的主要区别在于,ReST 通过拟合从当前模型采样的最佳响应来迭代地改进模型,而没有应用负梯度来惩罚不正确的响应。如图 10 所示,我们的方法表现出比 ReST 更好的样本复杂性,表明负梯度的加入显著提高了模型生成长 CoT 的效率。我们的方法不仅提高了推理质量,还优化了训练过程,在使用更少的训练样本的情况下实现了强大的性能。这一发现表明,在我们的设置中,选择合适的优化策略至关重要,因为 ReST 与其他基于 RL 的方法之间的性能差距在其他领域并不明显 (Gulcehre 等人,2023 年)。因此,我们的结果突出了选择合适的优化策略以最大限度地提高生成长 CoT 的有效性的重要性。
采样策略:我们进一步证明了我们在第 2.3.4 节中介绍的课程采样策略的有效性。我们的训练数据集 DD 包含各种难度水平的问题混合。使用我们的课程采样方法,我们首先使用 DD 进行预热阶段,然后仅关注难题来训练模型。这种方法与采用统一采样策略而不进行任何课程调整的基线方法进行了比较。如图 9 所示,我们的结果清楚地表明,所提出的课程采样方法显著提高了性能。这种改进可归因于该方法能够逐步挑战模型,使其能够更稳健地理解和掌握处理复杂问题。通过在初步的总体介绍之后将训练重点放在更困难的问题上,模型可以更好地增强其推理和问题解决能力。
4. 结论
我们介绍了我们的最新多模态 LLM——Kimi k1.5 的训练配方和系统设计。我们从实践中得出的一个关键见解是,上下文长度的扩展对于 LLM 的持续改进至关重要。我们采用优化的学习算法和基础设施优化,例如部分回溯,以实现高效的长期上下文 RL 训练。如何进一步提高长期上下文 RL 训练的效率和可扩展性仍然是未来一个重要的问题。
我们做出的另一个贡献是一系列技术的结合,这些技术实现了改进的策略优化。具体来说,我们使用 LLM 制定了长 CoT RL,并推导出在线镜像下降的变体以实现鲁棒优化。我们还尝试了采样策略、长度惩罚和优化数据配方,以实现强大的 RL 性能。
我们证明了,即使不使用更复杂的技术,如蒙特卡洛树搜索、价值函数和过程奖励模型,通过长上下文扩展和改进的策略优化也可以实现强大的性能。未来,研究在不损害模型探索能力的情况下改进信用分配和减少过度思考也将是很有趣的。
我们还观察到了长到短方法的潜力。这些方法在很大程度上提高了短 CoT 模型的性能。此外,可以将长到短方法与长 CoT RL 以迭代的方式结合起来,以进一步提高标记效率,并从给定的上下文长度预算中提取最佳性能。
参考链接: kimi_k1.5.pdf