NVIDIA NeMo RL ではじめる強化学習 – Gym 編

本記事では、NVIDIA NeMo RL ではじめる強化学習 – Math 編 に続き、NVIDIA NeMo RL と NVIDIA NeMo Gym を使用して、LLM に強化学習 (RL) を適用する方法をご紹介します。

NeMo Gym とは

NeMo Gym は、NVIDIA がオープン ソースとして公開している環境を使用して、モデル/エージェントを評価、改善するためのライブラリです。NeMo Gym は、環境開発、評価とトレーニングのスケーラブルな実行、人気のベンチマークやトレーニング環境のコレクションといったインフラストラクチャを提供します。

ここで出てくる環境とは、エージェントがタスクを完了するために相互作用するシステム全体のことです。環境は、データセット (解決すべきタスク)、エージェント ハーネス (モデルが世界とどのように相互作用するか)、検証器 (タスク完了のスコアリング)、および状態 (タスクごとの実行コンテキスト) で構成されます。

NeMo Gym を使うべきタイミング

• ステートフルな環境 (コード実行、ツール呼び出し、サンドボックスなど) でモデルやエージェントを評価する必要があるとき
• 共通の環境と検証ツールを使用して、チーム間で再現可能な評価を実施したいとき
• 大規模な環境を使用する必要があり、タスクごとに複数回の繰り返し、またはトレーニングのための数千の同時リクエストが必要なとき
• 評価、エージェント最適化、トレーニングの間をシームレスに移行したいとき

NeMo Gym が提供するもの

• エージェント、環境、タスク、検証器のためのモジュール式で拡張可能なインターフェイス
• 人気のベンチマークとトレーニング環境の環境ハブ
• 独自のエージェントを使用するか、組み込みハーネスから選択可能
• 数千の同時実行環境に対応可能
• お好みの強化学習フレームワークでトレーニング
• 実戦 (Nemotron の訓練) で実証済み

NeMo Gym は、広範なエコシステムに統合されています。詳細については、エコシステム ページをご覧ください。

• 環境ライブラリ: NeMo Gym 環境と他のライブラリの環境およびベンチマークをシームレスに組み合わせることができます。(例: Aviary、Harbor、OpenEnv、Reasoning Gym、verifiers)
• トレーニング フレームワーク ライブラリ: SFT および RL トレーニングに使用する環境。(例: NeMo RL、Unsloth、verl)
• エージェント ハーネス: 評価およびトレーニング用のエージェント ハーネスがすぐに利用可能です。(例: OpenHands、Mini SWE Agent、LangGraph)

NeMo RL と NeMo Gym の関係性

RL のトレーニング中、エージェントは環境内でタスクを実行し、その後、検証器によってスコアリングされて、報酬が計算されます。エージェントの軌跡と報酬スコアは、トレーニング エンジンによってモデルの重みを更新するために利用されます。トレーニング中は複数の環境が同時に使用されることが多く、これは一般的にマルチ環境トレーニングまたはマルチ検証器トレーニングと呼ばれます。

下図のように、NeMo RL と NeMo Gym は「モデルの学習 (最適化)」と「環境 (データ生成/検証)」を完全に分離する設計になっており、モデルの学習/推論エンジンには NeMo RL を使用し、OpenAI Responses API をインターフェイスに NeMo Gym の環境と通信し、ツールの利用や一連の軌跡の収集 (ロールアウト生成)、報酬計算を実行します。

これによって、学習アルゴリズムを変えずに環境を変更したり、環境を変えずに学習アルゴリズムを変更するなどの柔軟な利用が可能になります。

NeMo Gym は複数のトレーニング フレームワークをサポートしていますが、NeMo RL から使用すると Nemotron の学習で実証されたパイプライン、豊富なサンプルを活用することができるのでおすすめです。

NeMo RL + Gym チュートリアル

本記事では、NeMo RL と NeMo Gym を組み合わせてマルチ環境で GRPO を実行します。データ、スクリプトはサンプルとして用意されているため実行自体は簡単です。

モデル: nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16

• 学習バックエンド: Megatron
• 生成バックエンド: vLLM
• RLアルゴリズム: GRPO (一部改変)
• タスク: マルチタスク (nvidia/Nemotron-3-Nano-RL-Training-Blend)
• 同期学習 + コロケーション (学習と生成で同じ GPU を使用する)

nvidia/NVIDIA-Nemotron-Nano-9B-v2-Japanese で試行しましたが、Nemotron v2 モデルを使用する場合、RL と Gym リポジトリ内の複数箇所の実装を変更する必要があったため、チュートリアルでは軽量かつ簡潔に実行できる Nemotron-3-Nano-4B を採用しました。

今回のチュートリアルの検証環境は以下の条件で行っております。こちらの構成は一例であり、NeMo RL は単一 GPU からマルチノード環境で実行することが可能です。

• ハードウェア
  • NVIDIA DGX H200
  • GPU: 8 x NVIDIA H200 141 GB GPU (driver version: 580.126.09)
  • CPU:  Intel Xeon Platinum 8480C
  • システム メモリ: 2 TB
• ソフトウェア
  • OS: Ubuntu 24.04.3 LTS
  • Container: nvcr.io/nvidia/nemo-rl:v0.4.0.nemotron_3_nano

NVIDIA Nemotron 3 Nano 4B

Nemotron 3 Nano 4B は、Nemotron Elastic の技術を使用して、Nemotron Nano 9B v2 から派生したモデルです。わずか 40 億のパラメーターを持つ Nemotron 3 Nano 4B は、NVIDIA Jetson プラットフォーム (Jetson Thor/Jetson Orin Nano) だけでなく、NVIDIA DGX Spark や NVIDIA RTX GPU 上でエッジ コンピューティングを実行するのに十分なコンパクトさを備えています。これにより、応答時間の短縮、データ プライバシーの強化、柔軟な導入が可能になり、推論コストを低く抑えることができます。さらなる詳細はリリース ブログを参照してください。

事前準備

以下のコマンドで作業用のディレクトリを作成し、移動します。

mkdir rl-example
cd rl-example

リポジトリの Clone

以下のコマンドを実行し、NeMo-RL の実行環境を整えます。

git clone --branch nano-v3 --recursive git@github.com:NVIDIA-NeMo/RL.git nemo-rl-nano-v3
cd nemo-rl-nano-v3
uv venv

コンテナーの準備

このチュートリアルでは NGC 上にある NeMo-RL のコンテナーを使用します。NeMo-RL のコンテナーには、NeMo Gym が含まれており、特別な環境構築を必要とせずにそのまま利用できます。ベアメタル (コンテナー外) で使用する場合は、こちらを参照して別途環境構築してください。

docker pull nvcr.io/nvidia/nemo-rl:v0.4.0.nemotron_3_nano

リポジトリと NGC コンテナー イメージの tag を揃えておかないと実行時にエラーや不具合が発生する可能性があります。

データの準備

今回使用するデータセット (nvidia/Nemotron-3-Nano-RL-Training-Blend) をダウンロードします。

# Download RL data blend
uvx --from huggingface-hub hf download nvidia/Nemotron-3-Nano-RL-Training-Blend --repo-type dataset --local-dir=datasets

# Fill in placeholders in dataset
chmod +x datasets/create_nanov3_jsonl.py
./datasets/create_nanov3_jsonl.py --input datasets/train.jsonl --output datasets/train-full.jsonl

# Use the last 1000 rows for validation
head -n -1000 datasets/train-full.jsonl > datasets/train-split.jsonl
tail -n 1000 datasets/train-full.jsonl > datasets/val-split.jsonl

このデータセットは NVIDIA-Nemotron-3-Nano-30B-A3B の RL に使用されたデータセットで、SFT 後のチェックポイントを使用して、すでに正答できる問題のフィルタリングおよび、問題の難易度が徐々に高くなるように順序付けされています (カリキュラムになっています)。したがって、検証データセットに使用している末尾の 1,000 件は最も難易度の高いデータ ポイントになります。

Docker コンテナーの起動

以下のコマンドでコンテナーを起動します。

docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
  -v $PWD:$PWD \
  -w $PWD \
  -it nvcr.io/nvidia/nemo-rl:v0.4.0.nemotron_3_nano bash

Config ファイルの作成

NeMo RL で NeMo Gym タスクを実行するため、yaml ベースの Config ファイルを用意します。ベースはリポジトリ内の example にある nemo_gym/grpo_nanov3.yaml を参照しながら以下の変更を適用しています。

• 使用するモデルを NVIDIA-Nemotron-3-Nano-4B-BF16 へ変更
• max_total_sequence_length を 49152 から 8192 へ変更
• モデル並列化の設定を NVIDIA-Nemotron-3-Nano-4B-BF16やmax_total_sequence_length に併せて変更

この yaml ファイルを examples/configs の下に custom_grpo_nanov3_4b.yaml という名前で配置します。今回は examples/configs の下を選びましたが、特に指定はありません。

NeMo RL から NeMo Gym を使用する際は、config 内の env キーの should_use_nemo_gym を True にします。nemo_gym キーの下にあるものは NeMo Gym の設定で config_paths に各 Gym 環境の設定ファイルを指定します。各設定はこの yaml から上書きすることが可能です。さらなる詳細はドキュメントを参照してください。

実行するリポジトリのバージョンによって Yaml ファイルで使われているキーの名前に微妙な差があったり、新たなキーが追加されている場合があります。使用するバージョンの Config ファイルを参考にしてください。

# GRPO Algorithm Configuration for training Nano v3
grpo:
  num_prompts_per_step: 128
  num_generations_per_prompt: 16
  max_rollout_turns: 1 # for multi-turn rollouts. Math Environments just have 1 turn (answering the question)
  max_num_epochs: 1
  max_num_steps: 1000000
  normalize_rewards: true
  use_leave_one_out_baseline: true
  val_period: 5
  val_at_start: False
  val_at_end: False
  overlong_filtering: true
  max_val_samples: null
  val_batch_size: 256
  seed: 42
  async_grpo:
    enabled: false # Set to true to enable async training mode
    # Max age (in training steps) for trajectories used in training
    max_trajectory_age_steps: 1

  batch_multiplier: 1
  use_dynamic_sampling: False
  reward_shaping:
    enabled: False
  reward_scaling:
    enabled: False

  seq_logprob_error_threshold: 2

loss_fn:
  reference_policy_kl_penalty: 0
  reference_policy_kl_type: k3
  kl_input_clamp_value: null
  kl_output_clamp_value: null
  ratio_clip_min: 0.2
  ratio_clip_max: 0.28
  ratio_clip_c: null
  # (default off) loss formulation improvements (docs/guides/grpo.md#loss)
  use_on_policy_kl_approximation: false
  use_importance_sampling_correction: true
  sequence_level_importance_ratios: false
  token_level_loss: true
  truncated_importance_sampling_ratio: 2.0
  force_on_policy_ratio: false  # Set to true to force ratio=1.0 (requires train_global_batch_size == num_prompts_per_step * num_generations_per_prompt)

checkpointing:
  enabled: true
  checkpoint_dir: "results/grpo-nano3-4b"
  metric_name: "val:total_reward/mean"
  higher_is_better: true
  keep_top_k: 5
  save_period: 5
  checkpoint_must_save_by: null
  save_optimizer: true

policy:
  model_name: "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" # Nano 3 4Bに変更
  tokenizer:
    name: ${policy.model_name} ## specify if you'd like to use a tokenizer different from the model's default
  train_global_batch_size: 2048
  train_micro_batch_size: 1
  generation_batch_size: 64 # Only used when generating using HF backend
  logprob_batch_size: 1
  max_total_sequence_length: 8192  # max_model_lenと一致させる
  precision: "bfloat16"
  logprob_chunk_size: 2048

  dtensor_cfg:
    _v2: true
    enabled: false
    cpu_offload: False
    sequence_parallel: false
    activation_checkpointing: false
    tensor_parallel_size: 1
    context_parallel_size: 1
    custom_parallel_plan: null

  megatron_cfg:
    enabled: true
    empty_unused_memory_level: 1
    activation_checkpointing: true
    bias_activation_fusion: False
      # converter_type: "Qwen2ForCausalLM"
    tensor_model_parallel_size: 2
    expert_tensor_parallel_size: 1
    expert_model_parallel_size: 1
    pipeline_model_parallel_size: 1
    num_layers_in_first_pipeline_stage: null
    num_layers_in_last_pipeline_stage: null
    context_parallel_size: 1
    pipeline_dtype: ${policy.precision}
    sequence_parallel: false  # SP=falseに変更
    freeze_moe_router: true
      # moe_router_dtype: "fp64"
    moe_router_dtype: "fp32"
    moe_router_load_balancing_type: "none" # "seq_aux_loss" causes logprob error divergence for grpo
    moe_router_bias_update_rate: 1e-3
    moe_permute_fusion: true
    moe_enable_deepep: false
    moe_token_dispatcher_type: "alltoall"
    moe_aux_loss_coeff: 0.0
    moe_router_enable_expert_bias: true
    #gives ~20% training perf speedup with sequence packing
    apply_rope_fusion: True
    defer_fp32_logits: True
    track_moe_metrics: True
    moe_per_layer_logging: True
    moe_shared_expert_overlap: false
    gradient_accumulation_fusion: false

    optimizer:
      optimizer: "adam"
      lr: 1.0e-6
      min_lr: 1.0e-6
      weight_decay: 0.0
      bf16: true
      fp16: false
      params_dtype: "float32"

      #adam
      adam_beta1: 0.9
      adam_beta2: 0.999
      adam_eps: 1e-8

      #sgd
      sgd_momentum: 0.9

      clip_grad: ${policy.max_grad_norm}

      #distributed optimizer
      use_distributed_optimizer: true
      use_precision_aware_optimizer: true

      optimizer_cpu_offload: False
      optimizer_offload_fraction: 0

    scheduler:
      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
      weight_decay_incr_style: "constant"
      lr_decay_style: "constant"
      lr_decay_iters: null
      lr_warmup_iters: 10
      lr_warmup_init: 3.0e-7

    distributed_data_parallel_config:
      grad_reduce_in_fp32: false
      overlap_grad_reduce: true
      overlap_param_gather: true
      average_in_collective: false
      use_custom_fsdp: false
      data_parallel_sharding_strategy: "optim_grads_params"

    env_vars: null

  # See docs/design-docs/sequence-packing-and-dynamic-batching.md
  # for more details on dynamic batching and sequence packing.
  dynamic_batching:
    enabled: False
    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
    sequence_length_round: 64

  sequence_packing:
    enabled: True
    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
    algorithm: "modified_first_fit_decreasing"
    sequence_length_round: 64

  # makes the training sequence length divisible by the tensor parallel size
  # this is useful for sequence parallel training
  make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size}
  max_grad_norm: 1.0

  optimizer: null # remove default FSDP optimizer
  scheduler: null

  offload_optimizer_for_logprob: False

  generation:
    backend: "vllm"
    max_new_tokens: ${policy.max_total_sequence_length}
    temperature: 1.0
    top_p: 1.0
    top_k: null
    stop_token_ids: null
    stop_strings: null
    vllm_cfg:
      # NB: can re-enable prefix cache on vllm >= 0.11.2.
      # enable_prefix_caching: false
      async_engine: false
      kv_cache_dtype: auto
      precision: ${policy.precision}
      tensor_parallel_size: 1
      pipeline_parallel_size: 1
      expert_parallel_size: 1  # When EP > 1, EP must be a multiple of TP since vLLM's EP = DP * TP
      gpu_memory_utilization: 0.8  # 0.5から0.8に変更
      max_model_len: ${policy.max_total_sequence_length}
      # when enforce_eager is False, it is optional to set ++policy.generation.vllm_kwargs.compilation_config.backend=eager for better accuracy,
      # with the flag, vllm will use the custom CUDA kernels instead of the Triton kernels generated by torch.compile
      # for more details, see convergence issue https://github.com/NVIDIA-NeMo/RL/issues/998
      enforce_eager: False
      use_deep_gemm: False
      num_last_layers_in_bf16: 0
      num_first_layers_in_bf16: 0
      expose_http_server: true
      http_server_serving_chat_kwargs:
        enable_auto_tools: true
        tool_parser: qwen3_coder
        reasoning_parser: deepseek_r1

    vllm_kwargs:
      mamba_ssm_cache_dtype: "float32"
      compilation_config:
        # when enforce_eager is False, set ++policy.generation.vllm_kwargs.compilation_config.backend=eager for better accuracy,
        # with the flag, vllm will use the custom CUDA kernels instead of the Triton kernels generated by torch.compile
        # for more details, see convergence issue https://github.com/NVIDIA-NeMo/RL/issues/998
        backend: eager
    colocated:
      # true: generation shares training GPUs
      # false: uses dedicated generation resources
      enabled: true
      # only relevant when enabled is false
      resources:
        gpus_per_node: null # Decides num gpus to be dedicated to generation when there is one node in the cluster i.e cluster.num_nodes == 1
        num_nodes: null # Decides number of nodes to be dedicated to generation

data:
  max_input_seq_length: null
  shuffle: False
  num_workers: 1
  use_multiple_dataloader: false # 追加

  train_jsonl_fpath: "datasets/train-split.jsonl"  # データパス
  validation_jsonl_fpath: "datasets/val-split.jsonl"  # データパス

  default:
    dataset_name: NemoGymDataset
    env_name: "nemo_gym"
    prompt_file: null
    system_prompt_file: null
    processor: "nemo_gym_data_processor"

env:
  should_use_nemo_gym: true
  nemo_gym:
    config_paths:
    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml  # Required! And it must be *for_training
    - resources_servers/math_with_judge/configs/math_with_judge.yaml
    - resources_servers/code_gen/configs/code_gen.yaml
    - resources_servers/workplace_assistant/configs/workplace_assistant.yaml
    - resources_servers/mcqa/configs/mcqa.yaml
    - resources_servers/instruction_following/configs/instruction_following.yaml
    - resources_servers/structured_outputs/configs/structured_outputs_json.yaml
    math_with_judge:
      resources_servers:
        math_with_judge:
          judge_model_server:
            name: policy_model
          should_use_judge: false
    code_gen:
      resources_servers:
        code_gen:
          num_processes: 1024
          unit_test_timeout_secs: 10
          debug: false
    # 再実行回数を制限したいため追加
    code_gen_simple_agent:
      responses_api_agents:
        simple_agent:
          max_steps: 1
    math_with_judge_simple_agent:
      responses_api_agents:
        simple_agent:
          max_steps: 1
    mcqa_simple_agent:
      responses_api_agents:
        simple_agent:
          max_steps: 1
    instruction_following_simple_agent:
      responses_api_agents:
        simple_agent:
          max_steps: 1
    structured_outputs_simple_agent:
      responses_api_agents:
        simple_agent:
          max_steps: 1
    workplace_assistant_simple_agent:
      responses_api_agents:
        simple_agent:
          max_steps: 6

logger:
  log_dir: "logs/grpo-nano3-4b"  # Base directory for all logs
  num_val_samples_to_print: 0 # Number of validation samples to pretty print on terminal
  wandb_enabled: true
  tensorboard_enabled: false
  mlflow_enabled: false  # Disable MLflow logging
  monitor_gpus: true  # If true, will monitor GPU usage and log to wandb and/or tensorboard
  swanlab_enabled: false # Disable SwanLab logging
  wandb:
    project: "grpo-dev"
    name: "grpo-nano3-4b-logger"
  tensorboard: {}
  mlflow:
    experiment_name: "grpo-dev"
    run_name: "grpo-dev-logger"
  gpu_monitoring:
    collection_interval: 10  # How often to collect GPU usage metrics (in seconds)
    flush_interval: 10  # How often to flush GPU usage metrics to the loggers (in seconds)

cluster:
  gpus_per_node: 8
  num_nodes: 1  # 1ノードに変更

NeMo Gym で実行する各タスクについて

nvidia/Nemotron-3-Nano-RL-Training-Blend では、以下の 6 つの環境を使用しています。マルチステップのworkplace_assistant (デフォルト: 6) を除き、各タスクはシングル ステップです。上記の Config では、シングル ステップ タスクについては max_step を 1 に指定 (デフォルト: 指定なし) しており、これは予備実験にてツール利用の必要のないタスクでツール利用が繰り返し発生し、1 ステップの処理時間が増大したことから追加しています。

• math_with_judge
  • タスク: 数学タスク。
  • データセット: BytedTsinghua-SIA/DAPO-Math-17k, Skywork/Skywork-OR1-RL-Data
• code_gen
  • タスク: Python を使用した競技プログラミング形式のコード生成タスク。
  • データセット: nvidia/Nemotron-RL-coding-competitive_coding
• workplace_assistant
  • タスク: ビジネス シーンを想定したマルチステップのツール利用タスク。
  • データセット: nvidia/Nemotron-RL-agent-workplace_assistant
• mcqa
  • タスク: マルチ ドメイン (物理学、生物学、化学、数学、コンピューター科学、工学、人文科学、法律など) の多肢選択 QA タスク。
  • データセット: nvidia/Nemotron-RL-knowledge-mcqa
• instruction_following
  • タスク: 何かしらの出力制約を順守する指示追従タスク。
  • データセット: nvidia/Nemotron-RL-instruction_following
• structured_outputs
  • タスク: 構造化出力制約 (JSON 形式のスキーマ制約) での出力フォーマット指示追従タスク。
  • データセット: nvidia/Nemotron-RL-instruction_following-structured_outputs

GRPO の実行

以下のコマンドで GRPO を実行します。WANDB で学習ログを記録したい場合は WANDB_API_KEY を環境変数に追加してください。学習プロセスで収集されるメトリクスがとても多いため、利用を推奨します。

export WANDB_API_KEY="YOUR-WANDB-API-KEY"
uv run examples/nemo_gym/run_grpo_nemo_gym.py --config examples/configs/custom_grpo_nanov3_4b.yaml

ジョブの起動時に RuntimeError: Process `XXXXXX` finished unexpectedly! などが出た場合は、ジョブの実行前に以下を実際の実行環境の PATH に合わせて実行してください。`XXXXXX` には任意の nemo_gym 環境名が入ります (例: mcqa など)。

git config --global --add safe.directory /PATH/TO/nemotron_3_nano/3rdparty/Gym-workspace/Gym

以下が 720 steps (1 エポック) 実行した際の学習/検証データセットでの報酬の遷移になります。学習データセットでの報酬は、 (後に記載しているタスク毎にみると上昇/横ばい/下落が混在していますが) 横ばい傾向で綺麗に上がってはいきませんでした。いろいろな要因が考えられますが、データが徐々に難しくなっていくように順序付けされている影響も受けていると考えられます。

一方で、最も難易度の高いデータポイントで構成される検証データセットの報酬は順調に上がっていったことから、このまま学習を続行する判断をしました（もしかしたら潜在的な問題を内在している可能性もあります）。

先ほどは全てのタスクを含めた報酬推移でしたが、学習データセットでの各環境の報酬推移が以下になります。各環境ごとに見ると、報酬推移はまちまちで、横ばい傾向のタスクがあれば、改善/劣化していく傾向のタスクもありました。こちらもデータの順序付けの他に、各データセットの難易度の幅やモデルの能力など複数の要因が影響していると考えられ、学習がうまくいっているのかの判断が難しかったです。

図 7. 学習データセットの各環境での報酬推移

検証データセットでの各環境の報酬推移が以下になります。こちらは学習データセットとは異なり、いずれのタスクも改善していく傾向にありました。今回はこの検証データセットでの改善が学習を継続する大きな要因となりました。

図 8. 検証データセットの各環境での報酬推移

前回の Math 編と同様に、以下に学習中の学習/推論エンジン間の対数確率誤差を載せます。波打って推移していますが、値自体は 1.05 を超えることはありませんでした。

学習が終わったチェックポイントはリポジトリ内の変換スクリプトで HF 形式へ変換できます。今回使用したバージョンでは以下になります。評価には、720 steps の学習中で validation reward が最も高かった 685 steps の checkpoint を使用しました。

CKPT_DIR=results/grpo-nano3-4b/step_685
uv run --extra mcore examples/converters/convert_megatron_to_hf.py --config=$CKPT_DIR/config.yaml --megatron-ckpt-path=$CKPT_DIR/policy/weights/iter_0000000/ --hf-ckpt-path=output/NVIDIA-Nemotron-3-Nano-4B-BF16-Gym-step_685

学習バックエンドによって、使用するスクリプトが異なる点、リポジトリのバージョンによって引数が微妙に異なることがある点に気をつけてください。詳細は対応するバージョンのドキュメント内のセクションを参照してください。

学習後のモデルの評価

学習前後のモデルを swallow-evaluation-instruct Ver. 202604 内の複数のタスクで評価しました。最大コンテキスト長は 32,768 トークン、生成パラメーターは temperature=0.6, top-p=0.95 を使用しています。

結果を見ると、日本語/英語問わず多くのベンチマーク スコアが概ね改善しました。Nemotron 3 Nano の技術レポートでは単一環境での学習は他のベンチマークの低下を招くことがあり、マルチ環境での学習により、安定したベンチマークでの性能向上が得られたと報告されています。今回の実験においても同様の結果が得られました。

このアプローチは、Nemotron 3 Super でも利用されており、Nemotron 3 Super では、多段階 RL で 20 を超える環境が導入されています。さらに Nemotron 3 Ultra では、環境の数が増え続けると各ドメインのシグナルが希薄化するという問題意識から、RL に MOPD (Multi-Teacher On-Policy Distillation) を組み合わせたアプローチがとられています。

チェリーピックですが MMLU-ProX (日本語) の一部の結果を記載します。RL 後に精度が改善したサンプルを分析すると、大半のケースでは以下のように純粋な正答率の改善 (元モデルが誤答し、RL 後に改善) が挙げられます。

[プロンプト]
次の選択問題に答えてください。出力の最後の行には「回答: $選択肢」（鉤括弧は書かない）という形でA, B, C, D, E, F, G, H, I, Jから1つ選んで答えてください。ステップバイステップで考えてから回答してください。

25と55の最大公約数を求めなさい。

A) 275
B) 15
C) 20
D) 50
E) 11
F) 30
G) 5
H) 10
I) 25
J) 35

[NVIDIA-Nemotron-3-Nano-4B-BF16応答]

25と55の最小公倍数（LCM）を見つけるために、それぞれの数の素因数分解を行います。

- 25 = 5²
- 55 = 5 × 11

LCMは、各素因数の最高のべき乗の積です:
- 5の最高のべき乗は5²
- 11の最高のべき乗は11

したがって、LCM = 5² × 11 = 25 × 11 = 275。

したがって、正しい選択はAです。

回答: A

[NVIDIA-Nemotron-3-Nano-4B-BF16 + GRPO 685 steps checkpoint応答]

25と55の最大公約数（GCD）を求めるために、その両方の素因数を分解します：

- 25 = 5 × 5 = 5²
- 55 = 5 × 11

両方の数に共通する因数は5のみです。したがって、最大公約数は5です。

**解答: G**
解答：G

[正解] G

また、NVIDIA-Nemotron-3-Nano-4B-BF16 がすでに指示追従やツール呼び出しタスクで RL を適用したモデルであり改善幅は小さかった (改善全体の 3% ほど) ですが、以下のように元のモデルがコード実行やツール呼び出し、思考過程のまま最終応答に辿りつかないケースが RL によって改善した例もあります。

[プロンプト]
次の選択問題に答えてください。出力の最後の行には「回答: $選択肢」（鉤括弧は書かない）という
形でA, B, C, D, E, F, G, H, I, Jから1つ選んで答えてください。ステップバイステップで考えてから
回答してください。

残存期間18年の8%債券の利回りが9%です。この債券の価格はいくらですか？

A) 91.17　B) 100.00　C) 110.00　D) 88.34　E) 105.68
F) 85.23　G) 95.50　H) 102.56　I) 77.89　J) 120.45

[NVIDIA-Nemotron-3-Nano-4B-BF16応答]

Let's compute the price.

We'll use the standard bond price formula: \(P = \sum_{t=1}^{18} \frac{8}{(1.09)^t} + \frac{100}{(1.09)^{18}}\). Let's compute this.
</think>
Let's compute the price.

We'll use the standard bond price formula: \(P = \sum_{t=1}^{18} \frac{8}{(1.09)^t} + \frac{100}{(1.09)^{18}}\). Let's compute this.
```python
import math
y = 0.09
n = 18
C = 8
FV = 100

PV_coupons = C * (1 - (1+y)**-n) / y
PV_FV = FV * (1+y)**-n
price = PV_coupons + PV_FV
price
```

[NVIDIA-Nemotron-3-Nano-4B-BF16 + GRPO 685 steps checkpoint応答]

債券の価格は、将来のクーポン支払いと元額の現在価値の合計です。

クーポンの価値 = 8％ × 100 = 8
利回り（年率）= 9％ = 0.09
期間 = 18 年

価格 = 8 × (1 - (1.09)^-18)/0.09 + 100 × (1.09)^-18 ≈ 91.17

**回答: A**

[正解] A

モデルの学習中に遭遇したトラブル

ここでは本記事執筆にあたり、遭遇したトラブルをご紹介します。チュートリアルのジョブが複数回に分割されているのは、以下のトラブルに遭遇したことが原因です。

勾配ノルムの爆発

前回記事の Math の実験中には遭遇することがありませんでしたが、マルチ環境で学習した際に device 0, iteration 5: Unexpected result inf (message='found Inf in local grad norm for bucket #0 in backward pass before data-parallel communication collective') といった勾配ノルムの爆発によるエラーが発生することがありました。大抵のケースでは再実行により、前回止まったステップを通過していきましたが、モデルやデータ、Config の組み合わせによっては決まったステップで毎回エラーが発生するという状況も起こり得ます。

NeMo RL のドキュメントではこの要因の一つとして、モデルの入力/出力の埋め込みにおいて、値がほぼ 0 やほぼ同じ値の行 (埋め込みベクトル、トークン) がある場合に、学習中にこれらのトークンに遭遇すると数値不安定性が発生する可能性があるとしており、これらの埋め込みを事前学習時と同様の方法で初期化すると解決に役立つ可能性があるとしています。

NeMo RLのリポジトリにはこれらの埋め込みベクトルを検出、初期化するスクリプトがあるので、問題があった際には調査し、初期化を検討することをお勧めします。

Internal Server Error

処理するリクエスト数が多く、さらに長い生成が多い場合に、vLLMのリクエスト負荷が高まり、aiohttp.client_exceptions.ClientResponseError: 500, message='Internal Server Error', url='http://127.0.0.1:48629/run といったエラーでジョブが止まることがありました。こちらも突発的に発生し、再実行した際には前回止まったステップを通過していきましたが、もし頻繁に発生する場合はロールアウト生成数 (num_prompts_per_step と num_generations_per_prompt) を小さくするなどの変更を検討することをおすすめします。

まとめ

本記事では、NeMo RL と NeMo Gym を使用したマルチ環境での GRPO の実行を紹介しました。NeMo RL と NeMo Gym を使用して強化学習を試していただけると嬉しいです。

関連情報

• 関連ドキュメント:
  • NeMo RL Documentation
  • NeMo Gym Documentation
• リポジトリ:
  • GitHub – NVIDIA-NeMo/RL
  • GitHub – NVIDIA-NeMo/Gym
• コンテナー イメージ: NGC – NVIDIA nemo-rl
• 技術ブログ:
  • NVIDIA NeMo RL を用いた合成データによる Supervised Fine-Tuning (SFT)
  • NVIDIA NeMo RL ではじめる強化学習 – Math 編