# GRPO Algorithm Configuration
grpo:
  num_prompts_per_step: 32
  num_generations_per_prompt: 16
  max_rollout_turns: 1 # for multi-turn rollouts. Math Environments just have 1 turn (answering the question)
  max_num_steps: 1000000
  normalize_rewards: true
  use_leave_one_out_baseline: true
  val_period: 10
  val_at_start: false
  max_val_samples: 256
  val_batch_size: 256

loss_fn:
  reference_policy_kl_penalty: 0.01
  ratio_clip_min: 0.2
  ratio_clip_max: 0.2
  ratio_clip_c: null
  # (default off) loss formulation improvements (docs/guides/grpo.md#loss)
  use_on_policy_kl_approximation: false
  use_importance_sampling_correction: false
  token_level_loss: true

checkpointing:
  enabled: true
  checkpoint_dir: "results/grpo"
  metric_name: "val_reward"
  higher_is_better: true
  keep_top_k: 3
  save_period: 10

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
  refit_buffer_size_gb: 4 # used for refitting inference engine, the unit is GB

  dtensor_cfg:
    enabled: true
    cpu_offload: False
    sequence_parallel: false
    activation_checkpointing: false
    tensor_parallel_size: 1

  # dynamic_batching improves performance by ensuring logprob and training microbatches
  # have a sufficent number of tokens to maximize GPU utilization. Specifically, variable length
  # responses are sorted by sequence length and bucketed into microbatches with a total
  # amount of tokens is approximately close to 'train_mb_tokens' and 'logprob_mb_tokens' for the
  # training and logprob stages respectively.
  dynamic_batching:
    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}}
    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.dtensor_cfg.tensor_parallel_size}
  max_grad_norm: 1.0

  optimizer:
    name: "torch.optim.AdamW"
    kwargs:
      lr: 5.0e-6
      weight_decay: 0.01
      betas: [0.9, 0.999]
      eps: 1e-8
      # when using Dtensor, we need to set foreach
      # and fused to False
      foreach: False
      fused: False

  scheduler:
    - name: "torch.optim.lr_scheduler.LinearLR"
      kwargs:
        start_factor: 0.1
        end_factor: 1.0
        total_iters: 50
    - name: "torch.optim.lr_scheduler.ConstantLR"
      kwargs:
        factor: 1.0
        total_iters: 10000000000
    - milestones: [50]

  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:
      precision: ${policy.precision}
      tensor_parallel_size: 1
      gpu_memory_utilization: 0.6
      max_model_len: ${policy.max_total_sequence_length}

data:
  max_input_seq_length: ${policy.max_total_sequence_length} # upper bound, real truncation occurs at vllm.max_model_len
  prompt_file: "examples/prompts/cot.txt"
  system_prompt_file: null
  dataset_name: "OpenMathInstruct-2"

env:
  math:
    num_workers: 8

logger:
  log_dir: "logs"  # Base directory for all logs
  num_val_samples_to_print: 0 # Number of validation samples to pretty print on terminal
  wandb_enabled: false
  tensorboard_enabled: false
  monitor_gpus: false  # If true, will monitor GPU usage and log to wandb and/or tensorboard
  wandb:
    project: "grpo-dev"
    name: "grpo-dev-logger"
  tensorboard: {}
  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: 1
  num_nodes: 1
