用真实超算跑代码反馈性能,让大模型学会生成高效高性能计算代码。
Improving HPC Code Generation Capability of LLMs via Online Reinforcement Learning with Real-Machine Benchmark Rewards
- 在线强化学习:让大模型生成代码后在超算上实测运行时间作为奖励。
- 在双精度矩阵乘法任务中,生成代码的性能提升至12.3 GFLOPS。
- 分阶段多样性优化策略,帮助模型从多种角度学习代码加速技巧。
大型语言模型(LLMs)虽具备强大代码生成能力,但其生成代码的运行效率无法保证,且在高性能计算(HPC)领域极少使用运行时性能作为训练奖励。本文提出一种在线强化学习方法,将LLM生成的代码在超算上执行,并直接以测得的运行性能(GFLOPS)作为奖励反馈。进一步引入分阶段质量-多样性(SQD)算法,按问题动态调整允许的优化技术,使模型从多元视角学习代码优化。构建了一个连接GPU训练集群与CPU基准测试集群的分布式系统,基于群组相对策略优化(GRPO)对Qwen2.5 Coder 14B进行训练。通过两项实验验证,结合运行时性能反馈与分阶段优化的强化学习可显著提升LLM在HPC代码生成方面的能力。
原文摘要 · Abstract (English)
Large language models (LLMs) have demonstrated strong code generation capabilities, yet the runtime performance of generated code is not guaranteed, and there have been few attempts to train LLMs using runtime performance as a reward in the HPC domain. We propose an online reinforcement learning approach that executes LLM-generated code on a supercomputer and directly feeds back the measured runtime performance (GFLOPS) as a reward. We further introduce a Staged Quality-Diversity (SQD) algorithm that progressively varies the permitted optimization techniques on a per-problem basis, enabling the model to learn code optimization from diverse perspectives. We build a distributed system connecting a GPU training cluster with a CPU benchmarking cluster, and train Qwen2.5 Coder 14B on a double-precision matrix multiplication task using Group Relative Policy Optimization (GRPO). Through two experiments, we show that reinforcement learning combining runtime performance feedback with staged optimization can improve the HPC code generation capability of LLMs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。