rlhf_type: grpo

model: Qwen/Qwen3.5-2B
teacher_model: Qwen/Qwen3.5-9B

dataset: modelscope/gsm8k
dataset_num_proc: 4
split_dataset_ratio: 0

micro_batch_size: 2
global_batch_size: 16
num_generations: 1
steps_per_generation: 4
num_train_epochs: 1
logging_steps: 1
seed: 42
max_length: 2048
max_completion_length: 2048
padding_free: false
cross_entropy_loss_fusion: true
gradient_accumulation_fusion: false
lr: 3e-5
lr_warmup_fraction: 0.0
attention_backend: flash
temperature: 1.0
beta: 0
teacher_kl_coef: 1.0

use_vllm: true

colocate_groups: [[train, rollout]]
offload_model: true
offload_optimizer: true
offload_teacher_model: true
sleep_level: 1

save_steps: 100
no_save_optim: true
no_save_rng: true

train:
  gpus: 4
  tuner_type: lora
  lora_rank: 8
  lora_alpha: 32
  tensor_model_parallel_size: 1
  output_dir: megatron_output/ray_opd_rl_colocate

rollout:
  gpus: 4
  vllm_tensor_parallel_size: 1
  vllm_gpu_memory_utilization: 0.4
  vllm_max_model_len: 4096