pre dapo

Author: hiyouga

examples/config.yaml

@@ -7,7 +7,8 @@ data:
   image_dir: null
   max_prompt_length: 2048
   max_response_length: 2048
-  rollout_batch_size: 512
+  rollout_batch_size: 512  # equivalent to verl's data.train_batch_size
+  mini_rollout_batch_size: null  # equivalent to verl's data.gen_batch_size
   val_batch_size: 1024
   format_prompt: ./examples/format_prompt/math_format.jinja
   override_chat_template: null
@@ -26,9 +27,9 @@ algorithm:
 
 worker:
   actor:
-    global_batch_size: 128
-    micro_batch_size_per_device_for_update: 4
-    micro_batch_size_per_device_for_experience: 16
+    global_batch_size: 128  # equivalent to verl's actor.ppo_mini_batch_size
+    micro_batch_size_per_device_for_update: 4  # equivalent to verl's actor.ppo_micro_batch_size_per_gpu
+    micro_batch_size_per_device_for_experience: 16  # equivalent to verl's rollout.log_prob_micro_batch_size_per_gpu
     max_grad_norm: 1.0
     padding_free: true
     ulysses_size: 1

verl/trainer/config.py

@@ -42,6 +42,7 @@ class DataConfig:
     max_prompt_length: int = 512
     max_response_length: int = 512
     rollout_batch_size: int = 512
+    mini_rollout_batch_size: Optional[int] = None
     val_batch_size: int = -1
     format_prompt: Optional[str] = None
     override_chat_template: Optional[str] = None

verl/trainer/data_loader.py

@@ -47,9 +47,14 @@ def create_dataloader(config: DataConfig, tokenizer: PreTrainedTokenizer, proces
     else:
         sampler = SequentialSampler(data_source=train_dataset)
 
+    if config.mini_rollout_batch_size is not None:
+        train_batch_size = config.mini_rollout_batch_size
+    else:
+        train_batch_size = config.rollout_batch_size
+
     train_dataloader = StatefulDataLoader(
         dataset=train_dataset,
-        batch_size=config.rollout_batch_size,
+        batch_size=train_batch_size,
         sampler=sampler,
         num_workers=8,
         collate_fn=collate_fn,
@@ -72,9 +77,15 @@ def create_dataloader(config: DataConfig, tokenizer: PreTrainedTokenizer, proces
         max_pixels=config.max_pixels,
         filter_overlong_prompts=config.filter_overlong_prompts,
     )
+
+    if config.val_batch_size == -1:
+        val_batch_size = len(val_dataset)
+    else:
+        val_batch_size = config.val_batch_size
+
     val_dataloader = StatefulDataLoader(
         dataset=val_dataset,
-        batch_size=len(val_dataset) if config.val_batch_size == -1 else config.val_batch_size,
+        batch_size=val_batch_size,
         shuffle=False,
         num_workers=8,
         collate_fn=collate_fn,

verl/trainer/ray_trainer.py

@@ -192,7 +192,7 @@ def compute_log_prob(self, data: DataProto) -> torch.Tensor:
         )
         log_probs_lst = []
         if self.rank == 0:
-            micro_batches = tqdm(micro_batches, desc="Compute log probs", position=2)
+            micro_batches = tqdm(micro_batches, desc="Compute log probs", position=1)
 
         for micro_batch in micro_batches:
             model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch}
@@ -217,15 +217,15 @@ def update_policy(self, data: DataProto) -> Dict[str, Any]:
         metrics = defaultdict(list)
         for _ in range(self.config.ppo_epochs):
             if self.rank == 0:
-                mini_batches = tqdm(mini_batches, desc="Train mini-batches", position=2)
+                mini_batches = tqdm(mini_batches, desc="Train mini-batches", position=1)
 
             for mini_batch in mini_batches:
                 gradient_accumulation = (
                     self.config.global_batch_size_per_device // self.config.micro_batch_size_per_device_for_update
                 )
                 micro_batches = mini_batch.split(self.config.micro_batch_size_per_device_for_update)
                 if self.rank == 0:
-                    micro_batches = tqdm(micro_batches, desc="Update policy", position=3)
+                    micro_batches = tqdm(micro_batches, desc="Update policy", position=2)
 
                 for micro_batch in micro_batches:
                     model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch}

verl/workers/actor/dp_actor.py

@@ -149,7 +149,7 @@ def compute_values(self, data: DataProto) -> torch.Tensor:
         )
         values_lst = []
         if self.rank == 0:
-            micro_batches = tqdm(micro_batches, desc="Compute values", position=2)
+            micro_batches = tqdm(micro_batches, desc="Compute values", position=1)
 
         for micro_batch in micro_batches:
             model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch}
@@ -176,15 +176,15 @@ def update_critic(self, data: DataProto) -> Dict[str, Any]:
         metrics = defaultdict(list)
         for _ in range(self.config.ppo_epochs):
             if self.rank == 0:
-                mini_batches = tqdm(mini_batches, desc="Train mini-batches", position=2)
+                mini_batches = tqdm(mini_batches, desc="Train mini-batches", position=1)
 
             for mini_batch in mini_batches:
                 gradient_accumulation = (
                     self.config.global_batch_size_per_device // self.config.micro_batch_size_per_device_for_update
                 )
                 micro_batches = mini_batch.split(self.config.micro_batch_size_per_device_for_update)
                 if self.rank == 0:
-                    micro_batches = tqdm(micro_batches, desc="Update critic", position=3)
+                    micro_batches = tqdm(micro_batches, desc="Update critic", position=2)
 
                 for micro_batch in micro_batches:
                     model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch}

verl/workers/critic/dp_critic.py

@@ -339,6 +339,7 @@ def _build_rollout(self) -> None:
             module=self.fsdp_module,
             inference_engine=self.rollout.inference_engine,
             device_mesh=rollout_device_mesh,
+            use_param_offload=self._use_param_offload,
         )
         print_gpu_memory_usage("After vllm init")
 
@@ -518,9 +519,6 @@ def update_actor(self, data: DataProto):
     def generate_sequences(self, prompts: DataProto):
         assert self._has_rollout
 
-        if self._use_param_offload:
-            load_fsdp_model(self.fsdp_module)
-
         meta_info = {
             "eos_token_id": self.generation_config.eos_token_id
             if self.generation_config is not None
@@ -530,14 +528,8 @@ def generate_sequences(self, prompts: DataProto):
             else self.tokenizer.pad_token_id,
         }
         prompts.meta_info.update(meta_info)
+        self.rollout_sharding_manager.skip_vllm_sync_once = prompts.meta_info.get("skip_vllm_sync_once", False)
         with self.rollout_sharding_manager:
-            # after parameters sync with rollout, offload actor model to CPU
-            if self._use_param_offload:
-                offload_fsdp_model(self.fsdp_module)
-
-            if self._use_optimizer_offload:
-                offload_fsdp_optimizer(optimizer=self.optimizer)
-
             prompts = self.rollout_sharding_manager.preprocess_data(prompts)
             output = self.rollout.generate_sequences(prompts=prompts)
             output = self.rollout_sharding_manager.postprocess_data(output)

verl/workers/fsdp_workers.py

@@ -27,7 +27,8 @@
 from vllm.distributed import parallel_state as vllm_ps
 
 from ...protocol import DataProto, all_gather_data_proto
-from ...utils.model_utils import print_gpu_memory_usage
+from ...utils.fsdp_utils import load_fsdp_model, offload_fsdp_model
+from ...utils.model_utils import is_rank0, print_gpu_memory_usage
 from .base import BaseShardingManager
 
 
@@ -37,10 +38,13 @@ def __init__(
         module: FSDP,
         inference_engine: LLM,
         device_mesh: DeviceMesh,
+        use_param_offload: bool,
     ):
         self.module = module
         self.inference_engine = inference_engine
         self.device_mesh = device_mesh
+        self.use_param_offload = use_param_offload
+        self.skip_vllm_sync_once = False
 
         self.world_size = dist.get_world_size()
         self.tp_size = vllm_ps.get_tensor_model_parallel_world_size()
@@ -85,6 +89,24 @@ def _make_weight_iterator(
         for name, tensor in actor_weights.items():
             yield name, tensor.full_tensor() if self.world_size != 1 else tensor
 
+    def _sync_weight_to_vllm(self):
+        if self.use_param_offload:
+            load_fsdp_model(self.module)
+
+        actor_weights = get_model_state_dict(self.module)
+        actor_weights = self._rename_weight_keys(actor_weights, self.module._fsdp_wrapped_module)
+        print_gpu_memory_usage("After gather model weights in sharding manager")
+
+        model = self.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner.model
+        model.load_weights(self._make_weight_iterator(actor_weights))
+
+        del actor_weights
+        if self.use_param_offload:
+            offload_fsdp_model(self.module)
+
+        torch.cuda.empty_cache()
+        print_gpu_memory_usage("After sync model weights in sharding manager")
+
     def __enter__(self):
         # NOTE: Basically, we only need `torch.cuda.empty_cache()` before vllm wake_up and
         # after vllm sleep, since vllm has its own caching memory allocator CuMemAllocator.
@@ -94,27 +116,23 @@ def __enter__(self):
         # pytorch: https://pytorch.org/docs/stable/notes/cuda.html#memory-management
         # vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/device_allocator/cumem.py#L103
         torch.cuda.empty_cache()
-        print_gpu_memory_usage("Before state_dict() in sharding manager")
-        actor_weights = get_model_state_dict(self.module)
-        actor_weights = self._rename_weight_keys(actor_weights, self.module._fsdp_wrapped_module)
-        print_gpu_memory_usage("After state_dict() in sharding manager")
-
+        print_gpu_memory_usage("Before vllm wake up in sharding manager")
         if "tags" in inspect.signature(self.inference_engine.wake_up).parameters:
             self.inference_engine.wake_up(tags=["weights"])
         else:
             self.inference_engine.wake_up()
 
-        model = self.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner.model
-        model.load_weights(self._make_weight_iterator(actor_weights))
-        print_gpu_memory_usage("After sync model weights in sharding manager")
-
-        del actor_weights
-        torch.cuda.empty_cache()
+        if self.skip_vllm_sync_once:
+            self.skip_vllm_sync_once = False  # reset the flag
+            if is_rank0():
+                print("Skip vllm weight sync in sharding manager once.")
+        else:
+            self._sync_weight_to_vllm()
 
         if "tags" in inspect.signature(self.inference_engine.wake_up).parameters:
             self.inference_engine.wake_up(tags=["kv_cache"])
 
-        print_gpu_memory_usage("After del state_dict and empty_cache in sharding manager")
+        print_gpu_memory_usage("After vllm wake up in sharding manager")
         # important: need to manually set the random states of each tp to be identical.
         if self.device_mesh is not None:
             self.torch_random_states = torch.cuda.get_rng_state()