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

從頭編寫 GRPO 代碼

使用 Qwen2.5-1.5B-Instruct 的分布式實現

本教程將展示如何使用 GRPO 方法構建分布式強化學習（RL）流程，從而可以針對數學、邏輯和編程任務對語言模型進行微調。

首先需要明確，這些任務都存在一個唯一且正確的 ground truth 答案，可通過簡單的字符串比較輕松加以驗證。

GRPO 的發明者是 DeepSeek，最早是被用于微調 DeepSeek 的 R1 和 R1-Zero 模型 —— 它們可通過學習生成思維鏈（CoT）來更好地解決數學和邏輯問題。

本教程的目標是將通用語言模型 Qwen2.5-1.5B-Instruct 轉換為數學問題求解器。我們將從頭開始編寫 GRPO 代碼，然后將其與幾個流行的庫和工具集成起來，以實現分布式訓練管道流程，包括：

• PyTorch：用于張量運算和分布式訓練。
• Hugging Face Transformers：用于加載預訓練的語言模型和 tokenizer。
• FlashAttention2：優化的注意力機制，有助于減少內存使用量并提高訓練速度。
• Weights & Biases (wandb)：用于實驗跟蹤、可視化和模型版本控制。

本教程分為幾個部分。首先是基本設置和導入，然后是數據格式化和答案提取、數據集準備、評估函數、獎勵函數、訓練設置和執行，最后加載和測試模型。此過程中，我們將從頭實現 GRPO 算法。

Part 1：基礎設置和導入

首先是安裝并導入所有必要的模塊。下面是導入庫的一段代碼截圖。

部分代碼截圖。完整代碼塊參見 GitHub。

運行上述代碼（參考項目完整代碼），可以執行以下任務：

• 設置隨機種子：set_random_seed 函數通過為 Python 的隨機模塊、NumPy 和 PyTorch 設置種子，確?？蓮同F性；
• 環境變量配置：設置 WANDB_API_KEY 和 WANDB_PROJECT 環境變量，以啟用與 Weights & Biases 的實驗跟蹤；
• 導入必要的庫，包括 random、copy、re、torch 等等。

Part 2：數據格式以及答案提取

接下來，項目定義了數據格式，以及模型如何從輸出和數據集中提取答案段落。

為了確保模型輸出格式一致，項目還定義了一個系統提示。該提示指示模型生成包含 < reasoning > 和 < answer > 標簽的輸出。這一步通過兩個函數完成：

• extract_answer_from_model_output：此函數獲取模型的輸出文本，并提取 < answer > 標簽內的內容；
• extract_answer_from_dataset：此函數從 GSM8K 數據集中提取預期答案，該數據集使用 “####” 分隔符來分隔答案：

部分代碼截圖。完整代碼塊參見 GitHub。

Part 3：數據準備

該項目使用 GSM8K 數據集進行訓練。項目使用了該數據集中的示例來訓練模型，基于強化學習（RL）訓練范式，讓模型生成多個問題解答樣本，之后作者將這些解答與 GSM8K 示例中的標準答案進行對比，如果匹配，就為 RL 算法（GRPO）提供高獎勵，然后更新模型權重，以增加模型下次獲得高獎勵的可能性。

實驗過程是這樣的。首先從 Hugging Face 加載數據集，然后格式化每個示例，包括系統提示和用戶提示。這段實現代碼中還定義了兩個輔助函數：prepare_dataset 以及 build_prompt。

部分代碼截圖。完整代碼塊參見 GitHub。

Part 4：評估函數

評估對于跟蹤模型的進展至關重要。因此作者定義了一些函數，從而可以在一組示例上對模型進行評估。該項目的評估函數執行以下任務：

• token 化提示并生成響應：模型的輸出是在 token 化提示的基礎上生成的。
• 提取預測答案：從生成的響應中提取答案。
• 將預測答案與預期答案進行比較：這種比較是通過精確匹配以及數值等價檢查來完成的。

在這段代碼中，兩個輔助函數 _extract_last_number 和 _extract_single_number 被用來從文本中提取數字。評估函數 evaluate_model 使用這些輔助函數來確定預測答案是否正確：

部分代碼截圖。完整代碼塊參見 GitHub。

Part 5：獎勵函數

在強化學習中，獎勵函數是必不可缺的，作者定義了兩個獎勵函數：

correctness_reward：這個函數根據生成的答案是否正確來分配獎勵。采用兩種方式：精確的字符串匹配和數值等價檢查，將模型輸出的答案與預期答案進行比較。完全匹配會獲得更高的獎勵（2.0），而基于數值等價的匹配會獲得較小的獎勵（1.5）。

format_reward：這個函數鼓勵模型遵循所需的類似 XML 的輸出格式。它為生成文本中存在 < reasoning>、</reasoning>、<answer > 和 </answer > 標簽提供小額獎勵。

部分代碼截圖。完整代碼塊參見 GitHub。

Part 6：從頭開始實現 DataParallel GRPO

這一節，我們將從頭實現 GRPO 算法的所有構建模塊。首先，這里假設運行代碼的機器至少有 2 臺 GPU。為此，這里要使用 PyTorch 的 DataParallel API 來將策略模型放在多個 GPU 核心上，每個 GPU 核心都有該模型的一個副本。然后將批量數據分散在這些 GPU 核心上完成處理。

部分代碼截圖。完整代碼塊參見 GitHub。

Part 7：訓練設置和執行

這一節，我們將所有組件組合在一起，完成設置并開始訓練。

首先，加載預訓練的模型和 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 緩存，優化模型的內存效率。
• 步驟 1：在微調之前評估模型，以建立基線準確性。
• 步驟 2：使用 train_with_grpo 函數和我們定義的獎勵函數（format_reward 和 correctness_reward，合并為 combined_reward）執行強化學習微調。這里使用了多臺 GPU 訓練模型。
• 步驟 3：將最終的微調模型和 tokenizer 保存到磁盤。

GRPO 訓練流程使用的超參數如下。

訓練配置

以下參數設定了使用上面的 GRPO 算法實現強化學習微調運行的配置：

• num_iteratinotallow=1：從當前策略模型創建新參考模型的外部迭代次數。一次迭代是指在整個數據集上執行一次通過。
• num_steps=500：訓練循環將執行最多 500 個步驟，每個步驟處理一批樣本。
• batch_size=7：在 8 臺 GPU 的情況下，每個步驟每批處理 7 個樣本，每臺 GPU 上放置 1 個樣本。使用一個 GPU (0) 被 DataParallel 用作主節點來聚合梯度并收集輸出。
• num_generatinotallow=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 數據執行的策略更新次數。在這里，我們每批只執行一次更新。
• epsilnotallow=0.1：GRPO 的 PPO 組件的 clipping 參數。這可以防止策略在單次更新中發生太大的變化。

在微調之前和之后都會對模型進行評估，以衡量準確率的提高情況。最后，將微調后的模型保存到 grpo_finetuned_model 目錄中。

部分代碼截圖。完整代碼塊參見 GitHub。

教程中還給出了詳細的執行情況，可作參考。

下面我們也簡單看看其訓練過程。

首先，初始配置后，我們可以看到運行 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（完全分片數據并行）等模型并行工具。

下面加載和測試已經微調的模型：

完整代碼見原筆記本

加載完成后測試一下，首先問問 1+1 等于幾：

可以看到，模型反復思考了很多次，終于認定確實等于 2。

多次測試后還可以發現，該模型沒有學會生成序列結束（EOS）token，因此即使在 </answer> token 之后，輸出序列仍會繼續。這是預期的行為，因為我們使用的獎勵函數中沒有包含一個用于停止生成的獎勵。我們也沒有執行監督微調步驟 —— 該步驟可以讓模型學會在 </answer> 之后立即生成 EOS。

你對這篇代碼密集的教程怎么看？有沒有讓你產生在自己的電腦上實現 GRPO 的想法？