reorgnize the workder initializer (#63)

Author: Kunlun-Zhu

train_ppo.sh

@@ -1,14 +1,16 @@
 #!/bin/bash
 
 # --- Configuration (defaults, can be overridden via env vars) ---
-export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-1,2,3,4}
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,5,9}
 WAND_PROJECT=${WAND_PROJECT:-'OpenManus-rl'}
 export BASE_MODEL=${BASE_MODEL:-'Qwen/Qwen2.5-3B'}
 AGENTGYM_HOST=${AGENTGYM_HOST:-'0.0.0.0'} # Default to 0.0.0.0 for external access
 AGENTGYM_SQL_BIRD_PATH=${AGENTGYM_SQL_BIRD_PATH:-} # Used only for sqlgym
 export NCCL_IB_DISABLE=1
 export NCCL_P2P_DISABLE=1
 export PYTHONPATH="./openmanus_rl/agentgym/agentenv:${PYTHONPATH}"
+export VLLM_ATTENTION_BACKEND=XFORMERS # vllm + qwen2-7b with flash_attn has some issues
+
 
 # --- Argument Parsing ---
 usage() {
@@ -228,7 +230,6 @@ export EXPERIMENT_NAME="OpenManus-rl-ppo-${BASE_MODEL##*/}-${AGENTGYM_ENV_NAME}$
 
 # --- Run PPO Training in Base Environment ---
 echo -e "\\n[Trainer] Running PPO training in base environment '$BASE_CONDA_ENV'..."
-export VLLM_ATTENTION_BACKEND=${VLLM_ATTENTION_BACKEND:-XFORMERS}
 
 # Construct server base URL, adding path if needed
 AGENTGYM_SERVER_BASE="http://$AGENTGYM_HOST" # Base URL without port
@@ -283,8 +284,8 @@ hydra_overrides=(
     "data.env_ports=[${AGENTGYM_PORTS_STR}]"
     "data.train_data_num=null"
     "data.val_data_num=null"
-    "data.train_batch_size=4"
-    "data.val_batch_size=2"
+    "data.train_batch_size=6"
+    "data.val_batch_size=3"
     "data.max_prompt_length=4096"
     "data.max_response_length=1000"
     "data.max_start_length=2048"
@@ -296,8 +297,8 @@ hydra_overrides=(
     "actor_rollout_ref.model.enable_gradient_checkpointing=true"
     "actor_rollout_ref.model.use_remove_padding=True"
     "actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.95"
-    "actor_rollout_ref.actor.ppo_mini_batch_size=4"
-    "actor_rollout_ref.actor.ppo_micro_batch_size=4"
+    "actor_rollout_ref.actor.ppo_mini_batch_size=6"
+    "actor_rollout_ref.actor.ppo_micro_batch_size=6"
     "actor_rollout_ref.actor.fsdp_config.param_offload=true"
     "actor_rollout_ref.actor.fsdp_config.grad_offload=true"
     "actor_rollout_ref.actor.fsdp_config.optimizer_offload=true"
@@ -329,7 +330,7 @@ hydra_overrides=(
     "+trainer.val_only=false"
     "+trainer.val_before_train=true"
     "trainer.default_hdfs_dir=null"
-    "trainer.n_gpus_per_node=4"
+    "trainer.n_gpus_per_node=3"
     "trainer.nnodes=1"
     "trainer.save_freq=100"
     "trainer.test_freq=50"