Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Ryo Mikasa, Shun-ichiro Hayashi, Daichi Mukunoki, Tetsuya Hoshino, Takahiro Katagiri

TL;DR

The paper tackles the challenge of translating LLM-generated code into high-performance HPC kernels by using online reinforcement learning with real-device rewards. It combines GRPO with a Staged Quality-Diversity strategy to explore diverse optimization pathways under constraint, enabling the model to learn performance-oriented code as it interacts with actual hardware. Across two experiments on a GEMM kernel with a 256^3 problem size, the approach yields significant gains over baseline prompts, including a peak of around 549 GFLOPS, and reveals that guidance via staged constraints can stabilize learning and reveal robust optimization patterns. The work demonstrates a viable path toward performance-aware code generation in realistic HPC environments and points to future work on generalization, broader kernels, and cross-architecture portability.

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.

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

TL;DR

The paper tackles the challenge of translating LLM-generated code into high-performance HPC kernels by using online reinforcement learning with real-device rewards. It combines GRPO with a Staged Quality-Diversity strategy to explore diverse optimization pathways under constraint, enabling the model to learn performance-oriented code as it interacts with actual hardware. Across two experiments on a GEMM kernel with a 256^3 problem size, the approach yields significant gains over baseline prompts, including a peak of around 549 GFLOPS, and reveals that guidance via staged constraints can stabilize learning and reveal robust optimization patterns. The work demonstrates a viable path toward performance-aware code generation in realistic HPC environments and points to future work on generalization, broader kernels, and cross-architecture portability.

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.
Paper Structure (26 sections, 3 equations, 14 figures, 3 tables)

This paper contains 26 sections, 3 equations, 14 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Training Large Language Models
  4. RL for LLMs: From PPO to GRPO
  5. Proposed Method
  6. Reward Function Design
  7. Staged Optimization Constraints
  8. SQD Algorithm
  9. Distributed Training Architecture
  10. Online GRPO Training Loop
  11. Experiments
  12. Experimental Setup
  13. Experiment 1: Hyperparameter Search with Fixed Prompts
  14. Overview
  15. Results
  16. ...and 11 more sections

Figures (14)

  • Figure 1: SQD with online RL.
  • Figure 2: Distributed training architecture: GPU training cluster and CPU benchmark cluster connected via a coordinator
  • Figure 3: Experiment 1: Mean reward across all settings. Reward is raw GFLOPS for successful codes and a negative penalty for failures (Eq. \ref{['eq:reward']}). Solid lines are LR=2e-7; dashed lines are LR=5e-7.
  • Figure 4: Experiment 1: Max GFLOPS per step. Solid lines are LR=2e-7; dashed lines are LR=5e-7.
  • Figure 5: Experiment 1: Generated code before training (step 0, verification error).
  • ...and 9 more figures