Kimi K1.5:利用 LLM 规模化强化学习

25年1月来自月之暗面 Kimi 团队的论文“Kimi K1.5: Scaling Reinforcement Learning With LLMs”。

语言模型预训练与下一个 token 预测已被证明对于扩展计算有效,但受限于可用的训练数据量。规模化强化学习 (RL) 为人工智能的持续改进开辟新的方向,有望使大语言模型 (LLM) 能够通过学习带奖励的探索来扩展其训练数据。然而,之前发表的研究并没有产生有竞争力的结果。鉴于此,本文报告 Kimi k1.5 的训练实践,使用 RL 训练的多模态 LLM,包括其 RL 训练技术、多模态数据配方和基础设施优化。长上下文扩展和改进的策略优化方法,是该方法的关键要素,它建立一个简单、有效的 RL 框架,而无需依赖更复杂的技术,例如蒙特卡洛树搜索、价值函数和过程奖励模型。值得注意的是,该系统在多个基准和模态下均实现最先进的推理性能——例如,AIME 上 77.5、MATH 500 上 96.2、Codeforces 上 94 百分位、MathVista 上 74.9——与 OpenAI 的 o1 相匹配。此外,提出有效的 long2short 方法,使用长 CoT 技术来改进短 CoT 模型,产生最先进的短 CoT 推理结果——例如,AIME 上 60.8、MATH500 上 94.6、LiveCodeBench 上 47.3——大大优于现有的短 CoT 模型,例如 GPT-4o 和 Claude Sonnet 3.5(提升高达 +550%)。

如下两图所示分别是长CoT模型和短 CoT模型的结果:

请添加图片描述

请添加图片描述

在规模化定律的背景下,已经研究使用下一个token预测进行语言模型预训练,其中按比例规模化模型参数和数据大小可持续提高智能。(Kaplan,2020;Hoffmann,2022)然而,这种方法受限于可用的高质量训练数据量(Villalobos,2024;Muennighoff,2023)。

本报告介绍 Kimi k1.5 的训练方法,最新强化学习 (RL) 训练的多模态 LLM。目标是探索可能的新轴以继续扩展。使用带有 LLM 的 RL,模型可以学习使用奖励进行探索,因此不限于预先存在的静态数据集。

Kimi k1.5 的开发包括几个阶段:预训练、原始监督微调 (SFT)、长 CoT 监督微调和强化学习 (RL)。

强化学习提示集的挑选

强化学习提示集的质量和多样性对于确保强化学习的有效性起着至关重要的作用。精心构建的提示集不仅可以引导模型实现稳健的推理,还可以降低奖励黑客和过拟合那些肤浅模式的风险。具体来说,三个关键属性定义高质量的强化学习提示集:
• 多样化覆盖:提示应涵盖广泛的学科,例如 STEM、编码和一般推理,以增强模型的适应性并确保在不同领域的广泛适用性。
• 平衡难度:提示集应包括分布均匀的简单、中等和困难问题,以促进逐步学习并防止过拟合到特定的复杂程度。
• 准确的可评估性:提示应允许验证者进行客观可靠的评估,确保基于正确的推理而不是肤浅模式或随机猜测来衡量模型性能。

为了实现提示集的多样化覆盖,用自动过滤器来选择需要丰富推理且易于评估的问题。数据集包括来自各个领域的问题,例如 STEM 领域、竞赛和一般推理任务,包含纯文本和图像文本问答数据。此外,开发一个标签系统,按领域和学科对提示进行分类,确保在不同学科领域实现平衡表征(M. Li,2023;W. Liu,2023)。

采用基于模型的方法,利用模型自身的能力自适应地评估每个提示的难度。具体来说,对于每个提示,SFT 模型使用相对较高的采样温度生成十次答案。然后计算通过率并将其用作提示难度的智体——通过率越低,难度越高。这种方法使难度评估与模型的内在能力保持一致,使其对 RL 训练非常有效。通过利用这种方法,可以预先过滤大多数琐碎的情况,并在 RL 训练期间轻松探索不同的采样策略。

为了避免潜在的奖励黑客攻击(Everitt,2021;Pan,2022),需要确保每个提示的推理过程和最终答案都能得到准确验证。经验观察表明,一些复杂的推理问题可能有相对简单且容易猜到的答案,从而导致假正验证——模型通过错误的推理过程得出正确答案。为了解决这个问题,排除容易出现此类错误的问题,例如多项选择题、判断题和基于证明的问题。此外,对于一般的问答任务,提出一种简单而有效的方法来识别和删除容易被黑客攻击的提示。具体来说,提示模型猜测潜在答案,而无需任何 CoT 推理步骤。如果模型在 N 次尝试内预测正确答案,则认为该提示太容易被黑客攻击并被删除。设置 N = 8 可以删除大多数易于破解的提示。开发更高级的验证模型仍然是未来研究的一个开放方向。

长 CoT 监督微调

借助改进的 RL 提示集,采用提示工程来构建一个小而高质量的长 CoT 热身数据集,其中包含经过准确验证的文本和图像输入推理路径。这种方法类似于拒绝抽样 (RS),但侧重于通过提示工程生成长 CoT 推理路径。由此产生的热身数据集旨在封装对人类推理至关重要的关键认知过程,例如规划,其中模型在执行前系统地概述步骤;评估,涉及对中间步骤的批判性评估;反思,使模型能够重新考虑和改进其方法;探索,鼓励考虑替代解决方案。通过在这个热身数据集上执行轻量级 SFT,有效地为模型做好准备,以内化这些推理策略。因此,经过微调的长 CoT 模型表现出更好生成更详细和逻辑连贯响应的能力,从而提高其在各种推理任务中的表现。

强化学习

问题设置

给定一个训练数据集 D = {(x_i, y_i^∗ )},其中包含问题 x_i 和相应的真值答案 y_i^∗,目标是训练一个策略模型 π_θ 来准确解决测试问题。在复杂推理的背景下,问题 x 到解决方案 y 的映射并非易事。为了应对这一挑战,思维链 (CoT) 方法提出使用一系列中间步骤 z = (z_1, z_2, . . . , z_m) 来连接 x 和 y,其中每个 z_i 都是一个连贯的 token 序列,作为解决问题的重要中间步骤 (J. Wei,2022)。在解决问题 x 时,思维 z_t ∼ π_θ(·|x, z_1, . . . , z_t−1) 被自回归采样,然后是最终答案 y ∼ π_θ(·|x, z_1, . . . , z_m)。用 y, z ∼ π_θ 来表示此采样过程。注:想法和最终答案采样作为语言序列。

为了进一步增强模型的推理能力,规划算法被用于探索各种思维过程,在推理时产生改进的 CoT(Yao,2024;Y. Wu,2024;Snell,2024)。这些方法的核心见解,是明确构建由价值估计引导的思维搜索树。这使得模型能够探索思维过程的不同延续,或者在遇到死胡同时回溯以探索新方向。更详细地讲,假设 T 是一棵搜索树,其中每个节点代表一个部分解决方案 s = (x, z_1:|s|)。这里 s 由问题 x 和一系列指向该节点的思维 z_1:|s| = (z_1, . . . , z_|s|) 组成,|s| 表示该序列中的思维数。规划算法使用评论模型 v 来提供反馈 v(x, z_1:|s|),这有助于评估当前解决问题的进度并识别现有部分解决方案中的任何错误。反馈可以通过判别分数或语言序列提供(L. Zhang,2024)。在所有 s ∈ T 反馈的指导下,规划算法选择最有希望扩展的节点,从而扩大搜索树。上述过程不断重复,直到得出完整的解决方案。

还可以从算法的角度来研究规划算法。给定第 t 次迭代中可用的过去搜索历史 (s_1, v(s_1), … , s_t−1, v(s_t−1)),规划算法 A 迭代地确定下一个搜索方向 A(s_t|s_1, v(s_1), … , s_t−1, v(s_t−1)),并为当前搜索进度 A(v(s_t)|s_1, v(s_1), … , s_t) 提供反馈。由于想法和反馈都可以看作是中间推理步骤,并且这些组成部分都可以表示为语言token序列,因此使用 z 来替换 s 和 v 以简化符号。因此,将规划算法视为直接作用于推理步骤序列 A(·|z_1, z_2, … ) 的映射。在这个框架中,规划算法使用的搜索树中,存储的所有信息都被展平到算法的完整上下文中。这为生成高质量的 CoT 提供一个有趣的视角:不必明确构建搜索树并实现规划算法,而是可以训练一个模型来近似这个过程。在这里,思维(即语言token)的数量相当于传统上分配给规划算法的计算预算。长上下文窗口的最新进展促进训练和测试阶段的无缝可扩展性。如果可行,这种方法使模型能够通过自回归预测直接在推理空间上运行隐式搜索。因此,该模型不仅学会解决一组训练问题,而且还发展有效解决个别问题的能力,从而提高对未见过测试问题的泛化能力。

因此,考虑训练模型使用强化学习 (RL) (OpenAI 2024) 生成 CoT。让 r 成为一个奖励模型,该模型基于真值 y 证明针对给定问题 x 提出答案 y 的正确性,通过分配一个值 r(x, y, y) ∈ {0, 1}。对于可验证的问题,奖励直接由预定义的标准或规则确定。例如,在编码问题中,评估答案是否通过测试用例。对于具有自由形式真值的问题,训练一个奖励模型 r(x, y, y^∗),该模型预测答案是否与真值相匹配。给定一个问题 x,模型 π_θ 通过采样过程 z ∼ π_θ(·|x), y ∼ π_θ(·|x,z) 生成 CoT 和最终答案。生成的 CoT 的质量通过它是否可以得出正确最终答案来评估。

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

策略优化

采用在线策略镜像下降的变型作为训练算法(Abbasi-Yadkori 2019;Mei 2019;Tomar 2020)。该算法以迭代方式执行。在第 i 次迭代中,用当前模型 π_θ_i 作为参考模型,并优化一个相对熵正则化的策略优化问题。这个目标有一个闭式解。基于所定义约束的满足,在优化过程中利用离线-策略数据,其激发一种替代损失。

为了近似 τ logZ,采用一组样本 (y_1, z_1),…,(y_k, z_k)。使用采样奖励的经验平均值 r = mean(r(x, y_1, y), . . . , r(x, y_k, y^∗)) 可产生有效的实际结果。这是合理的,因为当 τ → ∞ 时,τ log Z 接近 π_θ_i 的预期奖励。最后,采取替代损失的梯度,结束该学习算法。对于每个问题 x,使用参考策略 π_θ_i 对 k 个响应进行采样。

对于熟悉策略梯度方法的人来说,这个梯度类似于使用采样奖励平均值作为优化问题的策略梯度(Kool,2019;Ahmadian,2024)。主要区别在于响应是从 π_θ_i 而不是在策略中采样的,并且应用 l2 正则化。因此,可以将其视为一个通常在线策略正则化的策略梯度算法到离线策略情况的自然延伸(Nachum,2017)。从 D 中抽取一批问题并将参数更新为 θ_i+1,随后将其作为下一次迭代的参考策略。由于每次迭代都会考虑由于参考策略变化而产生的不同优化问题,因此还会在每次迭代开始时重置优化器。

在训练系统中排除价值网络,这在以前的研究中也曾被利用过(Ahmadian,2024)。虽然这种设计选择显著提高了训练效率,但也假设在经典 RL 中通过传统的价值函数进行信誉分配可能不适合上下文。考虑这样一种情况,模型生成一个部分 CoT(z_1、z_2、…、z_t),并且有两个潜在的下一个推理步骤:z_t+1 和 z’_t+1。假设 z_t+1 直接导致正确答案,而 z’_t+1 包含一些错误。如果可以访问 oracle 价值函数,它将表明 z_t+1 保留比 z’_t+1 更高的值。根据标准信誉分配原则,选择 z’_t+1 将受到惩罚,因为它相对于当前策略具有负面优势。然而,探索 z’_t+1 对于训练模型生成一个长 CoT 非常有价值。通过使用从长 CoT 得出最终答案的理由作为奖励信号,只要模型成功恢复并得出正确答案,它就可以从 z’_t+1 中学习反复试验的模式。这个例子的关键是,应该鼓励模型探索不同的推理路径,以增强其解决复杂问题的能力。这种探索性方法产生丰富的经验,支持关键规划技能的发展。主要目标不仅限于在训练问题上获得高精度,还专注于为模型配备有效的问题解决策略,最终提高其在测试问题上的表现。

长度惩罚

有一种过度思考现象,即在强化学习训练期间,模型的响应长度显著增加。虽然这会带来更好的性能,但在训练和推理过程中,过长的推理过程代价高昂,而过度思考往往不是人类所喜欢的。为了解决这个问题,引入长度奖励来抑制 token 长度的快速增长,从而提高模型的 token 效率。给定问题 x 的 k 个采样响应 (y_1, z_1), . . . , (y_k, z_k),真实答案为 y^∗,让 len(i) 为 (y_i, z_i) 的长度,min_len = min_i len(i),max_len = max_i len(i)。如果 max_len = min_len,将所有响应的长度奖励设置为零,因为它们的长度相同。

本质上鼓励较短的回答,并惩罚正确答案中的较长回答,同时明确地惩罚错误答案的长回答。然后用加权参数将基于长度的奖励添加到原始奖励中。

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

采样策略

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

  1. 课程采样。从较简单的任务开始训练,逐渐进展到更具挑战性的任务。由于初始强化学习模型的性能有限,在非常困难的问题上,花费有限的计算预算通常会产生很少的正确样本,从而导致训练效率降低。同时,收集的数据自然包含等级和难度标签,使基于难度的采样成为一种直观有效提高训练效率的方法。
  2. 优先抽样。除了课程抽样外,还使用优先抽样策略来关注模型表现不佳的问题。跟踪每个问题 i 的成功率 s_i,并以 1 − s_i 的比例抽样问题,以便成功率较低的问题获得更高的抽样概率。这将模型的努力引向其最薄弱的领域,从而加快学习速度并提高整体性能。

有关训练配方的更多详细信息

编码测试用例生成。由于网络上没有针对许多编码问题的测试用例,因此设计一种方法来自动生成测试用例,作为使用 RL 训练模型的奖励。主要关注不需要特殊判断的问题。还假设这些问题有基本解决方案,以便可以利用这些解决方案生成更高质量的测试用例。

利用广受认可的测试用例生成库 CYaRon 来增强该方法。用基础 Kimi k1.5 根据问题陈述生成测试用例。CYaRon 的使用说明和问题描述作为生成器的输入。对于每个问题,首先使用生成器生成 50 个测试用例,并为每个测试用例随机抽取 10 个真值提交。根据提交运行测试用例。如果 10 个提交中至少有 7 个产生匹配结果,则测试用例被视为有效。经过这一轮筛选,获得一组选定的测试用例。如果 10 个提交中至少有 9 个通过整个选定的测试用例集,则问题及其相关的选定测试用例将添加到训练集中。

在统计方面,从 1,000 个在线竞赛问题样本中,大约 614 个不需要特殊评判。开发 463 个测试用例生成器,至少生成 40 个有效测试用例,从而将 323 个问题纳入训练集。

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

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

在人工抽检中,Classic RM 的准确率约为 84.4,而思维链 RM 的准确率达到 98.5。在强化学习训练过程中,采用 思维链 RM 来确保反馈更加正确。

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

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

每种类型的数据对于构建全面的视觉语言模型都至关重要,该模型可以有效地管理广泛的实际应用程序,同时确保在各种输入模态中保持一致的性能。

Long2short:短 CoT 模型的上下文压缩

虽然长 CoT 模型实现了强大的性能,但与标准短 CoT LLM 相比,它消耗了更多的测试时间 token。但是,可以将思维先验从长 CoT 模型转移到短 CoT 模型,这样即使在有限的测试时间 token 预算下也可以提高性能。针对这个 long2short 问题提出几种方法,包括模型合并(Yang,2024)、最短拒绝采样、DPO(Rafailov,2024)和 long2short RL。

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

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

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

Long2short RL。在标准 RL 训练阶段之后,选择一个在性能和 token 效率之间提供最佳平衡的模型作为基础模型,并进行单独的 long2short RL 训练阶段。在第二阶段,应用长度惩罚,并显著减少最大推出长度,以进一步惩罚超过所需长度但可能正确的响应。

其他训练细节

预训练

Kimi k1.5 基础模型是在多样化、高质量的多模态语料库上进行训练的。语言数据涵盖五个领域:英语、中文、代码、数学推理和知识。多模态数据(包括字幕、图文交错、OCR、知识和 QA 数据集)使模型能够获得视觉语言能力。严格的质量控制确保整个预训练数据集的相关性、多样性和平衡性。预训练分为三个阶段:(1)视觉语言预训练,建立强大的语言基础,然后逐步进行多模态集成;(2)冷却,使用精选和合成数据巩固能力,特别是对于推理和基于知识的任务;(3)长上下文激活,将序列处理扩展到 131,072 个tokens。

原始监督微调

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

原始 SFT 数据集包含大约 100 万个文本示例。具体来说,50 万个示例用于一般问答,20 万个用于编码,20 万个用于数学和科学,5000 个用于创意写作,2 万个用于长上下文任务,如摘要、文档问答、翻译和写作。此外,还构建 100 万个文本视觉示例,涵盖图表解释、OCR、基于图像的对话、视觉编码、视觉推理以及带有视觉辅助的数学/科学问题等各种类别。

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

RL 基础设施

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

在人工智能领域,强化学习 (RL) 已成为大语言模型 (LLM) 的关键训练方法 (Ouyang 2022) (Jaech 2024),其灵感来自于通过 AlphaGo (Silver 2017)、AlphaStar (Vinyals 2019) 和 OpenAI Dota Five (Berner 2019) 等系统成功掌握围棋、星际争霸 II 和 Dota 2 等复杂游戏。沿袭这一传统,Kimi k1.5 系统采用迭代同步 RL 框架,经过精心设计,可通过持续学习和适应来增强模型的推理能力。该系统的一项关键创新是引入部分展开技术,旨在优化复杂推理轨迹的处理。

如图 a 所示,RL 训练系统采用迭代同步方法运行,每次迭代都包含一个展开阶段和一个训练阶段。在展开阶段,由中央主机协调的展开工作器通过与模型交互生成展开轨迹,产生对各种输入的响应序列。然后,这些轨迹被存储在重放缓冲区中,通过破坏时间相关性,确保训练数据集的多样性和无偏性。在随后的训练阶段,训练工作器访问这些经验来更新模型的权重。这个循环过程允许模型不断从其行为中学习,随着时间的推移调整其策略以提高性能。

请添加图片描述

中央主机充当中央指挥者,管理展开工作器、训练工作器、奖励模型评估和重放缓冲区之间的数据流和通信。它确保系统和谐运行,平衡负载并促进高效的数据处理。

训练器可以访问这些展开轨迹(无论是在一次迭代中完成还是在多次迭代中划分),以计算梯度更新,从而改进模型的参数并提高其性能。这个过程由奖励模型监督,该模型评估模型输出的质量并提供必要的反馈来指导训练过程。奖励模型的评估对于确定模型策略的有效性以及引导模型实现最佳性能尤为重要。

此外,该系统还包含一个代码执行服务,该服务专门用于处理与代码相关的问题,是奖励模型不可或缺的一部分。该服务在实际编码场景中评估模型的输出,确保模型的学习与现实世界的编程挑战紧密结合。通过根据实际代码执行验证模型的解决方案,这个反馈循环对于改进模型的策略和提高其在代码相关任务中的性能至关重要。

长 CoT RL 的部分展开

工作的主要想法之一是扩展长上下文 RL 训练。部分展开是一项关键技术,它通过管理长轨迹和短轨迹的展开,有效地解决处理长 CoT 特征的挑战。该技术建立一个固定的输出token预算,限制每个展开轨迹的长度。如果轨迹在展开阶段超出token限制,则未完成的部分将保存到重放缓冲区并在下一次迭代中继续。它确保没有任何一条长轨迹独占系统资源。此外,由于展开工作器异步运行,当一些工作器参与长轨迹时,其他工作器可以独立处理新的、较短的展开任务。异步操作通过确保所有展开工作器都积极参与训练过程来最大限度地提高计算效率,从而优化系统的整体性能。

如上图 b 所示,部分展开系统通过将长响应分解为迭代中的段(从 n-m 到 n)来工作。重放缓冲区充当维护这些响应段的中央存储机制,其中只有当前迭代(n)需要基于策略的计算。先前的段(n-m 到 n-1)可以从缓冲区中高效重用,无需重复展开。这种分段方法显著减少计算开销:系统不是一次性推出整个响应,而是逐步处理和存储段,从而允许生成更长的响应,同时保持快速迭代时间。在训练期间,可以将某些段排除在损失计算之外,以进一步优化学习过程,使整个系统既高效又可扩展。

部分展开的实施还提供重复检测。系统识别生成内容中的重复序列,并提前终止它们,从而减少不必要的计算,同时保持输出质量。检测的重复可以被分配额外的惩罚,从而有效地阻止提示集中的冗余内容生成。

训练和推理的混合部署

RL 训练过程包括以下阶段:
• 训练阶段:一开始,Megatron(Shoeybi,2020)和 vLLM(Kwon,2023)在单独的容器中执行,由称为检查点引擎的填充进程封装。Megatron 开始训练过程。训练完成后,Megatron 卸载 GPU 内存并准备将当前权重传输到 vLLM。
• 推理阶段:在 Megatron 卸载之后,vLLM 从虚拟模型权重开始,通过 Mooncake 从 Megatron 传输的最新权重更新它们(Qin,2024)。完成部署后,检查点引擎将停止所有 vLLM 进程。
• 后续训练阶段:一旦分配给 vLLM 的内存被释放,Megatron 就会加载内存并启动另一轮训练。

现有工作很难同时支持以下所有特征。
• 复杂的并行策略:Megatron 可能与 vLLM 具有不同的并行策略。Megatron 中分布在多个节点中的训练权重可能难以与 vLLM 共享。
• 最小化空闲 GPU 资源:对于 On-Policy RL,最近的工作,例如 SGLang(L. Zheng,2024)和 vLLM),可能会在训练过程中保留一些 GPU,这反过来可能会导致空闲的训练 GPU。在训练和推理之间共享相同的设备会更有效率。
• 动态扩展能力:在某些情况下,通过增加推理节点的数量同时保持训练过程不变,可以实现显着的加速。系统可以在需要时有效利用空闲的 GPU 节点。

如图所示,在 Megatron 和 vLLM 之上实现这个混合部署框架,从训练到推理阶段只需不到一分钟,反之则大约十秒。

请添加图片描述

混合部署策略。提出一种用于训练和推理任务的混合部署策略,该策略利用共享所有可用 GPU 的 Kubernetes Sidecar 容器将两个工作负载共置在一个 pod 中。该策略的主要优点是:
• 它促进高效的资源共享和管理,当两个节点都部署在不同的节点上时,防止训练节点在等待推理节点时处于空闲状态。
• 利用不同的部署映像,训练和推理可以独立迭代以获得更好的性能。
• 该架构不仅限于 vLLM,还可以方便地集成其他框架。

检查点引擎。检查点引擎负责管理 vLLM 流程的生命周期,公开 HTTP API,以便在 vLLM 上触发各种操作。为了实现整体一致性和可靠性,用由 etcd 服务管理的全局元数据系统来广播操作和状态。

由于 CUDA 图、NCCL 缓冲区和 NVIDIA 驱动程序,通过 vLLM 卸载完全释放 GPU 内存可能具有挑战性。为了最大限度地减少对 vLLM 的修改,会在需要时终止并重新启动它,以提高 GPU 利用率和容错能力。

Megatron 中的工作程序将拥有的检查点转换为共享内存中的 Hugging Face 格式。此转换还考虑流水线并行性和专家并行性,以便这些检查点中只保留张量并行性。共享内存中的检查点,随后被分成分片并在全局元数据系统中注册。用 Mooncake 通过 RDMA 在对等节点之间传输检查点。需要对 vLLM 进行一些修改以加载权重文件并执行张量并行性转换。

代码沙盘

开发沙盘作为执行用户提交代码的安全环境,并针对代码执行和代码基准评估进行优化。通过动态切换容器镜像,沙盘通过 MultiPL-E(Cassano、Gouwar、D. Nguyen、S. Nguyen 等,2023 年)、DMOJ Judge Server、Lean、Jupyter Notebook 和其他镜像支持不同的用例。

对于编码任务中的强化学习,沙盘通过提供一致且可重复的评估机制来确保训练数据判断的可靠性。其反馈系统支持多步评估,例如代码执行反馈和存储库级别编辑,同时保持统一的上下文,以确保跨编程语言的公平和公正的基准比较。

在 Kubernetes 上部署服务以实现可扩展性和弹性,并通过 HTTP 端点将其公开以供外部集成。Kubernetes 的自动重启和滚动更新等功能可确保可用性和容错性。

为了优化性能并支持 RL 环境,在代码执行服务中采用多种技术来提高效率、速度和可靠性。这些包括:
• 使用 Crun:用 crun 作为容器运行时而不是 Docker,从而显著缩短容器启动时间。
• Cgroup 重用:预先创建用于容器的 cgroup,这在高并发场景中至关重要,因为为每个容器创建和销毁 cgroup 可能会成为瓶颈。
• 磁盘使用优化:使用上层挂载为 tmpfs 的覆盖文件系统来控制磁盘写入,提供固定大小的高速存储空间。这种方法对短暂的工作负载很有帮助。

请添加图片描述

这些优化提高 RL 在代码执行中的效率,为评估 RL 生成的代码提供一致且可靠的环境,这对于迭代训练和模型改进至关重要。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值