feat(fsdp): support pipeline mode with FSDP as training backend (RLinf#332)

Signed-off-by: Bo Dai <daibo@infini-ai.com>

Author: anHappyDog

.github/workflows/agent-e2e-tests.yml

@@ -62,6 +62,20 @@ jobs:
           source .venv/bin/activate
           bash tests/e2e_tests/reasoning/run.sh qwen2.5-1.5b-grpo-collocated-fsdp-vllm
 
+      - name: FSDP SGLang Pipeline mode
+        timeout-minutes: 20
+        run: |
+          export REPO_PATH=$(pwd)
+          source .venv/bin/activate
+          bash tests/e2e_tests/reasoning/run.sh qwen2.5-1.5b-grpo-pipeline-fsdp-sgl
+
+      - name: FSDP vLLM Pipeline mode
+        timeout-minutes: 20
+        run: |
+          export REPO_PATH=$(pwd)
+          source .venv/bin/activate
+          bash tests/e2e_tests/reasoning/run.sh qwen2.5-1.5b-grpo-pipeline-fsdp-vllm
+
       - name: Clean up
         run: |
           rm -rf .venv

examples/coding_online_rl/main_coding_online_rl.py

@@ -57,7 +57,7 @@ def main(cfg) -> None:
         placement_strategy=singleton_placement_strategy,
     )
 
-    rollout_worker_cls = get_rollout_backend_worker(cfg, component_placement)
+    rollout_worker_cls = get_rollout_backend_worker(cfg)
 
     # Rollout group
     rollout_placement_strategy = component_placement.get_strategy("rollout")

examples/coding_online_rl/main_coding_rl_llm_judge.py

@@ -43,7 +43,7 @@ def main(cfg) -> None:
     cluster = Cluster(cluster_cfg=cfg.cluster)
     component_placement = ModelParallelComponentPlacement(cfg, cluster)
 
-    rollout_worker_cls = get_rollout_backend_worker(cfg, component_placement)
+    rollout_worker_cls = get_rollout_backend_worker(cfg)
 
     # Rollout group
     rollout_placement_strategy = component_placement.get_strategy("rollout")

examples/embodiment/config/behavior_openvlaoft_eval.yaml

@@ -139,6 +139,7 @@ actor:
     adam_beta1: 0.9
     adam_beta2: 0.999
     adam_eps: 1.0e-05
+    weight_decay: 0.01
     clip_grad: 10.0
 
   tokenizer:

examples/embodiment/config/behavior_ppo_openvlaoft.yaml

@@ -146,6 +146,7 @@ actor:
     adam_beta2: 0.999
     adam_eps: 1.0e-05
     clip_grad: 10.0
+    weight_decay: 0.01
     critic_warmup_steps: 0
 
   tokenizer:

examples/embodiment/config/isaaclab_ppo_gr00t_demo.yaml

@@ -152,6 +152,7 @@ actor:
     adam_beta2: 0.95
     adam_eps: 1.0e-05
     clip_grad: 1.0
+    weight_decay: 0.01
     critic_warmup_steps: 0
 
   # Override the default values in training_backend/fsdp

examples/embodiment/config/libero_10_grpo_openpi.yaml

@@ -149,6 +149,7 @@ actor:
     adam_beta1: 0.9
     adam_beta2: 0.95
     adam_eps: 1.0e-05
+    weight_decay: 0.01
     clip_grad: 2.0
 
   # Override the default values in training_backend/fsdp

examples/embodiment/config/libero_10_grpo_openpi_pi05.yaml

@@ -151,6 +151,7 @@ actor:
     adam_beta2: 0.95
     adam_eps: 1.0e-05
     clip_grad: 1.0
+    weight_decay: 0.01
 
   # Override the default values in training_backend/fsdp
   fsdp_config:

examples/embodiment/config/libero_10_grpo_openvlaoft.yaml

@@ -151,6 +151,7 @@ actor:
     adam_beta1: 0.9
     adam_beta2: 0.999
     adam_eps: 1.0e-05
+    weight_decay: 0.01
     clip_grad: 1.0
 
   tokenizer:

examples/embodiment/config/libero_10_grpo_openvlaoft_eval.yaml

@@ -145,6 +145,7 @@ actor:
     adam_beta1: 0.9
     adam_beta2: 0.999
     adam_eps: 1.0e-05
+    weight_decay: 0.01
     clip_grad: 1.0
 
   tokenizer: