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

mkdir rl-example
cd rl-example

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

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

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

grpo:
  num_prompts_per_step: 64  # 1stepに時間がかかりすぎたため、512から64に変更
  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: 10000
  normalize_rewards: true
  use_leave_one_out_baseline: false
  val_period: 5
  val_at_start: false
  overlong_filtering: true  # false -> trueへ変更。生成長が長すぎるものはフィルタリング
  max_val_samples: 960
  val_batch_size: 960
  seed: 42
  use_dynamic_sampling: true
  batch_multiplier: 2 # Multiplier for dataloader batch size calculation (batch_multiplier × num_prompts_per_step). Following DAPO dynamic sampling, the actual training batch size equals num_prompts_per_step × num_generations_per_prompt.
  dynamic_sampling_max_gen_batches: 10
  reward_scaling:
    enabled: false  # true -> falseへ変更。報酬をスケーリングしない
    source_min: 0.0
    source_max: 1.0
    target_min: -1.0
    target_max: 1.0
  reward_shaping:
    enabled: false  # ソフトペナルティを設定する場合はTrue 例.16k + 4k (overlong_buffer_length) 生成が16kを超える場合、徐々に報酬を減らす
    overlong_buffer_length: 4096
    overlong_buffer_penalty: 1
    max_response_length: 20480

  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
    in_flight_weight_updates: false # Set to true to enable in-flight weight updates
    recompute_kv_cache_after_weight_updates: false # Set to true to recompute kv cache after in-flight-weight-updates

  seq_logprob_error_threshold: null

loss_fn:
  reference_policy_kl_penalty: 0.0
  # Can be set to k1, k2, k3
  # For more details, see http://joschu.net/blog/kl-approx.html
  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
  # Async GRPO requires importance sampling correction enabled
  # Set to true when async_grpo.enabled is true
  use_importance_sampling_correction: false
  truncated_importance_sampling_ratio: null
  sequence_level_importance_ratios: false
  token_level_loss: true
  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/dapo-nanov2-9b
  metric_name: "val:accuracy" # one of "val:" or "train:" followed by the metric name
  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-Nano-9B-v2-Japanese
  tokenizer:
    name: ${policy.model_name} ## specify if you'd like to use a tokenizer different from the model's default
  train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}
  train_micro_batch_size: 1
  generation_batch_size: 32 # Only used when generating using HF backend
  logprob_batch_size: 1
  max_total_sequence_length: 8192 # 今回の検証では8192に設定
  precision: "bfloat16"
  logprob_chunk_size: 2048

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

  megatron_cfg:
    enabled: true
    empty_unused_memory_level: 1
    activation_checkpointing: false
    bias_activation_fusion: false
    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: true
    freeze_moe_router: true
    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-06
      min_lr: 1.0e-06
      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
      optimizer_cpu_offload: false
      optimizer_offload_fraction: 0.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: 1.0e-07

    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:
      PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True"
      NCCL_DEBUG: "WARN"
      CUDA_DEVICE_MAX_CONNECTIONS: "1"

  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

  make_sequence_length_divisible_by: ${policy.dtensor_cfg.tensor_parallel_size}
  max_grad_norm: 1.0

  optimizer: null # remove default FSDP optimizer
  scheduler: null

  offload_optimizer_for_logprob: false # Only useful for non-colocated generation since colocated generation will always offload optimizer to cuda before refit

  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.85  # 0.6から0.85に変更
      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

    vllm_kwargs:
      mamba_ssm_cache_dtype: "float32"  # mambaベースのnemotron用
      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: 2048
  prompt_file: null
  system_prompt_file: null
  shuffle: true
  num_workers: 1
  dataset_name: DAPOMath17K

env:
  math:
    num_workers: 8
    math_verify_impl: "dapo_math_verify"

logger:
  log_dir: "logs/dapo-nanov2-9b"  # 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
  swanlab_enabled: false # Disable SwanLab logging
  monitor_gpus: true  # If true, will monitor GPU usage and log to wandb and/or tensorboard
  wandb:
    project: "grpo-dev"
    name: "dapo-nanov2-9b-dev-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

export WANDB_API_KEY="YOUR-WANDB-API-KEY"
uv run examples/run_grpo_math.py --config examples/configs/custom_dapo-nanov2-9b.yaml

CKPT_DIR=results/dapo-nanov2-9b/step_260
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-Nano-9B-v2-Japanese-DAPO-step_260