HARGO给HPC任务的RL后训练提了个新思路,不用标签就能按难度加权,四个任务上全面超过GRPO。
SFT虽能让LLM掌握HPC领域知识,但在数据竞争检测与基准问答上行为不精准,88.65%的分类正确率下仍有65.9%的MLPerf回答超40字符。HPC任务异构性显著,答案长度差异达58倍,覆盖三种奖励分布。HARGO通过置信度调制的优势加权,无需任务类型标签即可适配。在四个HPC任务和九种方法的对比中,HARGO在WinRate 54.62%、Data Race F1 91.30%和PLP Similarity 0.8558上均取得最佳。
HARGO: Heterogeneity-Aware Reward-Guided Optimization for RL Post-Training of LLMs on HPC Tasks
Supervised fine-tuning (SFT) can equip large language models (LLMs) with domain knowledge for high-performance computing (HPC) tasks such as data race detection and benchmark question answering. However, knowledge alone does not guarantee task-appropriate behavior: the same SFT model that correctly classifies 88.65\% of C/C++ data race samples produces verbose, imprecise answers to factual queries, with 65.9\% of MLPerf responses exceeding 40 characters. Reinforcement learning (RL) post-training addresses this gap by optimizing for task-specific rewards rather than token-level imitation. Yet HPC tasks exhibit extreme heterogeneity, with binary classification, factual QA, and semantic generation differing by 58x in answer length, spanning three distinct reward distributions, and showing widely varying SFT accuracy. This makes uniform-weight RL methods such as GRPO suboptimal. We propose HARGO, Heterogeneity-Aware Reward-Guided Optimization, which introduces per-response importance weighting via confidence-modulated advantage: computing a discrimination signal from group-level reward contrast and a confidence signal from reference model log-probabilities, then modulating the advantage before computing per-response weights, without requiring task-type labels. Across four HPC tasks and nine methods, HARGO achieves the best performance on all three primary metrics: WinRate 54.62\%, Data Race F1 91.30\%, and PLP Similarity 0.8558. Ablation confirms complementary contributions from both signals. HARGO establishes the best overall alignment quality among compared methods for heterogeneous HPC tasks.