set fsdp2 as the defualt data parallel mode for VeOmniEngine

Author: A1waysBeenHere

.github/workflows/e2e_ppo_trainer_veomni_vllm.yml

@@ -134,7 +134,7 @@ jobs:
       - name: Running GEO3K E2E training tests on 8 L20 GPUs with veomni engine (FSDP_SIZE=8, USP=1)
         run: |
           ray stop --force
-          MODEL_ID=Qwen/Qwen3-VL-2B-Instruct TRAIN_FILES=${HOME}/data/geo3k/train.parquet VAL_FILES=${HOME}/data/gsm8k/test.parquet VAL_BEFORE_TRAIN=True NUM_GPUS=8 FSDP_SIZE=8 SP_SIZE=1 EP_SIZE=1 VERL_EXP_NAME="qwen3-2b-vl-function-reward-minimal-fsdp-size4" bash tests/special_e2e/run_ppo_trainer_veomni.sh
+          MODEL_ID=Qwen/Qwen3-VL-2B-Instruct TRAIN_FILES=${HOME}/data/geo3k/train.parquet VAL_FILES=${HOME}/data/gsm8k/test.parquet VAL_BEFORE_TRAIN=True NUM_GPUS=8 FSDP_SIZE=8 SP_SIZE=1 EP_SIZE=1 VERL_EXP_NAME="qwen3-2b-vl-function-reward-minimal-fsdp-size8" bash tests/special_e2e/run_ppo_trainer_veomni.sh
 
   cleanup:
     runs-on: ubuntu-latest

tests/special_e2e/run_ppo_trainer_veomni.sh

@@ -31,7 +31,6 @@ python3 -m verl.trainer.main_ppo --config-path=config\
     actor_rollout_ref.model.enable_gradient_checkpointing=True \
     actor_rollout_ref.actor.veomni.param_offload=True \
     actor_rollout_ref.actor.veomni.optimizer_offload=True \
-    actor_rollout_ref.actor.veomni.data_parallel_mode=fsdp2 \
     actor_rollout_ref.actor.ppo_mini_batch_size=8 \
     actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
     actor_rollout_ref.actor.use_kl_loss=True \

tests/special_e2e/sft/run_sft_engine.sh

@@ -67,7 +67,6 @@ VEOMNI_ENGINE_CONFIG="\
     optim.lr_min=1e-6 \
     optim.lr_scheduler_type=cosine \
     engine.ulysses_parallel_size=${SP_SIZE} \
-    engine.data_parallel_mode=${FSDP_STRATEGY} \
     engine.data_parallel_size=${FSDP_SIZE}"
 
 

verl/trainer/config/engine/veomni.yaml

@@ -7,9 +7,6 @@ param_offload: False
 # Whether to offload optimizer state to CPU
 optimizer_offload: False
 
-# fsdp or fsdp2
-data_parallel_mode: fsdp2
-
 data_parallel_size: 1
 
 data_parallel_replicate_size: 1

verl/trainer/config/reward_model/veomni_reward_model.yaml

@@ -15,9 +15,6 @@ veomni:
   # Target configuration dataclass
   _target_: verl.workers.config.VeOmniEngineConfig
 
-  # fsdp or fsdp2
-  data_parallel_mode: fsdp2
-
   data_parallel_size: 1
 
   data_parallel_replicate_size: 1

verl/workers/config/engine.py

@@ -202,7 +202,6 @@ class VeOmniEngineConfig(EngineConfig):
         pipeline_parallel_size (int): Pipeline parallel size, default 1
         context_parallel_size (int): Ring-attn context parallel size, default 1
         ulysses_parallel_size (int): Ulysses sequence parallel size, default 1
-        data_parallel_mode (str): Data parallel mode, default "fsdp"
         init_device (str): Device to initialize model weights.
             1. `cpu`: Init parameters on CPU in rank0 only.
             2. `cuda`: Init parameters on GPU.
@@ -259,7 +258,6 @@ class VeOmniEngineConfig(EngineConfig):
     pipeline_parallel_size: int = 1
     context_parallel_size: int = 1
     ulysses_parallel_size: int = 1
-    data_parallel_mode: Literal["ddp", "fsdp1", "fsdp2"] = "fsdp"
     seed: int = 42
     full_determinism: bool = False
     mixed_precision: bool = False

verl/workers/engine/veomni/transformer_impl.py

@@ -75,8 +75,8 @@ def __init__(
         self.engine_config = engine_config
         self.optimizer_config = optimizer_config
         self.checkpoint_config = checkpoint_config
-        assert self.engine_config.data_parallel_mode == "fsdp2", "VeOmniEngine only supports fsdp2."
-
+        # VeOmniEngine only supports fsdp2.
+        self.data_parallel_mode = "fsdp2"
         self.rank = dist.get_rank()
 
         parallel_state.init_parallel_state(
@@ -88,7 +88,7 @@ def __init__(
             pp_size=self.engine_config.pipeline_parallel_size,
             cp_size=self.engine_config.context_parallel_size,
             ulysses_size=self.engine_config.ulysses_parallel_size,
-            dp_mode=self.engine_config.data_parallel_mode,
+            dp_mode=self.data_parallel_mode,
         )
 
         if self.engine_config.full_determinism:
@@ -155,7 +155,7 @@ def _build_optimizer(self, module):
         )
         get_optimizer_pre_hook = getattr(module, "get_optimizer_pre_hook", None)
         if get_optimizer_pre_hook is not None:
-            optimizer_pre_hook = get_optimizer_pre_hook(module, module.config, self.engine_config.data_parallel_mode)
+            optimizer_pre_hook = get_optimizer_pre_hook(module, module.config, self.data_parallel_mode)
             optimizer.register_step_pre_hook(optimizer_pre_hook)
 
         return optimizer