基于 DeepSeek GRPO 的 1.5B Rust 代碼生成模型訓練實戰 原創
編者按: 群組相對策略優化(GRPO)如何讓小型專用模型在特定任務上實現性能提升?我們今天為大家帶來的這篇文章展示了如何使用 GRPO,訓練一個僅有 1.5B 參數的 Rust 代碼生成模型,實現性能大幅提升。
文章詳細介紹了作者基于 Qwen2.5-Coder-1.5B-Instruct 模型使用 GRPO 技術進行訓練的實踐經驗。作者選擇 Rust 語言作為實驗對象,利用其嚴格的編譯器和完善的工具鏈作為反饋機制,構建了一套基于格式驗證與 Cargo 的獎勵函數。通過單次 GRPO 訓練,模型的代碼編譯通過率從 61% 提升至 80%,單元測試通過率從 22% 提升至 37%,成本不到 100 美元。
作者 | Greg Schoeninger
編譯 | 岳揚
群組相對策略優化(Group Relative Policy Optimization,GRPO)已被證明是一種有效的算法,可用于訓練大語言模型(LLMs),使其具備推理能力并在基準測試中持續提升性能表現。DeepSeek-R1 展示了如何通過監督式微調(Supervised Fine-Tuning)與 GRPO 技術的結合,引導模型達到與 OpenAI 的 o1 等頂尖模型相競爭的水平。
為了進一步探索其實踐應用,我們嘗試將這些技術應用于現實場景中。本文將概述如何利用 GRPO、自定義數據集和獎勵函數來訓練我們自己的定制化小型 LLM。下圖是后續將展示的部分模型訓練過程曲線預覽。觀察模型逐步學習生成代碼塊、優化可編譯的有效代碼,最終生成能通過單元測試的代碼,這一過程頗具趣味性。
若希望直接上手實踐,可訪問該 GitHub 倉庫[1]。
本文不會深入探討 GRPO 的基礎原理,若需了解其底層機制,請參閱??個人開發者也能訓練推理模型?GRPO 技術詳解??。
01 為什么選擇 Rust?
Rust 似乎是強化學習(RL)的絕佳舞臺,因為可以使用 Rust 編譯器和 Cargo 工具鏈。Rust 編譯器不僅能夠提供清晰的錯誤信息,還具備嚴格的類型檢查。
在本項目中,我們首先想要驗證的假設是:能否將 Cargo 作為反饋機制來教導模型成為更優秀的“程序員”。第二個實驗旨在探索用 GRPO 訓練,語言模型的最小可行規模是多少。這些實驗特意限制在單節點 H100 上運行,既控制了成本又展示了實驗的可實現性。
而且作為 ??Oxen.ai?? 的 Rust 開發團隊,我們也正在探索一些有趣的跨界應用 ?? x ??。
02 為何參數量選擇 1.5B?
最近,有大量研究探索了小型語言模型在特定任務上的性能邊界。當存在明確的反饋機制(如數學題的正確答案或程序的輸出結果)時,模型規模可在保持高競爭力的同時大大縮小。
微軟的 rStar-Math[2] 論文表明,在可驗證的數學問題領域,模型可以進行推理。其 1.5B 參數量的模型性能超越了 GPT-4o 和 o1-preview。
我的假設是:在編碼任務中我們同樣能實現類似水平的性能,因為該領域同樣存在類似的可驗證獎勵機制 —— 代碼能否通過編譯?能否通過單元測試?
03 小型語言模型的優勢
小型代碼模型具備諸多優勢,包括成本效益、處理效率、數據隱私性,以及針對自有代碼庫/編碼規范進行定制的能力。此外,這本身也是項有趣的挑戰。
我們的終極愿景是,讓這類小型模型能完整執行所有 cursor 能夠完成的任務:預測下一段代碼、補全中間的代碼內容,并在智能體的優化迭代循環中持續優化代碼。不過,還是讓我們先從最基礎開始吧。
04 明確定義待解決問題
要讓生成的代碼能夠通過單元測試,有多種不同的方法,我們最終嘗試了幾種方案。其中一種看似簡單的方法是,提供一組可驗證的單元測試,要求生成的代碼必須通過這些測試。這將形成一套黃金標準的可驗證答案集。
在嘗試了該流程后,我們發現了兩個主要問題。首先,若禁止模型在編寫代碼時查看單元測試,它就無法感知需要遵循的接口規范。在使用預構建的、已驗證的單元測試進行評估時,許多錯誤最終表現為代碼與單元測試之間的類型不匹配或命名不一致。
其次,若允許模型在編碼時查看單元測試,就會犧牲開發者體驗。除非您是一名嚴格的"測試驅動開發"實踐者(Test Driven Developer),否則您可能更希望直接輸入提示詞,而非預先設計函數定義或單元測試。
我們沒有試圖想出更聰明的辦法,而是最終選擇了優化方案,使其更加簡潔,將這個問題重構為要求模型在單次響應中同時生成代碼和單元測試。
不過,單次生成方案存在模型"作弊"的風險:例如通過 println! 輸出而非 assert 斷言,生成可通過編譯但無實際功能的代碼。我們將在后續環節為此添加防護機制。
最終,我們通過一個詳細的系統提示詞(system prompt),為模型提供任務指引。
系統提示詞(system prompt)以特定格式和風格為模型提供上下文,明確模型響應用戶查詢的預期形式。
05 數據集
在開始訓練之前,我們需要一個數據集。最開始時,我們發現專門針對 Rust 語言的數據集較少,現有的大語言模型(LLM)基準測試大多聚焦于 Python。因此,我們首先將包含 Python 編程問題的提示詞數據集轉換為 Rust 語言版本。
我們從 Ace-Code-87k 數據集中隨機選取了 20,000 個提示詞樣本,使用 Qwen 2.5 Coder 32B Instruct 模型生成對應的 Rust 代碼及單元測試。隨后通過編譯器和測試框架運行這些代碼,篩選出所有能通過單元測試的提示詞、代碼、單元測試三元組(prompt,code,unit_test triples),最終獲得 16,500 組有效數據。該數據集被劃分為 15,000 組訓練數據、1,000 組測試數據和 500 組評估數據。
最終數據集[3]如下所示:
您可以通過以下模型運行記錄跟蹤整個流程:
1)將代碼翻譯為 Rust 語言:??https://www.oxen.ai/ox/mbrp-playground/evaluations/ce45630c-d9e8-4fac-9b41-2d41692076b3??
2)編寫 Rust 代碼:??https://www.oxen.ai/ox/mbrp-playground/evaluations/febc562a-9bd4-4e91-88d7-a95ee676a5ed??
3)編寫 Rust 單元測試:??https://www.oxen.ai/ox/mbrp-playground/evaluations/b886ddd6-b501-4db8-8ed6-0b719d0ac595??
有趣的是,在最終的 GRPO 訓練方案中,我們移除了原始數據中的黃金標準 Rust 代碼和單元測試列。在強化學習的不斷迭代學習過程中,我們僅需將提示詞作為輸入即可完成訓練,這為未來數據的收集提供了便利。盡管我們舍棄了用于訓練的代碼和單元測試數據,但知道僅需這些提示詞就可以解決這個問題仍具有重要參考價值 —— 這一特性將在后續章節詳細解析。
06 設定 baseline
完成待解決問題的明確定義并構建完數據集后,我們需設定 baseline 以評估初始模型性能。我們將使用 Qwen/Qwen2.5-Coder-1.5B-Instruct 模型進行訓練引導。
并統計以下指標:構建通過率(譯者注:代碼編譯成功)、Clippy 檢查通過率(譯者注:通過 Rust 官方靜態分析工具 Clippy 檢測)、單元測試通過率。
由于 Clippy 檢測依賴于代碼的成功編譯,這兩項指標通常呈現強相關性。您可在 ??Oxen.ai?? 平臺查閱原始數據[4],進一步分析具體案例。
07 與 SOTA 模型進行對比,孰強孰弱?
我們同時希望了解主流大型基礎模型在該任務上的表現,以此為我們設定理論上的優化目標。經測試發現表現最佳的模型是 GPT4.5。
GPT4.5 的構建通過率高達 98%,他自己編寫的單元測試的通過率達 87%,表現令人矚目[5]。
Claude 3.7 Sonnet 以微弱差距位居次席。
至此,所有準備工作就緒 —— 已完成明確定義待解決問題、數據集構建、baseline 設定及目標確立,終于可以進入激動人心的模型訓練環節!
08 獎勵函數設計
GRPO(群組相對策略優化)的精妙之處在于,其獎勵機制可通過簡單的 Python 函數實現。工程師只需定義獎勵規則,模型便能自主優化策略。摩根士丹利的 Will Brown 將這種方法稱為 Rubric Engineering,它為工程師提供了通過強化學習引導模型的便捷方法。
GRPO 獎勵函數的輸入參數包括 prompts(提示詞)、responses(模型響應)及 target answers(目標答案)。實際上,部分評分標準甚至無需 target answers(目標答案),僅需基于生成內容本身的屬性(如格式規范)進行評分。典型的評分維度包括:
- 正確性(直接字符串匹配)
- 響應長度(token 數量)
- 響應格式(XML、JSON、代碼等)
- 外部工具調用(本例中為 Cargo 工具鏈)
- 使用 LLM 進行評判(是否真實、是否有幫助、是否有危害等)
在本項目中,我們將結合格式驗證與 Cargo 構建工具的執行結果構建獎勵函數。初始代碼基于 @willccbb 的優質代碼實現:
??https://gist.github.com/willccbb/4676755236bb08cab5f4e54a0475d6fb??
09 我們的“評分標準”
以下是我們針對該問題設計的“評分標準”。流程大致為:接收用戶提示詞 → LLM 生成代碼及單元測試 → 通過多個評分函數評估模型響應。
獎勵函數可以是一個簡單的正則表達式,驗證代碼中是否包含有效的測試模塊。
您會注意到,輸入參數 prompts 和 completions 都是列表,輸出的也是浮點數列表(list[float])。這是因為 GRPO 有一個名為 num_generations 的參數,用于控制每個 prompts 的生成次數。此外,在生成過程中可能會有多組 prompts + completions。GRPO 算法會為每個 prompts 生成 N 種不同響應,每個模型響應都會通過“評分標準”進行驗證。下圖展示了模型在嘗試不同迭代方案時各評分項的通過情況:
隨著時間的推移,模型將逐步掌握如何通過更多你定義的“評分標準”。借助 trl 庫,你可以定義多種獎勵函數并傳入 GRPOTrainer 進行模型訓練。
設計獎勵函數堪稱是搭建 GRPO 訓練流程中最具成就感的環節 —— 雙關語警告??。
10 Cargo 獎勵函數
正如文章開頭所述,我們將利用 Cargo 工具鏈構建獎勵機制。由于獎勵函數可定義為純 Python 函數,因此我們可直接通過 subprocess 模塊調用 Cargo 工具。
完整代碼實現詳見 GitHub:??https://github.com/Oxen-AI/GRPO-With-Cargo-Feedback/blob/main/train.py??
我們定義了一個 RustTool 類,支持執行 cargo build、cargo clippy 或 cargo test 命令。其 run() 方法將返回包含工具執行結果(通過/失敗狀態信息、錯誤信息)的字典。
該工具應用于為每個測試用例臨時創建的 Rust 項目目錄。項目的初始化與清理流程包括:創建目錄→寫入 ??main.rs??? 和 Cargo.toml 文件→在其中寫入代碼與單元測試→執行后清理。 欲知更多詳情,可查閱完整代碼中的 setup_and_test_rust_project 函數。
在具備初始化和清理 Rust 小型項目的能力后,需將其與 GRPO 獎勵函數對接。獎勵函數接收一批提示詞(prompts)與生成內容(completions),要求對每組 prompt+completion 進行評分。
持久化記錄 prompts 和 completions 有助于幫助我們理解 GRPO 算法的內部機制。每個 prompt 將生成 N 個 completions。假如我們設置 N=4,這樣模型就有 4 次生成內容的機會。
以下方日志中的 task_8347 為例,模型嘗試了 4 次函數代碼的實現,最終成功一次。GRPO 會對正確的解決方案進行獎勵,從而促使模型性能持續提升。 使用紅色標記圈出來的部分是三個單元測試失敗的案例,綠色標記的部分則為測試通過的案例:
使用相同的工具與邏輯,我們為 cargo build、test 和 clippy 分別構建獎勵函數。同時設置驗證機制確保代碼與單元測試不是空的,且單元測試必須包含 assert! 斷言,防止模型通過 println! 語句作弊。
所有訓練結果均實時記錄至 ??Oxen.ai??,以便繪制模型訓練曲線(如文章開篇所示)。在模型訓練過程中,我們會通過滾動平均(譯者注:Rolling Average,是一種統計方法,用于對時間序列數據中的短期波動進行平滑處理。) 的方式對數據進行平滑處理、計算,以此觀察模型在特定獎勵指標上的改進趨勢。
例如編譯通過率剛開始在 30-40% 之間波動,隨著訓練逐步攀升至 70%:
單元測試通過率的提升相對滯后,且通過率的波動范圍更大:
需要觀察模型在每個獎勵維度(如代碼編譯通過率等)上的改進情況,定期抽查模型接收的提示詞(輸入)及其生成的代碼(輸出)。為此我們編寫了一個 @experiment.log 裝飾器,在獎勵函數執行時自動將結果寫入 jsonl 文件并提交至 ??Oxen.ai??:
該裝飾器每次調用函數時將結果寫入指定文件,訓練循環(training loop)中的回調函數每 N 步將數據提交至 ??Oxen.ai???。由于數據按時間順序存儲,需翻頁至末尾查看訓練后期的生成樣本。部分輸出示例:??https://www.oxen.ai/ox/Rust/file/GRPO_82_2025-03-02_22-49-17_Qwen2.5-Coder-1.5B-Instruct/outputs/GRPO_82_2025-03-02_22-49-17_Qwen2.5-Coder-1.5B-Instruct/cargo_test_rewards.jsonl?page=497&ref=ghost.oxen.ai??
11 訓練效果評估
此前設定的 baseline —— 未經訓練的 Qwen/Qwen2.5-Coder-1.5B-Instruct 模型在代碼編譯通過率和單元測試通過率上分別僅為 61% 和 22%。
經過一輪 GRPO 訓練后,編譯通過率提升至 80%,單元測試通過率達 37% ??。 通過一輪訓練即實現 20% 與 15% 的絕對準確率提升。
僅僅定義了幾個獎勵函數和使用較小規模的數據集,這樣的訓練效果已經相當不錯了。您可通過此處[6]探索 1.5B 模型的原始輸出結果。
此次實驗成果令人鼓舞。整個訓練耗時約 24 小時,成本小于 100 美元(下圖是 H100 實例的每小時成本,僅供參考):
該實驗表明:GRPO 算法對普通開發者來說也比較實用 —— 開發者可靈活定義任意獎勵函數,即使在小模型上也能實現性能的明顯提升。
Thanks for reading!
Hope you have enjoyed and learned new things from this blog!
END
本期互動內容 ??
?如果讓你為代碼生成任務設計獎勵函數,您會優先考慮哪個指標?為什么?
??文中鏈接??
[1]??https://github.com/Oxen-AI/GRPO-With-Cargo-Feedback/tree/main??
[2]??https://arxiv.org/abs/2501.04519??
[3]??https://www.oxen.ai/ox/Rust/file/main/cargo_test_passed_train.parquet??
[5]??https://www.oxen.ai/ox/Rust/file/gpt4-5-results/results/GPT4.5/predictions.parquet??
原文鏈接:
??https://ghost.oxen.ai/training-a-rust-1-5b-coder-lm-with-reinforcement-learning-grpo/??
