強化學習 (RL) 是交互式 AI 的支柱。它對于教導智能體推理和從人類偏好中學習、實現多輪工具使用等至關重要。本文將介紹 NVIDIA NeMo-RL,這是一個新的開源后訓練庫,旨在支持從單 GPU 原型到數千個 GPU 大型模型的各種內容,并輕松編排多組件 RL 工作流。
NeMo-RL 是 NVIDIA NeMo 框架的一部分,包括與 Hugging Face 上的模型的原生集成、優化的訓練和推理、DPO 和 GRPO 等熱門算法,以及基于光線的編排。雖然當前的 v0.2.1 版本支持大小高達 320 億個參數的模型,但正在進行的開發旨在將支持擴展到更大的模型。
NeMo-RL 的一個關鍵設計原則是其靈活的后端架構,支持多個訓練和部署后端。對于訓練后端,該庫目前支持具有 PyTorch 原生并行的 Hugging Face 模型,并且即將推出 Megatron-Core 后端,以支持具有高級并行策略的大型模型。
NeMo-RL 使用 vLLM 后端進行生成,可輕松擴展到其他生成后端,包括 NVIDIA TensorRT-LLM 和 SGLang。整體設計可確保高級算法的實現始終與后端實現細節無關,每個后端都在自己的隔離環境中運行,同時遵循標準化的訓練或生成接口。此架構支持從單 GPU 原型無縫擴展到數千個 GPU 部署,而無需修改算法代碼。
在本文中,我們將專門探討如何使用 NeMo-RL,通過組相對策略優化 (GRPO) 強化學習算法無縫再現 DeepScaleR-1.5 B recipe。
使用 NeMo-RL 訓練高性能推理模型
最近,DeepSeek-R1 和 OpenAI O1 等長思維鏈 (CoT) 推理模型引起了廣泛關注。這些模型在各種具有挑戰性的領域具有非常先進的語言模型功能。以下各節將展示如何使用 NeMo-RL 訓練這些高性能推理模型。
我們遵循 DeepScaleR recipe,它提供了一個數據集和方法來訓練推理模型來解決困難的數學問題。特別是,我們使用 GRPO 訓練 Qwen-1.5 B,在競技性學術數學基準 AIME 2024 上達到 OpenAI O1 級別。
分步訓練過程
由于生成時間較長,訓練長 CoT 推理模型可能會非常緩慢,DeepScaleR 首先使用較短的最大序列長度來訓練,然后逐漸增加使用的最大序列長度。具體來說,DeepScaleR 有三個訓練步驟:8K 上下文長度、16K 上下文長度和 24K 上下文長度。這種方法還有助于控制推出序列長度的長尾分布。
使用 NeMo-RL 運行此訓練非常簡單,只需三個步驟。
第 1 步:設置
克隆存儲庫,然后安裝 uv Python 包。借助 uv,您可以快速創建獨立的虛擬環境,盡管依賴項可能存在沖突,同時還可以與 Ray 進行原生集成。
git clone git@github.com:NVIDIA - NeMo / RL.git cd nemo - rl pip install uv |
第 2 步:訓練
使用 DeepSeek-R1-XXDISTILXX-Qwen-1.5 B 進行訓練時,首先使用最大上下文長度 8K,然后使用最大上下文長度 16K,然后使用最大上下文長度 24K。NeMo-RL 以原生方式與 Hugging Face 模型集成,以便用戶直接指定模型選擇。配置文件用于指定 DeepScaleR 數據集和正確的 GRPO 超參數。
uv run examples / run_grpo_math.py - - config = examples / configs / grpo - deepscaler - 1.5b - 8K .yaml uv run examples / run_grpo_math.py - - config = examples / configs / grpo - deepscaler - 1.5b - 16K .yaml policy.model_name = / path / to / 8K / checkpoint / hf uv run examples / run_grpo_math.py - - config = examples / configs / grpo - deepscaler - 1.5b - 24K .yaml policy.model_name = / path / to / 16K / checkpoint / hf |
policy: # Qwen/Qwen2.5-1.5B has tied weights which are only supported with dtensor policy with tp size 1 (https: / / github.com / NVIDIA - NeMo / RL / issues / 227 ) model_name: "Qwen/Qwen2.5-1.5B" tokenizer: name: ${policy.model_name} ## specify if you'd like to use a tokenizer different from the model's default train_global_batch_size: 512 train_micro_batch_size: 4 generation_batch_size: 32 # Only used when generating using HF backend logprob_batch_size: 4 max_total_sequence_length: 512 precision: "bfloat16" fsdp_offload_enabled: false activation_checkpointing_enabled: false |
第 3 步:評估
將檢查點轉換為 Hugging Face 格式并評估模型。請注意,我們會在整個訓練過程中不斷評估模型。應如下所示指定 Hugging Face 檢查點的模型配置、模型位置和所需位置。
uv run examples / convert_dcp_to_hf.py - - config = results / grpo - deepscaler - 1.5b - 8K / step_xx / config.yaml - - dcp - ckpt - path = results / grpo - deepscaler - 1.5b - 8K / step_xx / policy / weights - - hf - ckpt - path = results / grpo - deepscaler - 1.5b - 8K / step_xx / hf uv run examples / run_eval.py \ generation.model_name = results / grpo - deepscaler - 1.5b - 8K / step_xx / hf |
結果
圖 2 顯示了使用 NeMo-RL 的訓練曲線。只需 400 步即可獲得 0.65 的訓練獎勵。

圖 3 顯示了 AIME 2024 在整個訓練過程中的評估結果,最終超過了 OpenAI O1。

如需先了解 DeepScaleR recipe 的功能并可視化 Tensorboard 日志,請查看 NeMo-RL DeepScaleR Tensorboard Viewer Google Colab。
開始使用 NeMo-RL
NeMo-RL 是一個可擴展的后訓練庫,專為單個 GPU 到數千個模型而設計。它包括與 Hugging Face 的無縫集成、具有靈活性的模塊化設計,以及使用 Ray 的高效資源管理。
要使用 NeMo-RL 開始您自己的強化學習實驗,請探索開源 NVIDIA NeMo – RL GitHub 資源庫,其中包含詳細的文檔、示例腳本和配置文件。您還可以嘗試 DeepScaleR 和 OpenMathInstruct-2 示例以了解詳情。
?