diff --git a/evals/redis_stat.py b/evals/redis_stat.py index 3e08f53..423079a 100644 --- a/evals/redis_stat.py +++ b/evals/redis_stat.py @@ -129,4 +129,4 @@ def analyze_episodes_with_positions(tag): } # Run the analysis -results = analyze_episodes_with_positions("Qwen2.5-7B-Instruct_vs_sft_qwen25_7b_bigtom_step_1500-bigtom_0402") +results = analyze_episodes_with_positions("grpo_direct_step_400_vs_sft_qwen25_7b_sft_round_1_bc_data_top_2_step_1500-0420") diff --git a/scripts/accelerate_config_ppo.yaml b/scripts/accelerate_config_ppo.yaml index d4c8212..775922f 100644 --- a/scripts/accelerate_config_ppo.yaml +++ b/scripts/accelerate_config_ppo.yaml @@ -2,11 +2,6 @@ compute_environment: LOCAL_MACHINE debug: true distributed_type: MULTI_GPU downcast_bf16: 'no' -dynamo_config: - dynamo_backend: EAGER - dynamo_mode: default - dynamo_use_dynamic: false - dynamo_use_fullgraph: false enable_cpu_affinity: true gpu_ids: all machine_rank: 0 diff --git a/scripts/accelerate_config_rm.yaml b/scripts/accelerate_config_rm.yaml index fb193c9..5b42349 100644 --- a/scripts/accelerate_config_rm.yaml +++ b/scripts/accelerate_config_rm.yaml @@ -8,7 +8,7 @@ machine_rank: 0 main_training_function: main mixed_precision: bf16 num_machines: 1 -num_processes: 5 +num_processes: 8 rdzv_backend: static same_network: true tpu_env: [] diff --git a/scripts/accelerate_config_sft.yaml b/scripts/accelerate_config_sft.yaml index e517fed..5b42349 100644 --- a/scripts/accelerate_config_sft.yaml +++ b/scripts/accelerate_config_sft.yaml @@ -8,7 +8,7 @@ machine_rank: 0 main_training_function: main mixed_precision: bf16 num_machines: 1 -num_processes: 2 +num_processes: 8 rdzv_backend: static same_network: true tpu_env: [] diff --git a/scripts/inference_rm.sh b/scripts/inference_rm.sh index f000421..18f8846 100644 --- a/scripts/inference_rm.sh +++ b/scripts/inference_rm.sh @@ -8,8 +8,8 @@ CUDA_VISIBLE_DEVICES=9 python inference_rm.py \ --example_path "/data/haofeiy2/sotopia-rl/data/sotopia_pi_gpt4_rm_overfit.json" -CUDA_VISIBLE_DEVICES=8 python inference_rm.py \ +CUDA_VISIBLE_DEVICES=5 python inference_rm.py \ --model_path "/mnt/data_from_server1/models/Qwen2.5-7B-Instruct" \ - --adapter_path "/data/haofeiy2/sotopia-rl/rm_overfit_test/checkpoint-100" \ + --adapter_path "/data/haofeiy2/sotopia-rl/rm_token_length/checkpoint-800" \ --template_path "/data/haofeiy2/sotopia-rl/evals/qwen2.5-7b.jinja" \ --example_path "/data/haofeiy2/sotopia-rl/data/sotopia_pi_gpt4_rm_overfit.json" diff --git a/scripts/train_grpo.py b/scripts/train_grpo.py index 55a9181..5ef8913 100644 --- a/scripts/train_grpo.py +++ b/scripts/train_grpo.py @@ -31,6 +31,8 @@ help="Maximum length of generated responses") parser.add_argument("--num_generations", type=int, default=4, help="Number of generations for GRPO") + parser.add_argument("--beta", type=float, default=0.04, + help="KL coefficient for GRPO") # Adapter parameters parser.add_argument("--policy_adapter_path", type=str, default=None, diff --git a/scripts/train_grpo.sh b/scripts/train_grpo.sh index d6370d6..7691347 100644 --- a/scripts/train_grpo.sh +++ b/scripts/train_grpo.sh @@ -1,4 +1,4 @@ -CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 accelerate launch \ +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch \ --config_file /data/disk0/sotopia-rl/scripts/accelerate_config_grpo.yaml \ --main_process_port 29511 \ /data/disk0/sotopia-rl/scripts/train_grpo.py \ @@ -14,4 +14,9 @@ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 accelerate launch \ --num_grpo_epochs 2 \ --use_lora_train_grpo \ --num_generations 16 \ +<<<<<<< HEAD + --beta 0.04 \ + --output_dir /data/disk0/sotopia-rl/grpo_rm_reward_direct_default_beta_004 +======= --output_dir /data/disk0/sotopia-rl/grpo_rm_reward_direct_default +>>>>>>> 6a5e62f4de6e9b376b6ed3f56e6745b8bb0f86bf diff --git a/scripts/train_ppo.py b/scripts/train_ppo.py index 7af4b70..b769cbe 100644 --- a/scripts/train_ppo.py +++ b/scripts/train_ppo.py @@ -20,7 +20,7 @@ help="Number of PPO epochs per update") parser.add_argument("--learning_rate", type=float, default=5e-6, help="Learning rate for optimizer") - parser.add_argument("--gamma", type=float, default=0.99, + parser.add_argument("--gamma", type=float, default=1.0, help="Discount factor") parser.add_argument("--lam", type=float, default=0.95, help="GAE lambda for advantage estimation") @@ -32,7 +32,7 @@ help="Number of steps to accumulate gradients before performing an update") parser.add_argument("--val_ratio", type=float, default=0.05, help="Ratio of validation data") - parser.add_argument("--response_length", type=int, default=128, + parser.add_argument("--response_length", type=int, default=256, help="Maximum length of generated responses") parser.add_argument("--local_rollout_forward_batch_size", type=int, default=16, help="Batch size for local rollout forward pass") @@ -69,6 +69,5 @@ help="Use LoRA for training PPO") args = parser.parse_args() - accelerator = Accelerator() - trainer = SotopiaPPOTrainer(args, accelerator) + trainer = SotopiaPPOTrainer(args) trainer.train() diff --git a/scripts/train_ppo.sh b/scripts/train_ppo.sh index b9d8f5e..6c04e89 100644 --- a/scripts/train_ppo.sh +++ b/scripts/train_ppo.sh @@ -1,21 +1,213 @@ -CUDA_VISIBLE_DEVICES=1,7,8 accelerate launch \ - --config_file /data/haofeiy2/sotopia-rl/scripts/accelerate_config.yaml \ - --main_process_port 29511 \ - /data/haofeiy2/sotopia-rl/scripts/train_ppo.py \ - --model_name /mnt/data_from_server1/models/Qwen2.5-7B-Instruct \ - --policy_adapter_path /data/haofeiy2/sotopia-rl/sft_qwen25_7b_sft_round_1_bc_data_top_2/checkpoint-1500 \ - --ref_adapter_path /data/haofeiy2/sotopia-rl/sft_qwen25_7b_sft_round_1_bc_data_top_2/checkpoint-1500 \ - --reward_adapter_path /data/haofeiy2/sotopia-rl/rm_reward_direct_default_without_that_n_error_as_the_end/checkpoint-4480 \ - --value_adapter_path /data/haofeiy2/sotopia-rl/rm_reward_direct_default_without_that_n_error_as_the_end/checkpoint-4480 \ +# parameter I used for final PPO checkpoint +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch \ + --config_file /projects/bdpw/haofeiy/sotopia-rl/scripts/accelerate_config_ppo.yaml \ + --main_process_port 29439 \ + /projects/bdpw/haofeiy/sotopia-rl/scripts/train_ppo.py \ + --model_name /projects/bdpw/haofeiy/models/Qwen2.5-7B-Instruct \ + --policy_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_qwen25_7b_sft_round_1_bc_data_top_2/checkpoint-1500 \ + --ref_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_qwen25_7b_sft_round_1_bc_data_top_2/checkpoint-1500 \ + --reward_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_reward_direct_default_without_that_n_error_as_the_end/checkpoint-4480 \ + --value_adapter_path /projects/bdpw/haofeiy/sotopia-rl/ppo_top_2_sft_step1500_for_pretrained_value_model_direct_rm/checkpoint-90 \ + --learning_rate 3e-6 \ + --per_device_train_batch_size 3 \ + --per_device_eval_batch_size 4 \ + --gradient_accumulation_steps 4 \ + --num_mini_batches 1 \ + --ppo_data_path /projects/bdpw/haofeiy/sotopia-rl/data/sotopia_pi_round1_qwen_sft_all_with_instruct_string.json \ + --template_path /projects/bdpw/haofeiy/sotopia-rl/evals/qwen2.5-7b.jinja \ + --num_train_epochs 30 \ + --max_length 4096 \ + --num_ppo_epochs 2 \ + --gamma 0.99 \ + --use_lora_train_ppo \ + --output_dir /projects/bdpw/haofeiy/sotopia-rl/ppo_top_2_sft_step1500_with_pretrained_value_model_gamma_099_direct_rm + +# param for pretrain value adapter +CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch \ + --config_file /projects/bdpw/haofeiy/sotopia-rl/scripts/accelerate_config_ppo.yaml \ + --main_process_port 29439 \ + /projects/bdpw/haofeiy/sotopia-rl/scripts/train_ppo.py \ + --model_name /projects/bdpw/haofeiy/models/Qwen2.5-7B-Instruct \ + --policy_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_qwen25_7b_sft_round_1_bc_data_top_2/checkpoint-1500 \ + --ref_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_qwen25_7b_sft_round_1_bc_data_top_2/checkpoint-1500 \ + --reward_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_goal_w_conversation_behavior_4_23/checkpoint-9400 \ + --value_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_goal_w_conversation_behavior_4_23/checkpoint-9400 \ + --learning_rate 5e-5 \ + --per_device_train_batch_size 1 \ + --per_device_eval_batch_size 4 \ + --gradient_accumulation_steps 4 \ + --num_mini_batches 1 \ + --ppo_data_path /projects/bdpw/haofeiy/sotopia-rl/data/sotopia_pi_round1_qwen_sft_all_with_instruct_string.json \ + --template_path /projects/bdpw/haofeiy/sotopia-rl/evals/qwen2.5-7b.jinja \ + --num_train_epochs 30 \ + --max_length 4096 \ + --num_ppo_epochs 2 \ + --gamma 1.00 \ + --use_lora_train_ppo \ + --output_dir /projects/bdpw/haofeiy/sotopia-rl/ppo_top_2_sft_step1500_for_pretrained_value_model_direct_rm + + + +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch \ + --config_file /projects/bdpw/haofeiy/sotopia-rl/scripts/accelerate_config_ppo.yaml \ + --main_process_port 29439 \ + /projects/bdpw/haofeiy/sotopia-rl/scripts/train_ppo.py \ + --model_name /projects/bdpw/haofeiy/models/Qwen2.5-7B-Instruct \ + --policy_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_qwen25_7b_sft_round_1_bc_data_top_2/checkpoint-1500 \ + --ref_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_qwen25_7b_sft_round_1_bc_data_top_2/checkpoint-1500 \ + --reward_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_checkpoint-800 \ + --value_adapter_path /projects/bdpw/haofeiy/sotopia-rl/ppo_top_2_sft_1_epoch_step160_default_kl_token_length_pretrained_value_model \ --learning_rate 1e-5 \ --per_device_train_batch_size 1 \ - --per_device_eval_batch_size 1 \ - --gradient_accumulation_steps 1 \ + --per_device_eval_batch_size 4 \ + --gradient_accumulation_steps 2 \ + --num_mini_batches 1 \ + --ppo_data_path /projects/bdpw/haofeiy/sotopia-rl/data/sotopia_pi_round1_qwen_sft_all_with_instruct_string.json \ + --template_path /projects/bdpw/haofeiy/sotopia-rl/evals/qwen2.5-7b.jinja \ + --num_train_epochs 5 \ + --max_length 4096 \ + --num_ppo_epochs 2 \ + --gamma 1.00 \ + --use_lora_train_ppo \ + --output_dir /projects/bdpw/haofeiy/sotopia-rl/ppo_top_2_sft_1_epoch_step160_default_kl_token_length_with_pretrained_value_model + +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch \ + --config_file /projects/bdpw/haofeiy/sotopia-rl/scripts/accelerate_config_ppo.yaml \ + --main_process_port 29449 \ + /projects/bdpw/haofeiy/sotopia-rl/scripts/train_ppo.py \ + --model_name /projects/bdpw/haofeiy/models/Qwen2.5-7B-Instruct \ + --policy_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_round_1_bc_data_top_2_ckpt/checkpoint-30 \ + --ref_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_round_1_bc_data_top_2_ckpt/checkpoint-30 \ + --reward_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --value_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --learning_rate 1e-6 \ + --per_device_train_batch_size 1 \ + --per_device_eval_batch_size 4 \ + --gradient_accumulation_steps 4 \ --num_mini_batches 1 \ - --ppo_data_path /data/haofeiy2/sotopia-rl/data/sotopia_pi_round1_qwen_sft_all_with_instruct_string.json \ - --template_path /data/haofeiy2/sotopia-rl/evals/qwen2.5-7b.jinja \ + --ppo_data_path /projects/bdpw/haofeiy/sotopia-rl/data/sotopia_pi_round1_qwen_sft_all_with_instruct_string.json \ + --template_path /projects/bdpw/haofeiy/sotopia-rl/evals/qwen2.5-7b.jinja \ + --num_train_epochs 5 \ + --max_length 4096 \ --num_ppo_epochs 2 \ + --gamma 1.00 \ + --use_lora_train_ppo \ + --output_dir /projects/bdpw/haofeiy/sotopia-rl/grpo_top_2_sft_step30_default_kl + + +CUDA_VISIBLE_DEVICES=2,3 accelerate launch \ + --config_file /projects/bdpw/haofeiy/sotopia-rl/scripts/accelerate_config_ppo.yaml \ + --main_process_port 29469 \ + /projects/bdpw/haofeiy/sotopia-rl/scripts/train_ppo.py \ + --model_name /projects/bdpw/haofeiy/models/Qwen2.5-7B-Instruct \ + --policy_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_round_1_bc_data_top_2_ckpt/checkpoint-50 \ + --ref_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_round_1_bc_data_top_2_ckpt/checkpoint-50 \ + --reward_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --value_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --learning_rate 1e-6 \ + --per_device_train_batch_size 1 \ + --per_device_eval_batch_size 4 \ + --gradient_accumulation_steps 4 \ + --num_mini_batches 1 \ + --ppo_data_path /projects/bdpw/haofeiy/sotopia-rl/data/sotopia_pi_round1_qwen_sft_all_with_instruct_string.json \ + --template_path /projects/bdpw/haofeiy/sotopia-rl/evals/qwen2.5-7b.jinja \ --num_train_epochs 5 \ - --gamma 0.99 \ - --lam 0.95 \ - --output_dir /data/haofeiy2/sotopia-rl/ppo_origin_qwen25_7b_reward_direct_default_no_goal_gpt-4o_without_goal_leak_with_sft_self_play_data_use_sotopia_pi_full_data_0408 \ No newline at end of file + --max_length 4096 \ + --num_ppo_epochs 2 \ + --gamma 1.00 \ + --use_lora_train_ppo \ + --output_dir /projects/bdpw/haofeiy/sotopia-rl/grpo_top_2_sft_step50_default_kl + + +CUDA_VISIBLE_DEVICES=0,1 accelerate launch \ + --config_file /projects/bdpw/haofeiy/sotopia-rl/scripts/accelerate_config_ppo.yaml \ + --main_process_port 29499 \ + /projects/bdpw/haofeiy/sotopia-rl/scripts/train_ppo.py \ + --model_name /projects/bdpw/haofeiy/models/Qwen2.5-7B-Instruct \ + --policy_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_round_1_bc_data_top_2_ckpt/checkpoint-70 \ + --ref_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_round_1_bc_data_top_2_ckpt/checkpoint-70 \ + --reward_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --value_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --learning_rate 1e-6 \ + --per_device_train_batch_size 1 \ + --per_device_eval_batch_size 4 \ + --gradient_accumulation_steps 4 \ + --num_mini_batches 1 \ + --ppo_data_path /projects/bdpw/haofeiy/sotopia-rl/data/sotopia_pi_round1_qwen_sft_all_with_instruct_string.json \ + --template_path /projects/bdpw/haofeiy/sotopia-rl/evals/qwen2.5-7b.jinja \ + --num_train_epochs 5 \ + --max_length 4096 \ + --num_ppo_epochs 2 \ + --gamma 1.00 \ + --use_lora_train_ppo \ + --output_dir /projects/bdpw/haofeiy/sotopia-rl/grpo_top_2_sft_step70_default_kl + + +CUDA_VISIBLE_DEVICES=5 accelerate launch \ + --config_file /projects/bdpw/haofeiy/sotopia-rl/scripts/accelerate_config_ppo.yaml \ + --main_process_port 29549 \ + /projects/bdpw/haofeiy/sotopia-rl/scripts/train_ppo.py \ + --model_name /projects/bdpw/haofeiy/models/Qwen2.5-7B-Instruct \ + --policy_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_round_1_bc_data_top_2_ckpt/checkpoint-30 \ + --ref_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_round_1_bc_data_top_2_ckpt/checkpoint-30 \ + --reward_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --value_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --learning_rate 1e-6 \ + --per_device_train_batch_size 1 \ + --per_device_eval_batch_size 4 \ + --gradient_accumulation_steps 32 \ + --num_mini_batches 1 \ + --ppo_data_path /projects/bdpw/haofeiy/sotopia-rl/data/sotopia_pi_round1_qwen_sft_all_with_instruct_string.json \ + --template_path /projects/bdpw/haofeiy/sotopia-rl/evals/qwen2.5-7b.jinja \ + --num_train_epochs 5 \ + --max_length 4096 \ + --num_ppo_epochs 2 \ + --gamma 1.00 \ + --use_lora_train_ppo \ + --output_dir /projects/bdpw/haofeiy/sotopia-rl/ + +CUDA_VISIBLE_DEVICES=6 accelerate launch \ + --config_file /projects/bdpw/haofeiy/sotopia-rl/scripts/accelerate_config_ppo.yaml \ + --main_process_port 29559 \ + /projects/bdpw/haofeiy/sotopia-rl/scripts/train_ppo.py \ + --model_name /projects/bdpw/haofeiy/models/Qwen2.5-7B-Instruct \ + --policy_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_qwen25_7b_bc/checkpoint-500 \ + --ref_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_qwen25_7b_bc/checkpoint-500 \ + --reward_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --value_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --learning_rate 1e-6 \ + --per_device_train_batch_size 1 \ + --per_device_eval_batch_size 4 \ + --gradient_accumulation_steps 32 \ + --num_mini_batches 1 \ + --ppo_data_path /projects/bdpw/haofeiy/sotopia-rl/data/sotopia_pi_round1_qwen_sft_all_with_instruct_string.json \ + --template_path /projects/bdpw/haofeiy/sotopia-rl/evals/qwen2.5-7b.jinja \ + --num_train_epochs 5 \ + --max_length 4096 \ + --num_ppo_epochs 2 \ + --gamma 1.00 \ + --use_lora_train_ppo \ + --output_dir /projects/bdpw/haofeiy/sotopia-rl/ppo_token_length_normalized_with_sft_testing_ckpt500 + +CUDA_VISIBLE_DEVICES=7 accelerate launch \ + --config_file /projects/bdpw/haofeiy/sotopia-rl/scripts/accelerate_config_ppo.yaml \ + --main_process_port 29569 \ + /projects/bdpw/haofeiy/sotopia-rl/scripts/train_ppo.py \ + --model_name /projects/bdpw/haofeiy/models/Qwen2.5-7B-Instruct \ + --policy_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_qwen25_7b_bc/checkpoint-700 \ + --ref_adapter_path /projects/bdpw/haofeiy/sotopia-rl/sft_qwen25_7b_bc/checkpoint-700 \ + --reward_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --value_adapter_path /projects/bdpw/haofeiy/sotopia-rl/rm_token_length_normalized/checkpoint-500 \ + --learning_rate 1e-6 \ + --per_device_train_batch_size 1 \ + --per_device_eval_batch_size 4 \ + --gradient_accumulation_steps 32 \ + --num_mini_batches 1 \ + --ppo_data_path /projects/bdpw/haofeiy/sotopia-rl/data/sotopia_pi_round1_qwen_sft_all_with_instruct_string.json \ + --template_path /projects/bdpw/haofeiy/sotopia-rl/evals/qwen2.5-7b.jinja \ + --num_train_epochs 5 \ + --max_length 4096 \ + --num_ppo_epochs 2 \ + --gamma 1.00 \ + --use_lora_train_ppo \ + --output_dir /projects/bdpw/haofeiy/sotopia-rl/ppo_token_length_normalized_with_sft_testing_ckpt700 diff --git a/scripts/train_rm.sh b/scripts/train_rm.sh index d934b76..1a8ae68 100644 --- a/scripts/train_rm.sh +++ b/scripts/train_rm.sh @@ -1,18 +1,18 @@ -CUDA_VISIBLE_DEVICES=5,6,7,8,9 accelerate launch \ - --config_file /data/haofeiy2/sotopia-rl/scripts/accelerate_config_rm.yaml \ +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch \ + --config_file /data/disk0/sotopia-rl/scripts/accelerate_config_rm.yaml \ --main_process_port 29500 \ - /data/haofeiy2/sotopia-rl/scripts/train_rm.py \ - --model_name /mnt/data_from_server1/models/Qwen2.5-7B-Instruct \ + /data/disk0/sotopia-rl/scripts/train_rm.py \ + --model_name /data/disk0/models/Qwen2.5-7B-Instruct \ --learning_rate 1e-5 \ --max_length 4096 \ - --train_batch_size 1 \ + --train_batch_size 4 \ --val_batch_size 1 \ - --accumulation_steps 8 \ + --accumulation_steps 2 \ --num_epochs 30 \ --evaluation_steps 50 \ - --reward_data_path /data/haofeiy2/sotopia-rl/data/sotopia_pi_bc_episodes_reward_token_length.json \ - --template_path /data/haofeiy2/sotopia-rl/evals/qwen2.5-7b.jinja \ - --checkpoint_dir /data/haofeiy2/sotopia-rl/rm_token_length + --reward_data_path /data/disk0/sotopia-rl/data/sotopia_pi_bc_episodes_reward_token_length_binary.json \ + --template_path /data/disk0/sotopia-rl/evals/qwen2.5-7b.jinja \ + --checkpoint_dir /data/disk0/sotopia-rl/rm_token_length_binary CUDA_VISIBLE_DEVICES=5,6,7,8,9 accelerate launch \ --config_file /data/haofeiy2/sotopia-rl/scripts/accelerate_config_rm.yaml \ diff --git a/scripts/train_sft.sh b/scripts/train_sft.sh index 4836d6a..9e767d9 100644 --- a/scripts/train_sft.sh +++ b/scripts/train_sft.sh @@ -1,8 +1,27 @@ -CUDA_VISIBLE_DEVICES=7,8 accelerate launch \ - --config_file /data/haofeiy2/sotopia-rl/scripts/accelerate_config_sft.yaml \ - --main_process_port 29512 \ - /data/haofeiy2/sotopia-rl/scripts/train_sft.py \ - --model_name /mnt/data_from_server1/models/Qwen2.5-7B-Instruct \ +# parameter I used for final SFT checkpoint, I use ckpt 160 +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch \ + --config_file /data/disk0/sotopia-rl/scripts/accelerate_config_sft.yaml \ + --main_process_port 29112 \ + /data/disk0/sotopia-rl/scripts/train_sft.py \ + --model_name /data/disk0/models/Qwen2.5-7B-Instruct \ + --learning_rate 1e-5 \ + --max_length 4096 \ + --train_batch_size 1 \ + --val_batch_size 4 \ + --accumulation_steps 1 \ + --num_epochs 1 \ + --use_lora \ + --evaluation_steps 5 \ + --sft_data_path /data/disk0/sotopia-rl/data/sft_round_1_bc_data_top_2.json \ + --template_path /data/disk0/sotopia-rl/evals/qwen2.5-7b.jinja \ + --checkpoint_dir /data/disk0/sotopia-rl/sft_round_1_bc_data_top_2_only_1_epoch + + +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch \ + --config_file /data/disk0/sotopia-rl/scripts/accelerate_config_sft.yaml \ + --main_process_port 29112 \ + /data/disk0/sotopia-rl/scripts/train_sft.py \ + --model_name /data/disk0/models/Qwen2.5-7B-Instruct \ --learning_rate 1e-4 \ --max_length 4096 \ --train_batch_size 1 \ @@ -10,9 +29,7 @@ CUDA_VISIBLE_DEVICES=7,8 accelerate launch \ --accumulation_steps 6 \ --num_epochs 20 \ --use_lora \ - --evaluation_steps 100 \ - --sft_data_path /data/haofeiy2/sotopia-rl/data/sotopia_pi_round1_qwen_sft_pi_with_instruct_string.json \ - --template_path /data/haofeiy2/sotopia-rl/evals/qwen2.5-7b.jinja \ - --lora_checkpoint_path /data/haofeiy2/sotopia-rl/sft_qwen25_7b_sft_round_1_bc_data_top_2/checkpoint-1500 \ - --checkpoint_dir /data/haofeiy2/sotopia-rl/sft_qwen25_7b_pi_round1_qwen_sft_pi \ - --use_qlora \ No newline at end of file + --evaluation_steps 10 \ + --sft_data_path /data/disk0/sotopia-rl/data/sft_round_1_bc_data_top_2.json \ + --template_path /data/disk0/sotopia-rl/evals/qwen2.5-7b.jinja \ + --checkpoint_dir /data/disk0/sotopia-rl/sft_round_1_bc_data_top_2_ckpt \ No newline at end of file diff --git a/sotopia_rl/data.py b/sotopia_rl/data.py index b7fff53..e440149 100644 --- a/sotopia_rl/data.py +++ b/sotopia_rl/data.py @@ -198,4 +198,4 @@ def __getitem__(self, idx: int) -> Dict[str, Any]: return { "prompt": rendered_prompt, "completion": item["output"] - } + } \ No newline at end of file diff --git a/sotopia_rl/grpo_trainer_math.py b/sotopia_rl/grpo_trainer_math.py new file mode 100644 index 0000000..d6f3a62 --- /dev/null +++ b/sotopia_rl/grpo_trainer_math.py @@ -0,0 +1,164 @@ +import os +import torch +import wandb +from datasets import load_dataset +from torch.utils.data import random_split +from transformers import ( + AutoModelForCausalLM, + AutoModelForSequenceClassification, + AutoTokenizer, + BitsAndBytesConfig, + GenerationConfig, +) +from accelerate import PartialState +from peft import PeftModelForCausalLM, PeftModelForSequenceClassification +from jinja2 import Environment, FileSystemLoader +from trl import get_kbit_device_map, GRPOConfig, GRPOTrainer +from accelerate import Accelerator +from sotopia_rl.data import GRPODataset +from functools import partial +from typing import List + +os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" +os.environ['NCCL_P2P_DISABLE'] = '1' +os.environ['TORCH_DISTRIBUTED_DEBUG'] = 'DETAIL' + +SIMPLE_CHAT_TEMPLATE = "{% for message in messages %}{{message['role'].capitalize() + ': ' + message['content'] + '\n\n'}}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}" + + +from transformers import GPTNeoXForCausalLM + +class PatchedGPTNeoXForCausalLM(GPTNeoXForCausalLM): + def forward(self, *args, logits_to_keep=None, **kwargs): + return super().forward(*args, **kwargs) + +class SotopiaGRPOTrainer: + def __init__(self, args, accelerator: Accelerator): + self.args = args + self.accelerator = accelerator + + self._init_wandb() + self._setup_tokenizer() + self._setup_dataset() + self._create_quantization_config() + + self._setup_grpo_trainer() + + def save_model(self, output_dir: str, _internal_call: bool = False): + self.model.save_pretrained(output_dir) + self.tokenizer.save_pretrained(output_dir) + print(f"Saved PEFT model to {output_dir}") + + self.grpo_trainer.save_model = save_model.__get__(self.grpo_trainer, type(self.grpo_trainer)) + + def _init_wandb(self): + wandb.init( + project=self.args.wandb_project, + name=self.args.wandb_run_name, + config={k: v for k, v in vars(self.args).items() if isinstance(v, (int, float, str))} + ) + + def _setup_tokenizer(self): + self.tokenizer = AutoTokenizer.from_pretrained("/data/disk0/models/EleutherAI_pythia-1b-deduped__sft__tldr") + self.tokenizer.add_special_tokens({'pad_token': '[PAD]'}) + self.tokenizer.pad_token_id = self.tokenizer.convert_tokens_to_ids('[PAD]') + if self.tokenizer.chat_template is None: + self.tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE + + + def _setup_dataset(self): + from datasets import load_dataset + + dataset = load_dataset("trl-internal-testing/tldr-preference-sft-trl-style") + print("processing") + train_dataset = dataset["train"] + eval_dataset = dataset["validation"] + + def prepare_dataset(dataset, tokenizer): + def tokenize(element): + input_ids = tokenizer.apply_chat_template( + element["messages"][:1], + padding=False, + add_generation_prompt=True, + ) + return {"input_ids": input_ids, "lengths": len(input_ids), "prompt": element["messages"][:1]} + + return dataset.map( + tokenize, + remove_columns=dataset.column_names, + num_proc=4, + ) + + with PartialState().local_main_process_first(): + train_dataset = prepare_dataset(train_dataset, self.tokenizer) + if eval_dataset is not None: + eval_dataset = prepare_dataset(eval_dataset, self.tokenizer) + train_dataset = train_dataset.filter(lambda x: x["lengths"] <= 512, num_proc=4) + if eval_dataset is not None: + eval_dataset = eval_dataset.filter(lambda x: x["lengths"] <= 512, num_proc=4) + + assert train_dataset[0]["input_ids"][-1] != self.tokenizer.eos_token_id, "The last token should not be an EOS token" + + self.train_dataset = train_dataset + self.val_dataset = eval_dataset + print(f"Dataset loaded and processed: {len(self.train_dataset)} train, {len(self.val_dataset or [])} validation") + + def _create_quantization_config(self): + self.quant_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=torch.bfloat16, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4" + ) + + def _setup_grpo_trainer(self): + num_processes = self.accelerator.num_processes + global_batch_size = self.args.per_device_train_batch_size * num_processes + + num_generations = 4 # manually chosen value + print(f"Using num_generations = {num_generations} (global_batch_size = {global_batch_size})") + + policy_model = AutoModelForCausalLM.from_pretrained( + "/data/disk0/models/EleutherAI_pythia-1b-deduped__sft__tldr", + torch_dtype='auto', + num_labels=1, + ) + + reward_model = AutoModelForSequenceClassification.from_pretrained( + "/data/disk0/models/EleutherAI_pythia-1b-deduped__reward__tldr", + torch_dtype='auto', + num_labels=1, + ) + + training_args = GRPOConfig( + logging_steps = 1, + report_to = "wandb", + per_device_train_batch_size=self.args.per_device_train_batch_size, + per_device_eval_batch_size=self.args.per_device_eval_batch_size, + gradient_accumulation_steps=self.args.gradient_accumulation_steps, + num_train_epochs=self.args.num_train_epochs, + learning_rate=self.args.learning_rate, + output_dir=self.args.output_dir, + save_steps=self.args.save_steps, + num_generations=num_generations + ) + + self.grpo_trainer = GRPOTrainer( + args=training_args, + model=policy_model, + reward_funcs=reward_model, + processing_class=self.tokenizer, + reward_processing_classes=self.tokenizer, + train_dataset=self.train_dataset, + eval_dataset=self.val_dataset, + ) + print("GRPOtrainer setup complete") + + def train(self): + try: + print("Starting GRPO training...") + train_stats = self.grpo_trainer.train() + return train_stats + except Exception as e: + print(f"Training error: {str(e)}") + raise diff --git a/sotopia_rl/ppo_trainer.py b/sotopia_rl/ppo_trainer.py index bde99fc..dd4a896 100644 --- a/sotopia_rl/ppo_trainer.py +++ b/sotopia_rl/ppo_trainer.py @@ -12,7 +12,10 @@ GenerationConfig, ) from trl import get_kbit_device_map, PPOConfig, PPOTrainer +from peft import prepare_model_for_kbit_training +from trl.trainer.utils import disable_dropout_in_model from accelerate import PartialState, Accelerator +import copy import wandb from sotopia_rl.data import PPODataset @@ -20,8 +23,7 @@ os.environ["TORCH_DISTRIBUTED_DEBUG"] = "DETAIL" class SotopiaPPOTrainer: - def __init__(self, args, accelerator: Accelerator): - self.accelerator = accelerator + def __init__(self, args): self.args = args self._init_wandb() @@ -31,25 +33,24 @@ def __init__(self, args, accelerator: Accelerator): self._setup_generation_models() self._setup_classification_models() - - self.policy, self.ref_policy, self.reward_model, self.value_model = self.accelerator.prepare( - self.policy, self.ref_policy, self.reward_model, self.value_model - ) - self.policy = self.accelerator.unwrap_model(self.policy) - self.ref_policy = self.accelerator.unwrap_model(self.ref_policy) - self.reward_model = self.accelerator.unwrap_model(self.reward_model) - self.value_model = self.accelerator.unwrap_model(self.value_model) - self._setup_ppo_trainer() + for m in [self.policy]: + m.config.use_cache = False + for m in [self.value_model, self.reward_model]: + m.config.use_cache = False + def save_model(self, output_dir: str, _internal_call: bool = False): if hasattr(self.model, "policy"): - model_to_save = self.model.policy - elif hasattr(self.model, "module") and hasattr(self.model.module, "policy"): - model_to_save = self.model.module.policy + policy_to_save = self.model.policy + value_to_save = self.model.value_model + elif hasattr(self.model, "module"): + policy_to_save = self.model.module.policy + value_to_save = self.model.module.value_model else: - model_to_save = self.model - model_to_save.save_pretrained(output_dir) + raise ValueError("Model does not have 'policy' or 'module' attribute") + policy_to_save.save_pretrained(output_dir) + value_to_save.save_pretrained(output_dir) self.tokenizer.save_pretrained(output_dir) print(f"Model saved to {output_dir}") @@ -75,16 +76,16 @@ def _setup_dataset(self): # Create and split dataset dataset = PPODataset( - self.args.ppo_data_path, - self.tokenizer, - template, + data_path=self.args.ppo_data_path, + tokenizer=self.tokenizer, + template=template, max_length=self.args.max_length ) print(f"dataset: {len(dataset)}") generator = torch.Generator().manual_seed(42) val_ratio = getattr(self.args, 'val_ratio', 0.05) - train_size = min(int(len(dataset) * (1 - val_ratio)), len(dataset) - 2) + train_size = min(int(len(dataset) * (1 - val_ratio)), len(dataset) - 10) val_size = len(dataset) - train_size self.train_dataset, self.val_dataset = random_split(dataset, [train_size, val_size], generator=generator) print(f"Dataset split: {len(self.train_dataset)} train, {len(self.val_dataset)} validation") @@ -98,28 +99,19 @@ def _create_quantization_config(self): ) def _setup_generation_models(self): - base_gen_ref = AutoModelForCausalLM.from_pretrained( + self.ref_policy = AutoModelForCausalLM.from_pretrained( self.args.model_name, torch_dtype='auto', - quantization_config=self.quant_config, - device_map=get_kbit_device_map(), - ) - self.ref_policy = PeftModelForCausalLM.from_pretrained( - base_gen_ref, - self.args.ref_adapter_path, - is_trainable=False, - adapter_name="ref_adapter" ) + if self.args.use_lora_train_ppo: - base_gen_policy = AutoModelForCausalLM.from_pretrained( + self.base_gen_policy = AutoModelForCausalLM.from_pretrained( self.args.model_name, torch_dtype='auto', - quantization_config=self.quant_config, - device_map=get_kbit_device_map(), ) self.policy = PeftModelForCausalLM.from_pretrained( - base_gen_policy, + self.base_gen_policy, self.args.policy_adapter_path, is_trainable=True, adapter_name="policy_adapter" @@ -129,6 +121,8 @@ def _setup_generation_models(self): self.args.model_name, torch_dtype='auto', ) + self.ref_policy.config.pad_token_id = self.tokenizer.pad_token_id + self.policy.config.pad_token_id = self.tokenizer.pad_token_id requires_grad_num = 0 for name, param in self.policy.named_parameters(): @@ -137,11 +131,15 @@ def _setup_generation_models(self): requires_grad_num += 1 print(f"Number of trainable parameters in policy: {requires_grad_num}") - requires_grad_num = 0 - for name, param in self.ref_policy.named_parameters(): - if param.requires_grad: - requires_grad_num += 1 - print(f"Number of trainable parameters in ref policy: {requires_grad_num}") + + #for name, param in self.policy.named_parameters(): + # if self.policy.active_adapter in name: + # param.requires_grad = False + #requires_grad_num = 0 + #for name, param in self.ref_policy.named_parameters(): + # if param.requires_grad: + # requires_grad_num += 1 + #print(f"Number of trainable parameters in ref policy: {requires_grad_num}") def _setup_classification_models(self): base_reward_model = AutoModelForSequenceClassification.from_pretrained( @@ -157,6 +155,7 @@ def _setup_classification_models(self): is_trainable=False, adapter_name="reward_adapter" ) + for name, param in self.reward_model.named_parameters(): if self.reward_model.active_adapter in name: param.requires_grad = False @@ -166,8 +165,6 @@ def _setup_classification_models(self): self.args.model_name, torch_dtype='auto', num_labels=1, - quantization_config=self.quant_config, - device_map=get_kbit_device_map(), ) self.value_model = PeftModelForSequenceClassification.from_pretrained( base_value_model, @@ -182,12 +179,9 @@ def _setup_classification_models(self): num_labels=1, ) - # VERY VERY IMPORTANT - # specifically designed for PPO training, - # based on the get_reward function - # it fill the input_ids paddings with 0s - self.value_model.config.pad_token_id = 0 - self.reward_model.config.pad_token_id = 0 + # need to set this with not None results + self.value_model.config.pad_token_id = self.tokenizer.pad_token_id + self.reward_model.config.pad_token_id = self.tokenizer.pad_token_id requires_grad_num = 0 for name, param in self.value_model.named_parameters(): @@ -219,13 +213,16 @@ def _setup_ppo_trainer(self): ddp_find_unused_parameters=True, response_length=self.args.response_length, stop_token='eos', + kl_estimator='k3', + vf_coef=1e-3, + kl_coef=0.05, ) self.ppo_trainer = PPOTrainer( args=training_args, model=self.policy, + ref_model=copy.deepcopy(self.policy), processing_class=self.tokenizer, - ref_model=self.ref_policy, reward_model=self.reward_model, value_model=self.value_model, train_dataset=self.train_dataset,