• <xmp id="om0om">
  • <table id="om0om"><noscript id="om0om"></noscript></table>
  • 開發與優化

    使用 NVIDIA NeMo-RL 進行強化學習:通過 GRPO 復制 DeepScaleR 配方

    強化學習 (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 的訓練獎勵。

    Chart showing training reward versus training steps for the DeepScaleR Qwen1.5B recipe using NeMo-RL. The curve demonstrates consistent reward improvement, reaching an average of 0.65 reward score around step 400.
    圖 2。在 DeepScaleR Qwen1.5 B recipe 上使用 NeMo-RL 訓練曲線

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

    Line chart showing AIME24 evaluation scores plotted against training steps for the Qwen1.5B recipe trained using NeMo-RL. The curve indicates progressive improvement in performance, eventually surpassing the OpenAI O1 baseline score of 40 on the AIME24 benchmark.
    圖 3。使用 NeMo-RL 訓練的 Qwen1.5 B recipe 的 AIME 2024 評估分數

    如需先了解 DeepScaleR recipe 的功能并可視化 Tensorboard 日志,請查看 NeMo-RL DeepScaleR Tensorboard Viewer Google Colab。

    開始使用 NeMo-RL

    NeMo-RL 是一個可擴展的后訓練庫,專為單個 GPU 到數千個模型而設計。它包括與 Hugging Face 的無縫集成、具有靈活性的模塊化設計,以及使用 Ray 的高效資源管理。

    要使用 NeMo-RL 開始您自己的強化學習實驗,請探索開源 NVIDIA NeMo – RL GitHub 資源庫,其中包含詳細的文檔、示例腳本和配置文件。您還可以嘗試 DeepScaleROpenMathInstruct-2 示例以了解詳情。

    ?

    0

    標簽

    人人超碰97caoporen国产