| #!/usr/bin/env bash |
| set -xeuo pipefail |
|
|
|
|
| NUM_GPUS=${NUM_GPUS:-8} |
| ACTOR_STRATEGY=${ACTOR_STRATEGY:-"fsdp"} |
|
|
| |
| MODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct} |
| MODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}} |
| huggingface-cli download "${MODEL_ID}" --local-dir "${MODEL_PATH}" |
|
|
|
|
| rollout_mode="async" |
| rollout_name="vllm" |
| if [ "$rollout_mode" = "async" ]; then |
| export VLLM_USE_V1=1 |
| return_raw_chat="True" |
| fi |
|
|
| |
| adv_estimator=grpo |
|
|
| use_kl_in_reward=False |
| kl_coef=0.0 |
| use_kl_loss=False |
| kl_loss_coef=0.0 |
|
|
| clip_ratio_low=0.2 |
| clip_ratio_high=0.28 |
|
|
| |
| max_prompt_length=512 |
| max_response_length=1024 |
| enable_overlong_buffer=True |
| overlong_buffer_len=128 |
| overlong_penalty_factor=1.0 |
|
|
| |
| loss_agg_mode="token-mean" |
|
|
| |
| temperature=1.0 |
| top_p=1.0 |
| top_k=-1 |
| val_top_p=0.7 |
|
|
| n_gpus_training=8 |
| train_prompt_bsz=128 |
| val_prompt_bsz=128 |
| n_resp_per_prompt=5 |
| train_prompt_mini_bsz=32 |
| test_freq=-1 |
|
|
| log_dir="./logs" |
| mkdir -p $log_dir |
| timestamp=$(date +"%Y%m%d%H%M%S") |
| log_file="${log_dir}/qwen2_5-0_5b_transferqueue_${timestamp}.log" |
|
|
| exp_name="$(basename "${MODEL_ID,}")-transferqueue-${ACTOR_STRATEGY}-minimal" |
|
|
| echo "Running transferqueue with ${ACTOR_STRATEGY} strategy" |
| echo "Total GPUs: ${NUM_GPUS}" |
|
|
| |
| common_params=( |
| data.train_files="${HOME}/data/gsm8k/train.parquet" |
| data.val_files="${HOME}/data/gsm8k/test.parquet" |
| data.prompt_key=prompt |
| data.truncation='error' |
| data.max_prompt_length=${max_prompt_length} |
| data.max_response_length=${max_response_length} |
| data.filter_overlong_prompts_workers=128 |
| data.filter_overlong_prompts=True |
| data.train_batch_size=${train_prompt_bsz} |
| data.val_batch_size=${val_prompt_bsz} |
| data.return_raw_chat=${return_raw_chat} |
| actor_rollout_ref.rollout.n=${n_resp_per_prompt} |
| algorithm.adv_estimator=${adv_estimator} |
| algorithm.use_kl_in_reward=${use_kl_in_reward} |
| algorithm.kl_ctrl.kl_coef=${kl_coef} |
| actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} |
| actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} |
| actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} |
| actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} |
| actor_rollout_ref.actor.clip_ratio_c=10.0 |
| actor_rollout_ref.actor.use_kl_loss=True |
| actor_rollout_ref.model.path="${MODEL_PATH}" |
| actor_rollout_ref.actor.optim.lr=1e-6 |
| actor_rollout_ref.actor.optim.lr_warmup_steps=-1 |
| actor_rollout_ref.actor.optim.weight_decay=0.1 |
| actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} |
| actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 |
| actor_rollout_ref.actor.entropy_coeff=0 |
| actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} |
| actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 |
| actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 |
| actor_rollout_ref.rollout.gpu_memory_utilization=0.80 |
| actor_rollout_ref.rollout.temperature=${temperature} |
| actor_rollout_ref.rollout.top_p=${top_p} |
| actor_rollout_ref.rollout.top_k=${top_k} |
| actor_rollout_ref.rollout.max_num_batched_tokens=10240 |
| actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} |
| actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} |
| actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} |
| actor_rollout_ref.rollout.val_kwargs.do_sample=True |
| actor_rollout_ref.rollout.val_kwargs.n=1 |
| actor_rollout_ref.rollout.enable_chunked_prefill=True |
| actor_rollout_ref.rollout.name=${rollout_name} |
| actor_rollout_ref.rollout.mode=${rollout_mode} |
| actor_rollout_ref.rollout.disable_log_stats=True |
| trainer.logger=console |
| trainer.project_name='verl-test-transferqueue' |
| trainer.experiment_name="${exp_name}" |
| trainer.test_freq="${test_freq}" |
| trainer.save_freq=-1 |
| trainer.resume_mode=disable |
| trainer.nnodes=1 |
| trainer.n_gpus_per_node=${n_gpus_training} |
| trainer.total_training_steps=2 |
| trainer.total_epochs=15 |
| trainer.val_before_train=True |
| ) |
|
|
| if [ "${ACTOR_STRATEGY}" == "fsdp" ]; then |
| echo "Running TransferQueue training with FSDP strategy..." |
| |
| gen_tp=1 |
| sp_size=1 |
| fsdp_size=-1 |
| ref_offload=True |
| actor_offload=False |
|
|
| python3 -m recipe.transfer_queue.main_ppo \ |
| --config-path=config \ |
| --config-name='transfer_queue_ppo_trainer' \ |
| "${common_params[@]}" \ |
| actor_rollout_ref.model.enable_gradient_checkpointing=True \ |
| actor_rollout_ref.actor.strategy=fsdp \ |
| critic.strategy=fsdp \ |
| actor_rollout_ref.actor.grad_clip=1.0 \ |
| actor_rollout_ref.model.use_remove_padding=True \ |
| actor_rollout_ref.actor.use_dynamic_bsz=True \ |
| actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ |
| actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ |
| actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \ |
| actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \ |
| actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \ |
| actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \ |
| actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \ |
| actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \ |
| actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \ |
| 2>&1 | tee "$log_file" $@ |
|
|
| elif [ "${ACTOR_STRATEGY}" == "megatron" ]; then |
| echo "Running TransferQueue training with Megatron strategy..." |
| |
| gen_tp=2 |
| train_tp=1 |
| train_pp=2 |
| ref_offload=True |
| actor_offload=False |
|
|
| |
| |
| |
| python3 -m recipe.transfer_queue.main_ppo \ |
| --config-path=config \ |
| --config-name='transfer_queue_ppo_megatron_trainer' \ |
| "${common_params[@]}" \ |
| actor_rollout_ref.actor.strategy=megatron \ |
| critic.strategy=megatron \ |
| actor_rollout_ref.actor.optim.lr_decay_steps=10000000 \ |
| actor_rollout_ref.actor.megatron.param_offload=${actor_offload} \ |
| actor_rollout_ref.actor.megatron.optimizer_offload=${actor_offload} \ |
| actor_rollout_ref.actor.megatron.grad_offload=${actor_offload} \ |
| actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \ |
| actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \ |
| actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \ |
| actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \ |
| actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \ |
| actor_rollout_ref.ref.megatron.param_offload=${ref_offload} \ |
| 2>&1 | tee "$log_file" $@ |
| else |
| echo "Error: Unknown strategy ${ACTOR_STRATEGY}. Please use 'fsdp' or 'megatron'" |
| exit 1 |
| fi |
|
|
| echo "TransferQueue test completed successfully with ${ACTOR_STRATEGY} strategy" |
