Kimi-k1.5:全新SOTA模型,追平满血OpenAI-o1!(论文解读)

1.摘要

1 月 20 日,月之暗面推出 Kimi 全新 SOTA 模型 ——k1.5 多模态思考模型,其实现了 SOTA (state-of-the-art)级别的多模态推理和通用推理能力。官方表示,在 short-CoT 模式下,Kimi k1.5 的数学、代码、视觉多模态和通用能力大幅超越了全球范围内短思考 SOTA 模型 GPT-4o 和 Claude 3.5 Sonnet 的水平,领先达到 550%。

-

在这篇报告中,作者介绍了最新的多模态LLM Kimi k1.5的训练方法,它是用强化学习(RL)训练的。作者的目标是探索一种可能的新方向,以便继续扩展。使用RL和LLM,模型学习以奖励进行探索,因此不限于预先存在的静态数据集。

关于k1.5的设计和训练有几个关键要素。

  • 长上下文缩放。我们将RL的上下文窗口缩放到128 k,并观察到随着上下文长度的增加,性能持续改善。方法背后的一个关键思想是使用部分部署来提高训练效率,即:通过重新使用先前轨迹来采样新轨迹,从而避免了从头开始重新生成新轨迹的成本。
  • 改进政策优化。作者推导了长CoT的RL公式,并采用online mirror descent的变体来实现鲁棒的策略优化。通过有效的采样策略、长度惩罚和数据配方的优化,对算法进行了进一步的改进。
  • 简单化框架。长上下文缩放结合改进的策略优化方法,建立了一个简化的学习资源模型学习框架。由于作者能够对上下文长度进行缩放,因此所学习的CoT表现出计划、反射和校正的属性。结果表明,在不依赖于诸如蒙特卡罗树搜索、价值函数和过程奖励模型等更复杂的技术的情况下,可以实现更强的性能。
  • 多模态。模型是在文本和视觉数据上联合训练的,它具有对两种模态进行联合推理的能力。

其中长上下文扩展和改进的策略优化方法是作者方法的关键组成部分,它建立了一个简单有效的RL框架,而不依赖于更复杂的技术,如蒙特卡洛树搜索,值函数和过程奖励模型。

-

此外,作者提出了有效的long2short方法,使用长CoT技术来改善短CoT模型。具体来说,作者的方法包括对长CoT激活和模型合并应用长度惩罚。

作者的长CoT版本在多个基准测试和模型中实现了最先进的推理性能-例如,77.5在AIME上为96.2,在MATH 500上为94%,在Codeforces上为74.9,与OpenAI的o 1匹配。作者的模型还实现了最先进的短CoT推理结果-例如,60.8在AIME上为94.6,在MATH 500上为94.6,在LiveCodeBench上为47.3,比现有的短CoT模型(如GPT-4 o和Claude Sonnet 3.5)表现更好(高达+550%)。

-

-

2.方法

Kimi k1.5的开发包括几个阶段:预训练,原始监督微调(SFT),长CoT监督微调和强化学习(RL)。本报告重点介绍RL。

RL提示设置

通过初步实验,作者发现RL提示集的质量和多样性在确保强化学习的有效性方面起着至关重要的作用。一个构造良好的提示集不仅可以引导模型进行稳健的推理,还可以减轻奖励黑客攻击和过度拟合表面模式的风险。具体来说,三个关键属性定义了一个高质量的RL提示集

  • 多样的覆盖范围:模型应该跨越广泛的学科,如STEM、编码和一般推理,以增强模型的适应性,并确保在不同领域的广泛适用性。
  • 平衡难度:提示集应该包括一系列分布均匀的简单、中等和困难的问题,以促进渐进式学习,并防止过度拟合特定的复杂性水平。
  • 准确的可评估性:验证者应该允许验证者进行客观和可靠的评估,确保模型性能是基于正确的推理而不是肤浅的模式或随机猜测来衡量的。

提示集策划的具体方法

  • 自动筛选(Automatic Filtering):作者开发了一个标记系统,按领域和学科对提示进行分类,确保不同学科领域的均衡代表性。

  • 难度评估(Difficulty Assessment):通过模型自身的回答成功率来评估每个提示的难度,成功率越低,提示难度越高。这种方法允许难度评估与模型的内在能力保持一致,使其对RL训练非常有效。通过利用这种方法,可以预过滤大多数琐碎的情况,并在RL训练期间轻松探索不同的采样策略。

  • 避免奖励欺骗(Avoiding Reward Hacking):一些复杂的推理问题可能很容易猜测到答案,导致模型通过不正确的推理过程达到正确的答案。因此作者排除容易被模型“猜测”正确答案的问题,例如多项选择题和是非题。确保每个提示的推理过程和最终答案都能得到准确的验证。

-

Long-CoT监督微调

通过改进的RL提示集,作者采用提示工程来构建一个小而高质量的长CoT预热数据集,其中包含文本和图像输入的精确验证的推理路径,其侧重于通过即时工程生成长CoT推理路径。由此产生的预热数据集旨在封装对类人推理至关重要的关键认知过程,包括

  • 规划,其中模型系统地概述了执行前的步骤;
  • 评估,涉及对中间步骤的关键评估;
  • 反思,使模型能够重新考虑和改进其方法;
  • 探索,鼓励考虑替代解决方案。

在热身数据集上进行轻量级的监督微调(SFT),帮助模型内化这些推理策略,从而在各种推理任务中表现得更好。因此,微调的长CoT模型在生成更详细和逻辑连贯的响应方面表现出更好的能力,这增强了其在不同推理任务中的性能。

-

强化学习

问题设置

给定问题xi的训练数据集D=\left \{ \left ( x_i,y_i^* \right ) \right \}_{i=1}^n和相应的真实答案yi*,目标是训练策略模型πθ来准确地解决测试问题。在复杂推理的上下文中,问题x到解决方案y的映射是非对应(non-trivial)的。为了应对这一挑战,思想链(CoT)方法提出使用一系列中间步骤z =(z1,z2,...,zm)来桥接x和y,其中每个zi是一个连贯的令牌序列,充当解决问题的重要中间步骤。当解决问题x时,思想z_t\sim \pi _\theta (\cdot |x,z_1,...,z_{t-1})进行自回归采样,然后是最终答案y\sim \pi _\theta (\cdot |x,z_1,...,z_{m})。我们用y,z\sim \pi _\theta来表示这个采样过程。请注意,思想和最终答案都是作为语言序列进行采样的。

为了进一步增强模型的推理能力,采用规划算法来探索各种思维过程,在推理时生成改进的CoT。这些方法的核心见解是在价值评估的指导下明确构建思想搜索树。这允许模型探索思维过程的不同延续,或者在遇到死胡同时回溯以研究新的方向。更详细地,设T是搜索树,其中每个节点表示部分解s=(x,z_{1:|s|})。这里s由问题x和一系列的想法z1组成:z_{1:|s|}=\left \{ z_1,...,z_{|s|} \right \},并在此节点上,|S|表示序列中思想的数量。规划算法使用评价模型v来提供反馈v(x,z_{1:|s|}),这有助于评估解决问题的当前进度,并识别现有部分解决方案中的任何错误。注意到,反馈可以通过区分性分数或语言序列来提供。在所有s ∈ T的反馈的指导下,规划算法选择最有希望的节点进行扩展,从而生长搜索树。上述过程迭代地重复,直到导出完整的解。

由于思想和反馈都可以被视为中间推理步骤,并且这些组件都可以表示为语言标记的序列,因此我们使用z来代替s和v来简化符号。因此,作者将规划算法视为直接作用于推理步骤序列A(\cdot |z_1,z_2,...).在这个框架中,存储在规划算法所使用的搜索树中的所有信息都被扁平化为提供给算法的完整上下文。

这为生成高质量CoT提供了一个有趣的视角:我们可以训练一个模型来近似这个过程,而不是显式地构建一个搜索树并实现一个规划算法。这里,思想的数量(即,语言标记)用作传统上分配给规划算法的计算预算的类比。如果可行,该方法使模型能够通过自回归预测直接在推理空间上运行隐式搜索。因此,该模型不仅学习解决一组训练问题,而且还开发了有效解决单个问题的能力,从而提高了对未知测试问题的泛化能力。

  • 这里提到的“近似”是指通过训练一个模型(如大型语言模型),让它学会模拟搜索树和规划算法的功能,而不是显式地构建这些结构。

  • 换句话说,模型通过学习大量的数据和推理路径,能够直接生成高质量的推理链,而不需要显式地构建搜索树。

因此,作者考虑使用强化学习(RL)来训练模型以生成CoT。设r是一个奖励模型,它通过分配一个值r(x,y,y^*)\in \left \{ 0,1 \right \},基于基本事实y来证明给定问题x的建议答案y的正确性。

  • 对于可验证的问题,奖励直接由预定义的标准或规则确定。例如,在编码问题中,评估答案是否通过测试用例。
  • 对于具有自由形式基础事实的问题,作者训练了一个奖励模型r,该模型预测答案是否与基础事实匹配。给定一个问题x,模型πθ通过采样过程z\sim \pi _\theta (\cdot |x)y\sim \pi _\theta (\cdot |x,z)生成CoT和最终答案。生成的CoT的质量通过它是否可以导致正确的最终答案来评估。

综上所述,作者认为优化政策的目标如下:max_{\theta }E_{(x,y^*)\sim D,(y,z)\sim \pi _{\theta }}[r(x,y,y^*)]

通过扩大RL训练,作者的目标是训练一个模型,利用简单的基于知识的CoT和计划增强的CoT的优势。该模型在推理过程中仍然自动回归采样语言序列,从而避免了在部署过程中高级规划算法所需的复杂并行化。然而,与简单的基于推理的方法的一个关键区别是,模型不应该仅仅遵循一系列推理步骤。相反,它还应该学习关键的规划技能,包括错误识别,回溯和解决方案的改进,利用整个探索的想法作为上下文信息。

策略优化

为了实现这一目标,作者采用了改进的在线镜面下降算法(Online Mirror Descent)。该算法通过迭代优化策略参数,同时引入相对熵正则化(Relative Entropy Regularization)来控制策略更新的幅度,避免过大的策略变化导致训练不稳定。

具体步骤如下:

  1. 迭代更新: 在每次迭代 i 中,使用当前策略 πθi​​ 作为参考策略,优化以下目标函数:

    其中:

    • KL(⋅∥⋅) 是Kullback-Leibler散度,用于衡量两个策略之间的差异。

    • τ>0 是正则化参数,控制策略更新的幅度。

  2. 闭式解: 该优化问题的闭式解为:

    其中 Z 是归一化因子,确保概率分布的和为1。

  3. 梯度计算: 通过采样 k 个响应 (y1​,z1​),…,(yk​,zk​) 来近似计算 logZ,并得到梯度:

    其中 rˉ 是采样奖励的均值,用于稳定训练。

  4. 优化器重置: 由于每次迭代的目标函数不同,作者在每次迭代开始时重置优化器,以避免优化器状态对新目标的影响。

通过改进的在线镜面下降算法和相对熵正则化,Kimi k1.5 的策略优化方法不仅能够高效地提升模型的推理能力,还能在训练过程中保持稳定性。此外,通过避免使用价值网络和鼓励探索多样化路径,模型能够在复杂任务中表现出更强的适应性和泛化能力。

长度惩罚

作者观察到一个过度思考的现象,虽然这会带来更好的性能,但过度冗长的推理过程在训练和推理过程中代价高昂,并且过度思考通常不是人类所喜欢的。

为了解决这个问题,作者引入了长度奖励来抑制令牌长度的快速增长,从而提高模型的令牌效率。给定k个采样响应(y1,z1),...,(yk,zk),设len(i)为(yi,zi)的长度,min_len = minilen(i),max_len = maxilen(i).如果max_len = min_len,作者将所有响应的长度奖励设置为零,因为它们具有相同的长度。否则,长度奖励为len\_reward(i)=\begin{cases} \lambda & \text{ if } r(x,y_i,y^*)=1 \\ min(0,\lambda ) & \text{ if } r(x,y_i,y^*)=0 \end{cases}\\,where \lambda =0.5-\frac{len(i)-min\_len}{max\_len-min\_len}

从本质上讲,作者促进较短的回答,并惩罚正确答案中较长的回答,同时明确惩罚不正确答案的长回答。然后,将该基于长度的奖励与具有加权参数的原始奖励相加。

在作者的初步实验中,长度惩罚可能会在初始阶段减慢训练速度。为了缓解这个问题,作者建议在训练过程中逐渐预热长度惩罚。具体来说,作者采用标准策略优化,没有长度惩罚,然后在其余的训练中使用恒定长度惩罚。

采样策略

尽管RL算法本身具有相对较好的采样特性(更困难的问题提供更大的梯度),但它们的训练效率有限。因此,一些定义良好的先验采样方法可能会产生更大的性能增益。作者利用多个信号,以进一步改善采样策略。首先,作者收集的RL训练数据自然带有不同的难度标签。例如,数学竞赛题比小学数学题更难。其次,由于RL训练过程多次对同一问题进行采样,还可以跟踪每个问题的成功率作为难度的度量。作者提出了两种采样方法来利用这些先验知识来提高训练效率。

Curriculum Sampling(周期采样):从较简单的任务开始训练,然后逐渐发展到更具挑战性的任务。由于初始RL模型的性能有限,因此在非常困难的问题上花费有限的计算预算通常会产生很少的正确样本,从而导致较低的训练效率。同时,收集的数据自然包括等级和难度标签,使得基于难度的采样成为提高训练效率的直观有效的方法。

Prioritized Sampling(优先级采样):除了周期采样之外,作者还使用优先抽样策略来关注模型表现不佳的问题。跟踪每个问题i的成功率si,并按1-si的比例对问题进行采样,因此成功率较低的问题会得到较高的采样概率。这将模型的努力引向其最薄弱的领域,从而实现更快的学习和更好的整体性能。

更多关于训练的细节信息

测试用例生成编码 由于测试用例不适用于网络上的许多编码问题,作者设计了一种方法来自动生成测试用例,作为使用RL训练模型的奖励。作者还假设这些问题有真实的解决方案,以便我们可以利用这些解决方案来生成更高质量的测试用例。

作者利用广泛认可的测试用例生成库,CYaRon,以提高性能。作者采用基础Kimi k1.5生成测试用例的基础上的问题陈述。CYaRon的使用说明和问题描述作为生成器的输入提供。对于每个问题,作者首先使用生成器生成50个测试用例,并为每个测试用例随机抽取10个真实提交结果。作者根据提交的资料进行测试。如果10次提交中至少有7次产生匹配结果,则测试用例被视为有效。经过这一轮的过滤,作者得到一组选定的测试用例。如果10个提交中至少有9个通过了整个选定测试用例集,则将问题及其相关的选定测试用例添加到我们的训练集。

从统计数据来看,作者开发了463个测试用例生成器,产生了至少40个有效的测试用例,训练集中包含323个问题。

数学奖励模型 评估数学解决方案的一个挑战是,不同的书面形式可以代表相同的基本答案。例如,a2 − 4和(a + 2)(a − 2)可能都是同一问题的有效解。作者采用了两种方法来提高奖励模型的评分精度:

  • 经典RM:从InstructGPT中汲取灵感,作者实施了一个基于价值头的奖励模型,并收集了大约80万个数据点进行微调。该模型最终将“问题”、“参考答案”和“响应”作为输入,并输出指示响应是否正确。
  • 思想链RM:最近的研究表明,用思想链(CoT)推理增强的奖励模型可以显著优于经典方法,特别是在数学等微妙的正确性标准很重要的任务上。因此,作者收集了一个同样大的数据集,其中包含大约80万个CoT标记的示例,以微调Kimi模型。基于与经典RM相同的输入,思想链方法在以JSON格式提供最终正确性判断之前显式生成一个逐步推理过程,从而实现更强大和可解释的奖励信号。

在作者的手动抽查中,Classic RM的准确度约为84.4,而Chain-ofThought RM的准确度达到98.5。在RL训练过程中,作者采用了思想链RM,以确保更正确的反馈。

为了提高模型的真实世界图像推理能力,并实现视觉输入和大型语言模型(LLM)之间更有效的对齐,作者的视觉强化学习(Vision RL)数据主要来自三个不同的类别:真实世界数据、合成视觉推理数据和文本渲染数据。

  1. 真实世界的数据 包括一系列跨不同年级的科学问题,这些问题需要图形理解和推理,需要视觉感知和推理的位置猜测任务,以及涉及理解复杂图表的数据分析,以及其他类型的数据。这些数据集提高了模型在现实世界场景中执行视觉推理的能力。
  2. 合成视觉推理数据是人工生成的,包括程序创建的图像和场景,旨在提高特定的视觉推理技能,如理解空间关系,几何图案和对象交互。这些合成数据集为测试模型的视觉推理能力提供了一个受控的环境,并提供了无穷无尽的训练示例。
  3. 文本呈现的数据是通过将文本内容转换为视觉格式来创建的,使模型能够在处理不同模态的基于文本的查询时保持一致性。通过将文本文档、代码片段和结构化数据转换为图像,我们确保模型提供一致的响应,无论输入是纯文本还是呈现为图像的文本(如屏幕截图或照片)。这也有助于增强模型在处理大量文本图像时的能力。

每种类型的数据对于构建一个全面的视觉语言模型都是必不可少的,该模型可以有效地管理各种现实世界的应用程序,同时确保各种输入模式的一致性能。

-

Long 2short:Short-CoT模型的上下文压缩

虽然长CoT模型实现了强大的性能,但与标准的短CoT LLM相比,它消耗了更多的测试时间令牌。然而,可以将思维先验从长CoT模型转移到短CoT模型,以便即使在有限的测试时间令牌预算下也可以提高性能。作者提出了几种解决这个long2short问题的方法,包括模型合并,最短拒绝采样,DPO和long2short RL。这些方法的详细说明如下:

模型合并 模型合并已被发现在保持泛化能力方面是有用的。作者还发现在合并长cot模型和短cot模型时,可以提高令牌的效率有效性。这种方法将长cot模型与较短的模型相结合,以获得一个新的模型,而无需训练。具体来说,作者通过简单地平均它们的权重来合并这两个模型。

最短拒绝采样:作者观察到,模型针对同一问题生成的响应长度变化很大。在此基础上,作者设计了最短拒绝抽样方法。该方法对同一问题进行n次采样(在实验中,n = 8),并选择最短的正确响应进行监督微调。

DPO 与最短拒绝抽样类似,作者利用长CoT模型生成多个响应样本。选择最短的正确解作为正样本,而较长的响应被视为负样本,包括错误的较长响应和正确的较长响应(比所选的正样本长1.5倍)。这些正负对形成用于DPO训练的成对偏好数据。

long2short RL 在标准的RL训练阶段之后,作者首先选择一个在性能和令牌效率之间提供最佳平衡的模型作为基础模型,然后进行单独的long2short RL训练阶段。在 Long2short RL 训练阶段中,作者引入长度惩罚机制,显著减少最大展开长度(rollout length),进一步惩罚那些超出期望长度但可能正确的响应,通过长度惩罚,鼓励模型生成更短的推理路径,同时保持较高的正确率,从而在测试时减少 Token 消耗。

-

其他训练细节

预训练

Kimi k1.5基础模型是在多样化的高质量多模态语料库上训练的。语言数据涵盖五个领域:英语,汉语,代码,数学推理和知识。

多模态数据,包括字幕,图像-文本交织,OCR,知识和QA数据集,使模型能够获得视觉语言能力。

预训练分为三个阶段:(1)视觉语言预训练,建立强大的语言基础,然后逐步进行多模态整合;(2)冷却,使用策划和合成数据巩固能力,特别是推理和基于知识的任务;(3)长上下文激活,将序列处理扩展到131,072个令牌。

原始监督微调

作者创建了覆盖多个领域的SFT语料库。对于非推理任务,包括问答、写作和文本处理,首先通过人工注释构建种子数据集。该种子数据集用于训练种子模型。随后,收集了各种各样的提示,并采用种子模型来生成对每个提示的多个响应。然后,注释器对这些响应进行排名,并优化排名最高的响应以生成最终版本。对于数学和编码问题等推理任务,基于规则和基于奖励模型的验证比人类判断更准确和有效,利用拒绝采样来扩展SFT数据集。

作者的SFT数据集包含大约100万个文本示例。具体来说,50万个示例用于一般问题回答,20万个用于编码,20万个用于数学和科学,5万个用于创意写作,20万个用于长上下文任务,如总结,doc-qa,翻译和写作。

此外,作者还构建了100万个文本视觉示例,包括图表解释、OCR、基于图像的对话、视觉编码、视觉推理以及使用视觉辅助工具的数学/科学问题。

作者首先以32 k token的序列长度训练模型,然后以128 k token的序列长度训练另一个epoch。在第一阶段(32 k),学习率从2×10−5衰减到2×10−6,然后在第二阶段(128 k)重新预热到1×10−5,最后衰减到1×10−6。为了提高训练效率,将多个训练样本打包到每个训练序列中。

-

RL基础设施

LLM大规模强化学习训练系统

作者介绍了 Kimi k1.5 的大规模强化学习训练系统,这是一个专为大语言模型(LLM)设计的高效训练框架。该系统采用迭代同步方法,每个迭代包含展开阶段和训练阶段。在展开阶段,多个展开工作器(rollout workers)与模型交互生成轨迹,并将这些轨迹存储在重放缓冲区(replay buffer)中。这些轨迹随后被训练工作器(trainer workers)用于更新模型权重,从而实现模型的持续优化。

系统的核心是中央主控(master),它负责协调展开工作器、训练工作器、奖励模型(reward models)和重放缓冲区之间的数据流和通信,确保整个训练过程的高效运行。此外,系统还集成了代码执行服务,用于验证模型生成代码的实际运行结果,并将这些结果反馈给奖励模型,以确保模型学习与实际编程挑战紧密对齐。

为了进一步提升效率,系统采用了部分展开技术(Partial Rollouts),允许将长推理轨迹分解为多个较短的片段,并在多个迭代中逐步完成整个推理过程。这种技术不仅提高了计算资源的利用率,还使得系统能够高效处理长推理链,同时保持快速的迭代速度。

总的来说,2.6.1 节详细描述了 Kimi k1.5 的强化学习训练系统,展示了其如何通过高效的资源管理和优化技术,实现大规模强化学习训练的高效性和可扩展性。

长CoT RL的部分rollout

作者详细介绍了 Kimi k1.5 模型中用于长推理链强化学习(Long CoT RL)的部分展开技术(Partial Rollouts)。这一技术旨在优化长文本推理的处理,通过将长推理轨迹分解为多个较短的片段,并在多个训练迭代中逐步完成整个推理过程,从而提高训练效率和资源利用率。

部分展开技术的核心在于设定一个固定的输出 Token 预算,限制每次展开的轨迹长度。如果某个轨迹在展开过程中超过 Token 预算,未完成的部分会被保存到重放缓冲区(Replay Buffer),并在后续迭代中继续展开。这种分段处理方式不仅避免了单个过长轨迹对系统资源的独占,还通过异步操作最大化了计算资源的利用率。展开工作器可以独立处理新的、较短的轨迹任务,而无需等待长轨迹完成。

重放缓冲区作为中央存储机制,保存这些轨迹片段。在训练过程中,只有当前迭代的片段需要进行策略计算,而之前的片段可以直接从缓冲区复用,避免了重复展开的开销。这种分段处理显著减少了计算开销,使得系统能够高效处理长推理链,同时保持快速的迭代速度。

此外,部分展开技术还引入了重复检测机制。系统能够识别生成内容中的重复序列,并提前终止这些重复部分,进一步减少不必要的计算。重复的序列还可以被赋予额外的惩罚,以有效遏制冗余内容的生成。

Partial Rollouts for Long CoT RL 是 Kimi k1.5 模型在强化学习训练中采用的关键技术之一。通过将长推理链分解为多个较短的片段,并在多个迭代中逐步完成整个推理过程,该技术显著提高了训练效率和可扩展性。这种优化方法不仅加快了训练速度,还使得模型能够处理更复杂的多步推理任务,同时保持计算资源的有效利用。

训练和推理的混合部署

本节中,作者介绍了 Kimi k1.5 模型的 “混合部署策略”(Hybrid Deployment of Training and Inference),这是一种将训练和推理任务集成在同一套基础设施中的高效部署方法。这种策略通过 Kubernetes 的 Sidecar 容器实现,允许训练和推理任务共享 GPU 资源,从而避免了在训练和推理阶段切换时的资源闲置问题。

混合部署框架的核心是 Checkpoint Engine,它负责管理 vLLM(推理引擎)的生命周期,通过 HTTP API 触发 vLLM 的各种操作。为了确保系统的整体一致性和可靠性,Checkpoint Engine 利用 etcd 服务作为全局元数据管理系统,用于广播操作指令和状态信息。

在训练阶段,Megatron(训练引擎)负责执行训练任务。训练完成后,Megatron 会将 GPU 内存卸载,并将当前权重传递给 vLLM。vLLM 在接收到权重后开始推理任务,并在完成推理后由 Checkpoint Engine 停止所有 vLLM 进程。随后,Megatron 重新加载 GPU 内存,开始新一轮的训练。

这种混合部署策略不仅提高了资源利用率,还支持动态扩展。例如,在保持训练过程不变的情况下,可以通过增加推理节点的数量来加速推理任务。此外,该架构不仅限于 vLLM,还可以方便地集成其他框架。

为了实现高效的权重传递,Megatron 的工作节点将检查点转换为 Hugging Face 格式,并存储在共享内存中。这些检查点随后被分割成多个分片,并在全局元数据系统中注册。通过 Mooncake 工具,这些分片可以在节点之间通过 RDMA(远程直接内存访问)高效传输。vLLM 需要进行一些修改,以加载权重文件并执行张量并行转换。

Code Sandbox

代码沙盒(Code Sandbox),这是一个为执行用户提交的代码而设计的安全环境。代码沙盒的核心目标是为代码执行和代码基准评估提供优化支持,确保在强化学习(RL)任务中,尤其是涉及编程任务时,训练数据的可靠性和一致性。

代码沙盒通过动态切换容器镜像来支持多种用例,例如 MultiPL-E、DMOJ Judge Server 2、Lean、Jupyter Notebook 等,从而适应不同的编程语言和任务需求。在强化学习的代码任务中,代码沙盒通过提供一致且可重复的评估机制,确保训练数据的可靠性。其反馈系统支持多阶段评估,例如代码执行反馈和仓库级编辑,同时保持统一的上下文,以确保跨编程语言的公平和公正的基准比较。

为了优化性能并支持 RL 环境,代码沙盒集成了多种技术来提高效率、速度和可靠性。这些技术包括:

  • 使用 crun 作为容器运行时,而不是 Docker,从而显著减少了容器启动时间。

  • Cgroup 重用:预先创建 cgroups 以供容器使用,这对于高并发场景至关重要,因为频繁创建和销毁 cgroups 会成为瓶颈。

  • 磁盘使用优化:采用覆盖文件系统,将上层挂载为 tmpfs,以控制磁盘写入,提供固定大小的高速存储空间,适用于短暂工作负载。

这些优化措施显著提高了代码执行的效率,为评估 RL 生成的代码提供了一个一致且可靠的环境,这对于迭代训练和模型改进至关重要。代码沙盒部署在 Kubernetes 上,利用其可扩展性和弹性,并通过 HTTP 端点暴露服务以供外部集成。Kubernetes 的自动重启和滚动更新功能确保了服务的可用性和容错能力。

-

-

3.实验

由于k1.5是一个多模态模型,作者对不同模态的各种基准进行了综合评估。的基准主要包括以下三个类别:

  • 文本基准:MMLU,IF-Eval,CLUEWSC、C-EVAL
  • 理解基准:HumanEval-Mul,LiveCodeBench,Codeforces,AIME 2024,MATH500
  • 视觉基准:MMMU,MATH-Vision,MathVista

K1.5 long-CoT型号Kimi k1.5 long-CoT型号的性能如表2所示。通过长期CoT监督微调和视觉-文本联合强化学习,模型的长期推理能力得到了显著增强。测试时间计算缩放进一步增强了其性能,使模型能够在一系列模态中实现最先进的结果。我们的评估显示,该模型在扩展上下文中推理、理解和综合信息的能力有了显着提高,这代表了多模态AI能力的进步。

K1.5 short-CoT车型Kimi k1.5 short-CoT车型的性能如表3所示。该模型集成了多种技术,包括传统的监督微调,强化学习和长到短蒸馏。结果表明,与领先的开源和专有模型相比,k1.5短CoT模型在多个任务中提供了具有竞争力或上级的性能。这些挑战包括文本,视觉和推理挑战,在自然语言理解,数学,编码和逻辑推理方面具有显着优势。

作者采用一个中等规模的模型来研究RL与LLM的标度特性。图5说明了在数学提示集上训练的小模型变体的训练精度和响应长度在训练迭代中的演变。随着训练的进行,可以观察到响应长度和性能准确性同时增加。值得注意的是,更具挑战性的基准在响应长度上表现出更大的增长,这表明模型学会为复杂问题生成更精细的解决方案。图6显示了m输出上下文长度与其解决问题的能力之间的强相关性。最终运行的k1.5扩展到128k上下文长度,并观察到硬推理基准的持续改进。

作者将所提出的long 2short RL算法与DPO,最短拒绝采样和模型合并方法进行了比较,重点关注long 2short问题的令牌效率,特别是获得的长cot模型如何使短模型受益。如图7所示,与其他方法(如DPO和模型合并)相比,所提出的long 2short RL算法展示了最高的令牌效率。值得注意的是,k1.5系列中的所有型号(以橙色标记)与其他型号(以蓝色标记)相比,都表现出上级令牌效率。例如,k1.5-short w/ rl在AIME 2024上获得了60.8的Pass@1分数(平均8次运行),而平均只使用了3,272个令牌。类似地,k1.5-shortest在MATH 500上获得了88.2的Pass@1分数,同时消耗了与其他短模型大致相同数量的令牌。

4.总结

本文介绍了 Kimi k1.5,这是一个通过强化学习(RL)训练的多模态大语言模型(LLM),旨在通过扩展上下文长度和改进策略优化方法来提升模型的推理能力。文章详细描述了 Kimi k1.5 的训练过程,包括其 RL 训练技术、多模态数据配方和基础设施优化。Kimi k1.5 的训练方法不依赖于复杂的蒙特卡洛树搜索、价值函数或过程奖励模型,而是通过长上下文扩展和改进的策略优化方法,建立了一个简洁有效的 RL 框架。

文章指出,传统的语言模型预训练方法受限于可用的高质量训练数据量,而强化学习为继续提升 AI 的能力提供了新的方向。通过 RL,模型可以通过奖励机制学习探索,从而突破静态数据集的限制。Kimi k1.5 在多个基准测试中取得了与 OpenAI 的 o1 模型相当的性能,例如在 AIME 2024 中达到 77.5 的通过率,在 MATH 500 中达到 96.2 的准确率,在 Codeforces 中达到 94 百分位的排名,在 MathVista 中达到 74.9 的通过率。

此外,文章还介绍了如何将长推理链(Long-CoT)技术应用于短推理链(Short-CoT)模型的优化,通过长度惩罚和模型合并等方法,显著提升了短-CoT 模型的性能。例如,在 AIME 2024 中,Kimi k1.5 的短-CoT 模型达到了 60.8 的通过率,在 MATH 500 中达到了 94.6 的准确率,在 LiveCodeBench 中达到了 47.3 的通过率,这些结果均优于现有的短-CoT 模型,如 GPT-4o 和 Claude Sonnet 3.5。

文章还探讨了长上下文扩展的重要性,通过实验表明,随着上下文长度的增加,模型的性能持续提升。此外,文章介绍了 Kimi k1.5 的训练流程,包括预训练、监督微调(SFT)、长-CoT 监督微调和强化学习阶段。在强化学习部分,作者详细讨论了 RL 提示集的策划、长-CoT 监督微调以及 RL 训练策略。

文章还介绍了 Kimi k1.5 的训练基础设施,包括大规模 RL 训练系统的设计,以及部分展开技术(Partial Rollouts)如何优化长文本推理的处理。此外,文章还讨论了混合部署策略,这种策略允许训练和推理任务共享 GPU 资源,从而提高资源利用率和训练效率。

最后,文章总结了 Kimi k1.5 的训练方法和系统设计,强调了上下文长度扩展和改进策略优化方法的重要性,并展望了未来进一步提升长上下文 RL 训练效率和可扩展性的研究方向。

-

如果你觉得这些内容对你有帮助,或者让你眼前一亮,不妨点个赞、关注一下,或者收藏起来慢慢看!你的支持是我最大的动力,也让我知道这些内容对你有价值。感谢你的陪伴,我会继续努力带来更多有趣、有用的东西!🎉📚🚀

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值