NeMo RL: 深度强化学习库的最佳实践

1. 项目介绍

NeMo RL 是由 NVIDIA 开发的一个可扩展且高效的深度强化学习库，适用于从单个 GPU 到数千个 GPU，从小型模型到超过 1000 亿参数的大模型。NeMo RL 提供了与 Hugging Face 无缝集成的能力，支持各种并行技术的高性能实现，以及使用 Ray 进行资源管理，确保了在不同硬件配置下的可扩展性和灵活性。

2. 项目快速启动

环境准备

首先，需要克隆 NeMo RL 项目：

git clone [e-mail]:NVIDIA/NeMo-RL.git nemo-rl
cd nemo-rl

安装 uv 工具，用于环境隔离和快速启动：

pip install uv

单节点运行

以下是一个单节点运行的示例，使用 1B 参数的模型进行数学任务训练：

uv run python examples/run_grpo_math.py

默认情况下，这个命令使用了 examples/configs/grpo_math_1B.yaml 配置文件。如果需要自定义参数，可以使用命令行覆盖。例如，要在 8 个 GPU 上运行，可以使用：

uv run python examples/run_grpo_math.py \
cluster.gpus_per_node=8

多节点运行

以下是一个多节点运行的示例，使用 8B 参数的模型进行数学任务训练：

NUM_ACTOR_NODES=2
COMMAND="uv run ./examples/run_grpo_math.py --config examples/configs/grpo_math_8B.yaml cluster.num_nodes=2 checkpointing.checkpoint_dir='results/llama8b_2nodes' logger.wandb_enabled=True logger.wandb.name='grpo-llama8b_math'"
CONTAINER=YOUR_CONTAINER
MOUNTS="$PWD:$PWD"
sbatch --nodes=${NUM_ACTOR_NODES} --account=YOUR_ACCOUNT --job-name=YOUR_JOBNAME --partition=YOUR_PARTITION --time=4:0:0 --gres=gpu:8 ray.sub

请根据实际情况替换 YOUR_CONTAINER, YOUR_ACCOUNT, YOUR_JOBNAME, YOUR_PARTITION 等变量。

3. 应用案例和最佳实践

案例一：数学任务

NeMo RL 提供了一个数学任务案例，使用 GRPO 算法在 OpenInstructMath2 数据集上进行训练。

案例二：滑动拼图游戏

NeMo RL 也支持多轮生成和训练，例如训练模型玩滑动拼图游戏：

uv run python examples/run_grpo_sliding_puzzle.py

最佳实践

• 使用 Ray 进行资源管理，确保在不同的硬件配置下都能高效运行。
• 利用 Hugging Face 集成，轻松使用各种预训练模型。
• 采用 PyTorch 原生的 FSDP2、TP 和 SP 技术，实现高效的模型训练。

4. 典型生态项目

• Hugging Face: 提供了广泛的预训练模型和工具，与 NeMo RL 无缝集成。
• Ray: 用于分布式训练的资源管理系统，与 NeMo RL 深度集成，提高训练效率。
• Megatron Core: 支持大型模型的高性能并行技术，适用于超过 100B 参数的模型。