Improving HPC Code Generation Capability of LLMs via Online Reinforcement Learning with Real-Machine Benchmark Rewards

2026-02-12

AI Summary

This work addresses the challenge that code generated by large language models for high-performance computing (HPC) often lacks runtime performance guarantees, failing to meet HPC’s stringent efficiency requirements. The authors propose a novel optimization framework based on online reinforcement learning that, for the first time, uses real-world supercomputer performance metrics—specifically GFLOPS—as the reward signal. By integrating a staged quality-diversity algorithm, the framework dynamically expands the strategy space to encourage the model to learn efficient code generation from multiple perspectives. The system employs Group Relative Policy Optimization (GRPO) within a distributed architecture that coordinates GPU-based training with CPU-based benchmarking, enabling end-to-end fine-tuning of Qwen2.5 Coder 14B. Experiments demonstrate significant improvements in code performance on key HPC tasks such as double-precision matrix multiplication.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Search and Optimization: Learning to Search
• Machine Learning: Hardware-aware ML

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• HPC code generation
• runtime performance
• large language models
• reinforcement learning
• code optimization

Innovation

• Reinforcement Learning
• HPC Code Generation
• Real-Machine Benchmarking
• Staged Quality-Diversity
• Runtime Performance Reward