[rollout] feat: support best-of-n generation in vLLM (Jiayi-Pan#80)

Author: PeterSH6

examples/split_placement/config/ppo_trainer_split.yaml

@@ -66,6 +66,8 @@ actor_rollout_ref:
     log_prob_micro_batch_size: 128
     # for hf rollout
     do_sample: True
+    # number of responses (i.e. num sample times)
+    n: 1 # > 1 for grpo
 
 critic:
   strategy: fsdp

verl/trainer/config/ppo_megatron_trainer.yaml

@@ -72,6 +72,8 @@ actor_rollout_ref:
     layer_name_map:
       qkv_layer_name: qkv
       gate_proj_layer_name: gate_up
+    # number of responses (i.e. num sample times)
+    n: 1
 
 critic:
   strategy: megatron

verl/trainer/config/ppo_trainer.yaml

@@ -66,6 +66,8 @@ actor_rollout_ref:
     log_prob_micro_batch_size: 128
     # for hf rollout
     do_sample: True
+    # number of responses (i.e. num sample times)
+    n: 1 # > 1 for grpo
 
 critic:
   strategy: fsdp

verl/trainer/ppo/ray_trainer.py

@@ -486,6 +486,8 @@ def fit(self):
                 with _timer('gen', timing_raw):
                     gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch)
 
+                # repeat to align with repeated responses in rollout
+                batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)
                 batch = batch.union(gen_batch_output)
 
                 if self.use_reference_policy:

verl/workers/fsdp_workers.py

@@ -81,10 +81,14 @@ def __init__(self, config: DictConfig, role: str):
         if self._is_actor:
             self.config.actor.ppo_mini_batch_size //= self.device_mesh.shape[0]
             self.config.actor.ppo_micro_batch_size //= self.device_mesh.shape[0]
+            self.config.actor.ppo_mini_batch_size *= self.config.rollout.n
+            self.config.actor.ppo_micro_batch_size *= self.config.rollout.n
         if self._is_rollout:
             self.config.rollout.log_prob_micro_batch_size //= self.device_mesh.shape[0]
+            self.config.rollout.log_prob_micro_batch_size *= self.config.rollout.n
         if self._is_ref:
             self.config.ref.log_prob_micro_batch_size //= self.device_mesh.shape[0]
+            self.config.ref.log_prob_micro_batch_size *= self.config.rollout.n
 
     def _build_model_optimizer(self,
                                model_path,

verl/workers/rollout/vllm_rollout/vllm_rollout.py

@@ -167,6 +167,7 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
                 'top_k': -1,
                 'min_p': 0.0,
                 'temperature': 0,
+                'n': 1  # if greedy, only 1 response
             }
 
         # users can customize different sampling_params at different run
@@ -177,13 +178,20 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
                 prompt_token_ids=idx_list,
                 use_tqdm=False)
 
-        response = output[0].to(idx.device)  # (bs, response_length)
-        log_probs = output[1].to(idx.device)  # (bs, response_length)
+        # TODO(sgm): disable logprob when recompute_log_prob is enable
+        # if n = 1: (bs, response_length) ; if n > 1: (bs * n, response_length)
+        response = output[0].to(idx.device)
+        log_probs = output[1].to(idx.device)
 
         if response.shape[1] < self.config.response_length:
             response = pad_sequence_to_length(response, self.config.response_length, self.pad_token_id)
             log_probs = pad_sequence_to_length(log_probs, self.config.response_length, self.pad_token_id)
 
+        if self.config.n > 1 and do_sample:
+            idx = idx.repeat_interleave(self.config.n, dim=0)
+            attention_mask = attention_mask.repeat_interleave(self.config.n, dim=0)
+            position_ids = position_ids.repeat_interleave(self.config.n, dim=0)
+            batch_size = batch_size * self.config.n
         seq = torch.cat([idx, response], dim=-1)
 
         response_length = response.size(1)