Kimi 1.5 技术报告解读

文摘   2025-01-21 16:00   北京  

最近简直是提前过年了,就在昨晚,DeepSeek和Kimi都发布了自己的o1版,Kimi还首先发布了自己的技术报告,有点意思… 提起Kimi,大家一直有的印象是,有技术先发优势,最早把长文本吃透,产品化领先。他们发布的第一个技术报告,肯定要仔细学习一番。

刚准备动手的时候,发现大神Jim Fan也在看,他是这么评价的:

首先看效果,一句话是多模态长推理里面最强的,并且能把模型推理做短,且性能不明显下降。这也侧面说明Kimi1.5系列潜力无穷,很难想象他的纯文本,彻底的infrence scaling起来效果得多炸裂。

看榜单的话,在长推理上,Kimi 1.5在数学,纯文本,多模态上,都可以都远超openai的o1模型;在codeforces评测上与其持平,LiveCode评测上表现上略差,但相比于QVQ和QwQ有较大的优势。

在短推理上,得益于他们搞的一个核心秘密武器,成为Kimi1.5的甜点能力,整个榜单都全面领先。数学能力无论是gpt-4o还是claude3.5-sonnet都远不如Kimi 1.5,尤其是在AIME榜单上,Kimi 1.5有60.8,而最高的deepseek-v3只有39.2,堪称断层式碾压。并且Kimi 1.5在其他场景也基本是一流水平,超过了目前的开源模型。

最最重要的是技术报告的原文,发布出来还不到24个小时,在Kimi的帮助下,把它仔细读了一遍。然后不得不感慨,Kimi是有高人的。

报告地址:https://github.com/MoonshotAI/Kimi-k1.5/blob/main/Kimi_k1.5.pdf

Kimi1.5开头核心归纳了自己的4大创新点:

1. 长上下文扩展 2. 改进的策略优化 3. 简化框架 4. 多模态

如果你觉得,就这?那就大错特错了,围绕这4大类,至少写了30多个非常有讲究的trick,都是让大家非常痛苦和纠结的点。

比如,o系列技术路线一定要搞那么长的cot吗?

mcts和prm吹那么响,似乎钦定了o系列了技术路线,我怎么就做不出效果来?

以及强化学习的infra怎么做?

这篇技术报告写的可以加细致入微了。

比如这里long2short,把探索过程的细节基本都写出来了,一些特别细的地方:

模型合并:之前都是通过模型合并来提高模型的泛化性,k1.5发现long-cot模型和short-cot模型也可以合并,从而提高输出效率,中和输出内容,并且无需训练。

最短拒绝采样:对于模型输出结果进行n次采样(实验中n=8),选择最短的正确结果进行模型微调。

DPO:与最短拒绝采样类似,利用long-cot模型生成多个输出结果,将最短的正确输出作为正样本,而较长的响应(包括:错误的长输出、比所选正样本长 1.5 倍的正确长输出)作为负样本,通过构造的正负样本进行DPO偏好学习。

Long2Short的强化学习:在标准的强化学习训练阶段之后,选择一个在性能和输出效率之间达到最佳平衡的模型作为基础模型,并进行单独的long-cot到short-cot的强化学习训练阶段。在这一阶段,采用长度惩罚,进一步惩罚超出期望长度,但保证模型仍然可能正确。

long2short部分在实验中也证明了,这个可以用极小的代价,极大地提高推理效率,在输出token长度和准确率上,达到了均衡的最优解。

mcts和prm到底需不需要?

Kimi1.5和DeepSeek的技术报告都提到了这个事儿,Kimi也更大胆在ppo上进一步做简化,其中的value network也不需要。要知道,把技术做花哨是显得你牛逼,但大家实践来实践去并不一定真的有用,遵从奥卡姆剃刀原则,把模型做work的拿出来那才叫真本事。

在RL部分Kimi仔细讲解了他们强化学习的设置和推演过程,以及让大家感到头痛的长度问题怎么解决,各种采样设置的细节,数学和代码RM建模的细节。

在prompt集的设定上,Kimi 1.5遵循了3个原则:

多样化覆盖:prompt应涵盖广泛的学科,如 STEM、编码和一般推理,以增强模型的适应性并确保在不同领域的广泛应用。

平衡难度:prompt应包含易、中、难问题的良好分布,以促进逐步学习并防止过拟合到特定复杂性水平。

准确可评估性:prompt应允许验证者进行客观且可靠的评估,确保模型性能是基于正确推理而非特定模式或随机猜测来衡量的。(这就是为啥不依赖MCTS)。

在核心的RL设定上,应了一种在线策略梯度下降的变体(简化了PPO)作为Kimi 1.5的训练算法。算法迭代执行。在第i次迭代中,使用当前模型作为参考模型,并优化以下相对熵正则化策略优化问题:

其中τ > 0是控制正则化程度的参数。此目标有一个闭式解:

这里

是归一化因子。对两边取对数,得到对于任何(y, z)满足以下约束,能够在优化过程中利用离策略数据:

可以推导出下面的损失函数:

为了逼近,使用样本

技术报告中还发现,使用采样奖励的经验均值 产生有效的实际结果。这是合理的,因为时接近下的期望奖励。最后,通过优化上面的损失函数来优化这个RL算法。对于每个问题x,从参考策略采样k个响应,梯度由以下公式给出:

实际上,除了核心的RL算法设定,技术报告中还提到了一些RL中大家头疼的问题,比如reward hacking,长度惩罚设定,采样策略,RM建模细节,多模态数据的融合等,值得一看。

最后不得不提的是infra部分。

infra成为下半年最火和最牛的技术方向,主要是这块大家都没有太多的参考经验,并且一般的框架为了高度的复用性,不会考虑做特别深入和细致的优化。如果说算法的优化是加法的话,infra的优化就是乘法,有很强的倍增效应。 在训练效率,稳定性,实验效率上的改进,都可以直接映射为降低的卡时成本。别家的算法工程师不睡觉才能卷出来的效果,infra大佬们直接优化到安安心心睡大觉。更别提在模型scaling到一定程度的时候,只有infra才能帮助传统的算法工程师把实验跑起来了。

如上图所示,RL infra整体通过迭代同步的方式运行,每个迭代包括一个轨迹生成(rollout)阶段和一个训练阶段。在轨迹生成阶段,由中央主控单元(master)协调的轨迹生成工作单元(rollout workers)通过与模型的交互生成轨迹,产生对各种输入的响应序列。这些轨迹随后被存储在一个回放缓冲区(replay buffer)中,该缓冲区通过打破时间相关性来确保训练数据的多样性和无偏性。在随后的训练阶段,训练工作单元(trainer workers)访问这些经验以更新模型的权重。这一循环过程使模型能够持续从其行为中学习,随着时间的推移调整其策略以提升性能。

中央主控单元作为整个系统的指挥中心,管理着轨迹生成工作单元、训练工作单元、奖励模型评估以及回放缓冲区之间的数据流动和通信。它确保系统协调运行,平衡负载并促进高效的数据处理。训练工作单元访问这些轨迹,无论是单次迭代完成的轨迹,还是分散在多次迭代中的轨迹,以计算梯度更新,优化模型参数并提升性能。这一过程由奖励模型监督,奖励模型评估模型输出的质量并提供反馈以指导训练过程。奖励模型的评估对于确定模型策略的有效性以及引导模型向最优性能发展至关重要。

系统还整合了一个代码执行服务,专门用于处理与代码相关的问题,是奖励模型的重要组成部分。该服务在实际编程场景中评估模型的输出,确保模型的学习与现实世界的编程挑战紧密对齐。通过将模型的解决方案与实际代码执行进行验证,这一反馈循环对于精炼模型策略和提升其在代码相关任务中的性能至关重要。

具体到混合部署的infra构建上,拆了几个阶段:

训练阶段:Megatron(Shoeybi et al. 2020)和 vLLM(Kwon et al. 2023)分别在独立的容器中运行,这些容器被一个名为检查点引擎(checkpoint-engine)的外壳进程封装(详见第 2.6.3 节)。Megatron 首先启动训练过程。训练完成后,Megatron 会释放 GPU 内存,并准备将当前权重传递给 vLLM。

推理阶段:Megatron和 vLLM分别在独立的容器中运行,容器称为checkpoint-engine的外壳进程封装。Megatron 首先启动训练过程,训练完成后,Megatron 会释放 GPU 内存,并准备将当前权重传递给 vLLM。在 Megatron 释放内存后,vLLM 以虚拟模型权重启动,并通过 Mooncake 从 Megatron 接收最新的权重更新。完成回放后,checkpoint-engine会停止所有 vLLM 进程。

后续训练阶段:释放 vLLM 所占用的内存后,Megatron 重新加载内存并开始下一轮训练。

而现有框架很难同时满足以下所有特性:

复杂的并行策略:Megatron 和 vLLM 可能采用不同的并行策略。Megatron 中分布在多个节点上的训练权重很难与 vLLM 共享。

最小化闲置 GPU 资源:对于在线策略强化学习,SGLang和 vLLM可能会在训练过程中保留一些 GPU,导致训练 GPU 的闲置。如何共享相同的设备,将训练变得更加高效。

动态扩展能力:通过增加推理节点的数量,同时保持训练过程不变,可以显著加速训练。如何高效利用闲置的 GPU 节点。

最后,Kimi1.5现在也正在进行灰度上线,你们可以看看自己的Kimi界面有没有,有的话就可以提前体验到k1.5的整体效果。

本技术报告由Kimi辅助阅读。

思源数据科学
Towards AGI
 最新文章