DeepSeek關鍵RL演算法GRPO，有人從頭跑通了，貢獻完整程式碼
机器之心發表於2025-03-02
GRPO（Group Relative Policy Optimization）是 DeepSeek-R1 成功的基礎技術之一，我們之前也多次報導過該技術，比如《DeepSeek 用的 GRPO 佔用大量記憶體？有人給出了些破解方法》。簡單來說，GRPO 演算法丟棄了 critic model，放棄了價值函式近似，轉而透過組內樣本的相對比較來計算策略梯度，從而有效降低了訓練的不穩定性，同時提高了學習效率。既然 GRPO 如此有效，那麼，你知道如何從頭開始實現 GRPO 嗎？近日，AI 工程師和技術作家 Andriy Burkov 釋出了一份「從頭開始寫 GRPO 程式碼」的教程，其中介紹瞭如何基於 Qwen2.5-1.5B-Instruct 模型構建一個使用 GRPO 的分散式強化學習流程。不過，在我們深入這份教程之前，先簡單介紹一下它的作者。Andriy Burkov 算得上是 AI 領域的一位著名科普作家，在加拿大拉瓦爾大學取得了電腦科學博士學位，還曾發表過兩本頗受歡迎的 AI 主題著作：《100 頁語言模型書》和《100 頁機器學習書》；書中一步步詳實地介紹了相關概念，並附帶了簡明的實現程式碼。接下來我們就來看看這份 GRPO 從頭實現教程吧。教程地址：https://github.com/aburkov/theLMbook/blob/main/GRPO_From_Scratch_Multi_GPU_DataParallel_Qwen_2_5_1_5B_Instruct.ipynb使用 Qwen2.5-1.5B-Instruct 的分散式實現本教程將展示如何使用 GRPO 方法構建分散式強化學習（RL）流程，從而可以針對數學、邏輯和程式設計任務對語言模型進行微調。首先需要明確，這些任務都存在一個唯一且正確的 ground truth 答案，可透過簡單的字串比較輕鬆加以驗證。GRPO 的發明者是 DeepSeek，最早是被用於微調 DeepSeek 的 R1 和 R1-Zero 模型 —— 它們可透過學習生成思維鏈（CoT）來更好地解決數學和邏輯問題。本教程的目標是將通用語言模型 Qwen2.5-1.5B-Instruct 轉換為數學問題求解器。我們將從頭開始編寫 GRPO 程式碼，然後將其與幾個流行的庫和工具整合起來，以實現分散式訓練管道流程，包括：- Hugging Face Transformers：用於載入預訓練的語言模型和 tokenizer。
- FlashAttention2：最佳化的注意力機制，有助於減少記憶體使用量並提高訓練速度。
- Weights & Biases (wandb)：用於實驗跟蹤、視覺化和模型版本控制。
本教程分為幾個部分。首先是基本設定和匯入，然後是資料格式化和答案提取、資料集準備、評估函式、獎勵函式、訓練設定和執行，最後載入和測試模型。此過程中，我們將從頭實現 GRPO 演算法。首先是安裝並匯入所有必要的模組。下面是匯入庫的一段程式碼截圖。執行上述程式碼（參考專案完整程式碼），可以執行以下任務：- 設定隨機種子：set_random_seed 函式透過為 Python 的隨機模組、NumPy 和 PyTorch 設定種子，確保可復現性；
- 環境變數配置：設定 WANDB_API_KEY 和 WANDB_PROJECT 環境變數，以啟用與 Weights & Biases 的實驗跟蹤；
- 匯入必要的庫，包括 random、copy、re、torch 等等。
接下來，專案定義了資料格式，以及模型如何從輸出和資料集中提取答案段落。為了確保模型輸出格式一致，專案還定義了一個系統提示。該提示指示模型生成包含 < reasoning > 和 < answer > 標籤的輸出。這一步透過兩個函式完成：- extract_answer_from_model_output：此函式獲取模型的輸出文字，並提取 < answer > 標籤內的內容；
- extract_answer_from_dataset：此函式從 GSM8K 資料集中提取預期答案，該資料集使用 “####” 分隔符來分隔答案：
該專案使用 GSM8K 資料集進行訓練。專案使用了該資料集中的示例來訓練模型，基於強化學習（RL）訓練正規化，讓模型生成多個問題解答樣本，之後作者將這些解答與 GSM8K 示例中的標準答案進行對比，如果匹配，就為 RL 演算法（GRPO）提供高獎勵，然後更新模型權重，以增加模型下次獲得高獎勵的可能性。實驗過程是這樣的。首先從 Hugging Face 載入資料集，然後格式化每個示例，包括系統提示和使用者提示。這段實現程式碼中還定義了兩個輔助函式：prepare_dataset 以及 build_prompt。評估對於跟蹤模型的進展至關重要。因此作者定義了一些函式，從而可以在一組示例上對模型進行評估。該專案的評估函式執行以下任務：- token 化提示並生成響應：模型的輸出是在 token 化提示的基礎上生成的。
- 將預測答案與預期答案進行比較：這種比較是透過精確匹配以及數值等價檢查來完成的。
在這段程式碼中，兩個輔助函式 _extract_last_number 和 _extract_single_number 被用來從文字中提取數字。評估函式 evaluate_model 使用這些輔助函式來確定預測答案是否正確：在強化學習中，獎勵函式是必不可缺的，作者定義了兩個獎勵函式：correctness_reward：這個函式根據生成的答案是否正確來分配獎勵。採用兩種方式：精確的字串匹配和數值等價檢查，將模型輸出的答案與預期答案進行比較。完全匹配會獲得更高的獎勵（2.0），而基於數值等價的匹配會獲得較小的獎勵（1.5）。format_reward：這個函式鼓勵模型遵循所需的類似 XML 的輸出格式。它為生成文字中存在 < reasoning>、</reasoning>、<answer > 和 </answer > 標籤提供小額獎勵。Part 6：從頭開始實現 DataParallel GRPO這一節，我們將從頭實現 GRPO 演算法的所有構建模組。首先，這裡假設執行程式碼的機器至少有 2 臺 GPU。為此，這裡要使用 PyTorch 的 DataParallel API 來將策略模型放在多個 GPU 核心上，每個 GPU 核心都有該模型的一個副本。然後將批次資料分散在這些 GPU 核心上完成處理。這一節，我們將所有元件組合在一起，完成設定並開始訓練。首先，載入預訓練的模型和 tokenizer，準備評估資料，然後使用上面從頭實現的 train_with_grpo 進行強化學習微調。- 模型和 tokenizer 初始化：使用最佳化設定（使用 torch.bfloat16 和 FlashAttention2）載入模型 Qwen/Qwen2.5-1.5B-Instruct。tokenizer 也要載入，其填充 token 設定為序列末尾 token。使用 torch.bfloat16 載入模型會將其引數轉換為每個數值使用 16 位而不是 32 位的形式，這可將模型的記憶體使用量減少一半，並且可加快在現代 GPU 上的訓練速度。
- 初步評估：在微調之前，根據幾個示例對模型進行評估，以確定基準效能。
- 強化學習微調：為從頭開始實現 GRPO 的訓練函式 train_with_grpo 配置適當的訓練引數和獎勵函式。然後，在剩餘的訓練資料上執行強化學習訓練。
- 最終評估和模型儲存：強化學習微調後，再次評估模型，並儲存最終模型。
- 確定裝置（如果有 GPU 就用 GPU，否則就用 CPU）。
- 載入預訓練版 Qwen2.5-1.5B-Instruct 模型和 tokenizer。tokenizer 的 pad token 設定為 eos_token。
- 透過啟用梯度檢查點和停用 KV 快取，最佳化模型的記憶體效率。
- 步驟 2：使用 train_with_grpo 函式和我們定義的獎勵函式（format_reward 和 correctness_reward，合併為 combined_reward）執行強化學習微調。這裡使用了多臺 GPU 訓練模型。
- 步驟 3：將最終的微調模型和 tokenizer 儲存到磁碟。
以下引數設定了使用上面的 GRPO 演算法實現強化學習微調執行的配置：- num_iterations=1：從當前策略模型建立新參考模型的外部迭代次數。一次迭代是指在整個資料集上執行一次透過。
- num_steps=500：訓練迴圈將執行最多 500 個步驟，每個步驟處理一批樣本。
- batch_size=7：在 8 臺 GPU 的情況下，每個步驟每批處理 7 個樣本，每臺 GPU 上放置 1 個樣本。使用一個 GPU (0) 被 DataParallel 用作主節點來聚合梯度並收集輸出。
- num_generations=14：對於訓練資料中的每個提示詞，訓練器將生成 14 個不同的完成結果。這些生成結果將被用於計算指導強化學習更新的相對優勢（或獎勵訊號）。如果你的 GPU 的 VRAM 較少，請減少此數字。
- max_completion_length=400：在生成完成結果（序列的 response 部分）時，生成上限為 400 個 token。這限制了模型在 RL 階段生成的輸出的長度。如果你的 GPU 的 VRAM 較少，請減少此數字。
- beta=0.04：GRPO 損失函式中 KL 散度懲罰的係數。這控制的是模型與參考模型的偏差程度。
- learning_rate=5e-6：RL 微調的學習率。為了實現穩定的策略更新，這裡使用了相對較低的學習率。
- mu=1：對每批 rollout 資料執行的策略更新次數。在這裡，我們每批只執行一次更新。
- epsilon=0.1：GRPO 的 PPO 元件的 clipping 引數。這可以防止策略在單次更新中發生太大的變化。
在微調之前和之後都會對模型進行評估，以衡量準確率的提高情況。最後，將微調後的模型儲存到 grpo_finetuned_model 目錄中。首先，初始配置後，我們可以看到執行 GRPO 之前的準確度為 23.33%。然後經過 500 步的 1 輪 GRPO 迭代，下圖展示了相關的訓練動態：訓練完成後，自然還需要對模型進行新一輪的評估。這裡採用了 30 個評估樣本來進行評估，以下展示了其中一個模型回答正確的示例：整體表現如何呢？可以看到，經過一輪 GRPO 之後，Qwen-2.5-1.5B-Instruct 模型答對了 30 問題中的 27 題，實現了 90% 的準確度。相較於 GRPO 之前的 23.33%，可說是實現了效能飛躍。上面兩張圖展示了模型的學習過程動態，可以看到：平均獎勵在 2.25 左右就趨於穩定了（理論最大值為 0.8 + 2.0 = 2.8）。相比於另一處微調的 Qwen-2.5-0.5B-Instruct（獲得的平均獎勵為 1.4），這個數字相當高了，參閱：https://github.com/aburkov/theLMbook/blob/main/GRPO_Qwen_0_5_Instruct.ipynb如果使用更大的模型並允許更長的生成時間，模型正確解答問題的能力還將進一步提升。但是，如果要訓練更大的模型，不僅需要將資料分佈在多臺 GPU 上，還需要將模型分開放在多臺 GPU 上，這需要用到 DeepSpeed 或 FSDP（完全分片資料並行）等模型並行工具。可以看到，模型反覆思考了很多次，終於認定確實等於 2。多次測試後還可以發現，該模型沒有學會生成序列結束（EOS）token，因此即使在 </answer> token 之後，輸出序列仍會繼續。這是預期的行為，因為我們使用的獎勵函式中沒有包含一個用於停止生成的獎勵。我們也沒有執行監督微調步驟 —— 該步驟可以讓模型學會在 </answer> 之後立即生成 EOS。你對這篇程式碼密集的教程怎麼看？有沒有讓你產生在自己的電腦上實現 GRPO 的想法？ 
