搞不懂CUDA的人有救了，Devin开发商开源Kevin，强化学习生成CUDA内核
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
机器之心报道
编辑：蛋酱、泽南
本周三，知名AI创业公司，曾发布「全球首个AI软件工程师」的CognitionAI开源了一款使用强化学习，用于编写CUDA内核的大模型Kevin-32B。
Kevin-32B基于QwQ-32B在KernelBench数据集上使用GRPO进行了多轮强化学习训练，实现了超越o3和o4-mini的顶级推理表现。
对此，机器学习社区表现出了极大的兴趣。有人表示期待DeepSeekR1风格的训练方法用来提升代码效率已久，这回终于有人站出来了。
在一篇博客中，CognitionAI详细介绍了新模型强化学习训练的机制。
代码是一个不断迭代的过程——需要我们编写、执行程序，评估结果，并根据反馈优化代码。大语言模型（LLM）在代码生成方面的最新进展尝试将此过程融入推理阶段，并使用并行采样等方法。虽然这些方法是有效的，但它们依赖于搜索而非实际学习——在这其中模型权重被冻结。
CognitionAI探索了多轮强化学习，使用来自环境的中间反馈，并屏蔽模型思维以避免在多轮训练中上下文爆炸。
他们提出的模型Kevin-32B（即KernelDevin）在内核生成方面优于前沿推理模型。此外他们通过实验结果表明，与单轮训练相比，多轮训练使模型在自我优化方面更有效。
多轮训练方法
在新模型上，作者使用了KernelBench，这是一个包含250个基于PyTorch的经典深度学习任务的数据集。它衡量优化CUDA内核替换PyTorch算子的能力。训练专注于前两个级别，每个级别包含100个任务。级别1包含矩阵乘法、卷积和损失函数等基础任务，级别2则包含融合算子，再使用这两个级别的180个任务进行训练，并留出20个任务作为保留集。
在训练过程中，模型会经历一个迭代反馈循环：从生成的内核中提取反馈，并再对其进行优化。如果内核编译失败，错误轨迹会传递给模型并要求其修复。如果编译正确，系统则会测量运行时间并要求模型进一步改进。
初始方法的机制如下。从初始提示开始，先在每个优化步骤后附加思路链、内核和评估信息，再为整个轨迹分配一个奖励（定义为任何内核获得的最高分数），并使用该序列进行训练。
不过，这种方法存在一些挑战：
上下文窗口爆炸：推理模型会生成很长的思维链。使用这种方法，轨迹的长度在几次传递后很容易达到5-10万个token，这会导致训练不便；
样本效率低下和credit分配问题：即使我们生成了多个内核，仍然只为整个轨迹分配一个奖励。这无法表明哪个细化步骤真正提高了正确性或性能。奖励应该根据细化步骤对最终结果的贡献来分配。
为了解决上下文长度爆炸的问题，Kevin-32B丢弃了轨迹中最长的部分——思维链。现在每个token将只包含先前生成的内核和评估结果。为了仍然保留上一步思考过程的信息，我们特意要求模型生成其自身思考过程的简要摘要，然后将其传递给后续上下文。
移除推理的思路链。
为了解决样本效率低下的问题，Kevin-32B选择了一个更具表现力的奖励函数，将内核的细化建模为马尔可夫决策过程，将给定响应的奖励设置为当前内核与所有后续内核得分的折扣总和。至此，每个细化步骤都变成了一个训练样本。
将奖励设为分数的折扣总和。
结果
对于每项任务，作者并行采样了16条轨迹，并进行8个连续的细化步骤。对于每项任务的正确性或性能，他们将best@16定义为所有轨迹的最大值，avg@16定义为所有轨迹的平均值。
经过8个优化步骤，Kevin-32B在整个数据集上的平均正确率为65%，显著超越了QwQ-32B和前沿模型。它解决了89%的数据集，而o4-mini和o3分别只解决了53%和51%。在整个数据集中，Kevin-32B实现了1.41倍的best@16加速比，优于前沿模型。
Kevin-32B在二级任务上尤其有效，平均正确率为48%（o4-mini为9.6%，o3为9.3%）。这表明多轮训练可以提高模型解决更具挑战性、更长期任务的能力。同样，我们能注意到模型实现了1.74倍的best@16加速比（是o4-mini和o3为1.2倍）。
多轮训练vs单轮训练
Kevin-32B也展现出比QwQ-32B和单轮训练模型显著的提升。在4个优化步骤下，Kevin-32B的表现略优于单轮模型，但随着优化步骤增加到8个，两者之间的差距进一步扩大。这表明，通过鼓励更积极的优化，多轮训练在串行轴上具有更好的扩展性。
我们可能会想，单轮训练模型是否可以通过采样并行轨迹来实现更好的加速，然而在这种环境下并非如此。在计算预算固定的情况下，即使对于单轮训练模型，多轮推理也比单轮推理更占优势。
奖励黑客攻击
最初的实验使用了规模较小的模型，例如DeepSeek-R1-Distill-Qwen-7B，这导致了多起奖励黑客攻击事件：
该模型只是复制了PyTorch参考实现，因此因生成正确答案而获得奖励，且加速比提高了1.0倍；
该模型将错误的CUDA内核实现包装在try-except语句中，并调用PyTorch实现函数作为回退；
该模型继承自参考实现，绕过了对CUDA实现的需求。
奖励黑客攻击示例
为了防止奖励黑客攻击，作者对响应施加了更严格的格式检查。对于使用PyTorch函数或不包含CUDA内核的响应，奖励被设置为0。
我们能观察到，当模型能力与数据集难度之间的差距很大时，就会发生奖励黑客攻击。模型难以解决任何任务，因此当它生成一个被黑客入侵的内核时，它是唯一具有正向优势的操作，因此会得到显著的强化。
无意义和重复的生成
多次运行后，我们可以观察到在大约35-40步，模型开始生成重复或无意义的响应。假设这是因为模型偏离了稳定区域，令人惊讶的是，作者偶然发现了一个可以有效预测未来垃圾响应的指标——思维链不是以「OK」开头的响应比例——我们称之为「NotOkayRatio」。
模型训练越多，CoT的开头就越不稳定：
“OkayAmigos,soIneedtooptimizethis3Dtensor-matrixmultiplication..”
“OkayHolycrap,Ineedtogetthiscodeoptimized….”
在精炼步骤的第8遍时：
“OkaySHIT,thisisgettingfrustrating.Letmesee.Theerrorisabout…”
为了解决这个问题，作者尝试使用KL系数0.001和0.01进行运行，但发现这会减慢学习速度，同时又不能防止垃圾数据的产生。最后，作者通过使用Dr.GRPO提出的恒定长度损失正则化（这显著降低了梯度范数），并将梯度范数大幅削减至0.05，将垃圾数据的出现延迟到了第100步。
示例内核
作为示例，以1级任务40（LayerNorm）的轨迹为例，作者重点介绍了模型如何迭代地整合来自内核评估的反馈，并利用这些反馈来改进其生成过程。
步骤1：模型生成一个正确的内核。它使用共享内存，并将均值/方差计算与归一化融合。这实现了0.6倍的加速。
步骤2：模型从评估反馈中注意到「当前的加速比是0.6，这意味着它实际上比PyTorch版本慢。这不太好。让我想想瓶颈可能在哪里」。它将大部分推理token用于寻找更合适的块大小，并相应地修改代码。修改是正确的，加速比现在达到了4.27倍。
步骤3和4：模型尝试进行更积极的优化。它考虑了循环展开和减少库冲突，但最终选择了Warp级内联函数。但最终失败了——首先是因为归约逻辑存在正确性问题，其次是因为思路太长。
步骤5：模型注意到Warp归约实现中存在错误的问题并进行了修复。它成功实现了两级Warp归约。最终加速比达到了9.61倍。
训练设置
这里使用的是组相对策略优化算法（GRPO），该算法由DeepSeek提出，是近端策略优化算法（PPO）的一种变体。
GRPO不使用值网络来估算基线和计算优势，而是将从同一提示中采样的响应组内的奖励归一化。
作者使用vLLM进行推理，使用DeepSpeedZero-3卸载优化器状态。每批训练8个任务，每个任务16个轨迹。我们使用GRPO，每批2个梯度步骤。基础模型是QwQ-32B。
响应生成完成后，每个GPU将其vLLM引擎卸载到CPU内存，并评估其生成的内核。对于每个响应，作者都会检查响应格式是否正确，并提取CUDA内核。然后，编译并执行代码，使用随机张量测试其正确性。如果正确，则会对内核的运行时间进行剖析。
通过正确性检查的响应将获得0.3的奖励，而额外的性能奖励则相当于与参考实现相比所获得的速度提升。
内核评估和基准问题
作者对评估进行了沙盒化处理，以免出现致命错误（如CUDA非法内存访问），导致训练过程崩溃。
由于KernelBench中的许多任务使用非常小的输入张量，该基准最终测量的是内核启动开销而非实际内核执行时间。为了解决这个问题，作者扩大了受影响任务的张量维度。
KernelBench评估工具中还有一个更隐蔽的错误，导致被测内核将参考实现的输出张量作为自己的张量输出循环使用。因此，只计算（正确）部分输出张量的内核仍能通过正确性检查。为了解决这个问题，我们首先运行测试过的内核，然后再运行参考实现，从而避免了这一问题。
单轮训练设置
我们能观察到，单轮模型比基础模型有明显改善，但在25步之后，奖励开始趋于稳定。
多轮训练设置
作者根据每个轨迹的最佳内核对其正确性和性能进行评分。在最终训练运行中，每个前向路径由16个平行轨迹组成，每个轨迹有4个细化步骤，折扣系数为0.4。与单轮训练不同的是，现在的奖励会稳步增加。
随着模型学会更高效地使用推理token来生成内核，响应长度最初会减少。第25步之后，随着模型尝试更复杂的解决方案，响应长度会增加。按照DeepScaleR的做法，作者在第30步将最大响应长度从16K扩展到22Ktoken。
更多结果&消融研究
推理时间缩放
作者还研究了多轮模型在并行和串行轴上的扩展。在推理时间的第一次实验中，其使用了16个并行轨迹和8个细化步骤。作者再次发现，随着细化步骤的增加，多轮模型的扩展性更好。
在第二次实验中，作者将并行轨迹的数量增加到64个，同时只保留4个细化步骤。这样，best@64的正确率达到89.5%，性能提高了1.28倍，比8个细化步骤的best@16稍差。
作者研究了沿并行轴或串行轴缩放推理的效果。此处使用pass@k性能来表示k生成的估计性能。他们使用类似于Chen等人的无偏估计器计算该指标，其方差低于avg@k。
然后，作者试图找到一个合适的定律来模拟实验数据，并注意到在这个（小）数量级上，细化步骤和平行轨迹的贡献看起来都像一个幂律。此外，由于内核加速是有限的，因此性能指标应该达到饱和。因此，作者决定拟合以下定律（该定律在小数量级时呈现幂律表现，随着计算量的增加，收益会逐渐减少）：
作者发现，给定一个固定的、非微不足道的推理计算预算（例如，细化步骤并行轨迹≥8），最佳计算分配会转向串行细化，而不是并行生成。
单轮模型推理
作者之前在相同的多轮推理设置中比较了多轮模型（Kevin-32B）和单轮模型。但由于单轮模型是在单轮环境下训练的，
因此，一些问题就产生了：在单轮环境下训练的模型计算量固定的情况下，我们能用单轮推理或多轮推理获得更好的推理结果吗？
在这种特定环境下，即使使用单轮训练模型，多轮推理的结果一般也优于单轮推理（平均正确率除外）。
为了比较这两种方法，作者使用64个并行轨迹和仅1个步骤来评估单轮模型，然后将结果与使用16个并行轨迹和每个轨迹4个细化步骤的多轮推理进行比较。作者将单轮推理的64个并行轨迹分成16组，每组4个内核，取每组的best@4，然后取16组的平均值。这样就能将这一指标与多轮推理中的avg@16进行比较（因为在这种情况下，我们是在单个轨迹中取best@4）。最后，作者将单轮推理的best@64与多轮推理的best@16（4个细化步骤）进行了比较。
奖励塑造
作者尝试了奖励塑造。在较小模型上运行时，其添加了中间奖励（成功解析、编译、执行……）来引导模型。然而，作者发现这些奖励可能会分散模型的注意力，使其无法朝着真正的目标更新——生成正确且性能良好的内核。作者还按照Kimi的建议尝试了长度惩罚，但发现在已有的环境中，长度惩罚会降低模型的性能。
对于多轮训练，作者对不同的奖励函数进行了消融。他们尝试了不同的伽玛值（0.4与0.8），以及如何汇总单个轨迹的奖励——求和或取最大值。
在内核生成中，作者从根本上关心的是获得具有最大轨迹的内核（而不是优化多个内核的分数折扣总和）。因此，作者认为使用最大奖励公式会带来更好的加速效果。
然而，我们可以发现，将整个MDP的奖励相加，gamma=0.4的效果最好。
奖励塑造消融—多轮训练是在奖励总和与伽马值=0.4（sum_gamma_0_4）的条件下进行的。
并行轨迹
作者还尝试在训练过程中将并行轨迹的数量从16个增加到64个，即训练批次大小为644*8=2048。理论上讲，这将使优势估计的噪声更小，但他们发现在评估的best@16和avg@16性能上没有明显差异。此外，在第30步（60梯度步）左右，无用的东西开始提前出现。
数据分布
在项目初期，作者尝试在DeepSeek-R1-Distill-Qwen-14B上运行level1任务的简单子集。实验发现，奖励会出现高原现象，模型会过拟合单一难度级别。此外，模型只能学习有限的优化技术。因此，作者认为在数据集中均衡、多样地分布不同难度的任务非常重要。
测试时搜索
在训练过程中，为了保持合理的训练时间，必须限制搜索。但在测试时，可以自由地使用更复杂的测试时间技术来进一步提高性能。
为此，Kevin-32B使用了改进版的束搜索，其工作原理如下。作者首先在16条轨迹上执行4个串行细化步骤，就像在训练时间所做的那样。在这一过程结束时，根据生成的最快内核对轨迹进行排序，并保留最好的4个轨迹。然后，作者将每条轨迹复制4次（共16条轨迹），并重复这一过程。作为一种通用方法，它还显著提高了模型性能，在整个数据集上平均提速1.56倍。
不出所料，随着测试时间计算量的增加，得到的收益也在减少，尽管平均性能在几个小时后仍有提高。
束搜索测试时间推理。计算时间是在单个H200GPU上按内核计算的。请注意，其中还包括评估时的停滞时间，有效推理时间约占总计算时间的45%。
作者通过修改轨迹数、束宽度和细化步数，尝试了波束搜索的几种不同变体。有趣的是，他们发现中位加速度较高的设置平均加速度较低，反之亦然。
作者将这种行为归因于每种技术在探索/利用的帕累托前沿中的不同位置：通过更多的探索，内核可以实现一致的加速，尽管它们可能会错过非常激进的、显著提高平均速度的优化。
未来工作
我们相信，这项工作只是探索编程智能体训练方法的开始。如果有更多的时间和计算能力，我们希望尝试以下方法：
学习价值网络并使用PPO进行训练。事实上，基线估计器是按引用步骤而不是按提示计算的；
在训练时整合更复杂的搜索方法，如束搜索（而不仅仅是并行+串行搜索）；
将多轮训练法应用于更普遍的编程环境。
总结
在这项工作中，作者提出了一种方法，该方法适用于任何具有中间奖励的多轮环境。研究表明，这种方法比单圈GRPO取得了更好的结果。
端到端训练将是未来智能体的重要组成部分。虽然手工制作的多LLM工作流程可以带来短期收益，但它们在很大程度上依赖于人类的启发式方法，而且无法扩展。相比之下，更通用的方法可以让模型自由探索不同的轨迹（搜索），然后学习这些长期动态。通过与编程环境的交互，模型产生经验流，并通过反馈不断调整。我们希望这样的方法能成为迈向自主编程智能体的第一步。
HuggingFace:https ://huggingface.co/cognition-ai/Kevin-32B
最后，让我们来复习一下《TheBitterLesson》的道理：
我们必须吸取惨痛的教训，从长远来看，在我们的思维方式中植入知识是行不通的。这一惨痛教训基于以下历史观察：
人工智能研究人员经常试图在其智能体中构建知识；
这在短期内总是有帮助的，而且会让研究人员个人感到满意；
从长远来看，它会停滞不前，甚至抑制进一步的进展；
通过搜索和学习来扩展计算规模的相反方法最终会带来突破性进展。
参考内容：
https ://cognition.ai/blog/kevin-32b
©THEEND
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin.com
