强化学习 (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 示例以了解详情。
