diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py index 5aaa1111477..b8a8f7f530e 100644 --- a/tests/test_grpo_trainer.py +++ b/tests/test_grpo_trainer.py @@ -792,6 +792,41 @@ def test_training_vllm_guided_decoding(self): new_param = trainer.model.get_parameter(n) self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.") + @require_vllm + @unittest.skip("We should add a mock for the vLLM server.") + def test_training_vllm_importance_sampling_correction(self): + """Test that training works with vLLM for generation with guided decoding.""" + dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + + training_args = GRPOConfig( + output_dir=self.tmp_dir, + learning_rate=0.1, + per_device_train_batch_size=3, + num_generations=3, + max_completion_length=8, + report_to="none", + use_vllm=True, + vllm_importance_sampling_correction=True, + vllm_importance_sampling_cap=3.0, + ) + trainer = GRPOTrainer( + model="Qwen/Qwen2.5-0.5B-Instruct", # tiny model is too small for vLLM + reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5", + args=training_args, + train_dataset=dataset, + ) + + previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()} + + trainer.train() + + self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"]) + + # Check that the params have changed + for n, param in previous_trainable_params.items(): + new_param = trainer.model.get_parameter(n) + self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.") + def test_training_with_additional_generation_kwargs(self): """Test that training works with additional generation kwargs.""" dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index feb2dd4261a..aa2238c0eba 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -209,8 +209,11 @@ def generate( will override them. Returns: - `list[list[int]]`: - List of lists of token IDs representing the model-generated completions for each prompt. + `dict` with keys: + - `completion_ids` (`list[list[int]]`): + List of lists of token IDs representing the model-generated completions for each prompt. + - `logprobs` (`list[list[float]]`): + List of lists of log probabilities for each generated token. """ url = f"{self.base_url}/generate/" @@ -240,7 +243,8 @@ def pil_to_base64(image): }, ) if response.status_code == 200: - return response.json()["completion_ids"] + json_response = response.json() + return {"completion_ids": json_response["completion_ids"], "logprobs": json_response["logprobs"]} else: raise Exception(f"Request failed: {response.status_code}, {response.text}") diff --git a/trl/scripts/vllm_serve.py b/trl/scripts/vllm_serve.py index 41efe030527..12e8dc2dd61 100644 --- a/trl/scripts/vllm_serve.py +++ b/trl/scripts/vllm_serve.py @@ -377,6 +377,17 @@ def chunk_list(lst: list, n: int) -> list[list]: return [lst[i * k + min(i, r) : (i + 1) * k + min(i + 1, r)] for i in range(n)] +def sanitize_logprob(logprob): + import math + + value = logprob.logprob + if math.isnan(value): + logger.warning(f"Generated NaN logprob, token logprob '{logprob}' will be ignored") + return None + + return value + + def main(script_args: ScriptArguments): if not is_fastapi_available(): raise ImportError( @@ -468,6 +479,7 @@ class GenerateRequest(BaseModel): class GenerateResponse(BaseModel): completion_ids: list[list[int]] + logprobs: list[list[float]] @app.post("/generate/", response_model=GenerateResponse) async def generate(request: GenerateRequest): @@ -500,6 +512,8 @@ async def generate(request: GenerateRequest): Returns: `GenerateResponse`: - `completion_ids` (list of list of `int`): A list of lists of token IDs for each generated completion. + - `logprobs` (list of list of `float`): A list of lists of log probabilities for each token in the + generated completions. Example request: ```json @@ -508,7 +522,7 @@ async def generate(request: GenerateRequest): Example response: ```json - {"completion_ids": [[101, 102, 103], [201, 202, 203]]} + {"completion_ids": [[101, 102, 103], [201, 202, 203]], "logprobs": [[-0.1, -0.2, -0.3], [-0.4, -0.5, -0.6]]} ``` """ request.images = request.images or [None] * len(request.prompts) @@ -535,6 +549,7 @@ async def generate(request: GenerateRequest): "min_p": request.min_p, "max_tokens": request.max_tokens, "guided_decoding": guided_decoding, + "logprobs": 0, } generation_kwargs.update(request.generation_kwargs) sampling_params = SamplingParams(**generation_kwargs) @@ -561,7 +576,12 @@ async def generate(request: GenerateRequest): # Flatten and combine all results all_outputs = list(chain.from_iterable(all_outputs)) # from list of list to single list completion_ids = [list(output.token_ids) for outputs in all_outputs for output in outputs.outputs] - return {"completion_ids": completion_ids} + logprobs: list[list[float]] = [ + [sanitize_logprob(next(iter(logprob.values()))) for logprob in output.logprobs] + for outputs in all_outputs + for output in outputs.outputs + ] + return {"completion_ids": completion_ids, "logprobs": logprobs} class InitCommunicatorRequest(BaseModel): host: str diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index 86461889321..5261420c01f 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -220,6 +220,15 @@ class GRPOConfig(TrainingArguments): `mask_truncated_completions=True`, only tokens from non-truncated completions are considered. use_liger_loss (`bool`, *optional*, defaults to `False`): Whether to use the Liger GRPO loss. + vllm_importance_sampling_correction (`bool`, *optional*, defaults to `True`): + Whether to apply Truncated Importance Sampling (TIS) between vLLM completion logprobs and recomputed + logprobs. [Your Efficient RL Framework Secretly Brings You Off-Policy RL + Training](https://fengyao.notion.site/off-policy-rl) highlights that using a separate generation framework + (such as vLLM) can introduce off-policy effects due to subtle implementation differences between generation + and training backends. TIS is proposed as a remedy for this issue. + vllm_importance_sampling_cap (`float`, *optional*, defaults to `2.0`): + Truncation parameter C for Truncated Importance Sampling (TIS). This sets an upper bound on the importance + sampling ratio, improving training stability. > Parameters that control the logging @@ -589,6 +598,23 @@ class GRPOConfig(TrainingArguments): default=False, metadata={"help": "Whether to use the Liger GRPO loss."}, ) + vllm_importance_sampling_correction: bool = field( + default=True, + metadata={ + "help": "Whether to apply Truncated Importance Sampling (TIS) between vLLM completion logprobs and " + "recomputed logprobs. Your Efficient RL Framework Secretly Brings You Off-Policy RL " + "Training highlights that using a separate generation framework (such as vLLM) can introduce off-policy " + "effects due to subtle implementation differences between generation and training backends. TIS is " + "proposed as a remedy for this issue." + }, + ) + vllm_importance_sampling_cap: float = field( + default=2.0, + metadata={ + "help": "Truncation parameter C for Truncated Importance Sampling (TIS). This sets an upper bound on the " + "importance sampling ratio, improving training stability." + }, + ) # Parameters that control the logging log_completions: bool = field( diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index c7b29e151ad..3bb1a0ab4d7 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -350,6 +350,8 @@ def __init__( self.vllm_mode = args.vllm_mode self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization # only applies to colocation mode self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size # only applies to colocation mode + self.vllm_importance_sampling_correction = args.vllm_importance_sampling_correction + self.vllm_importance_sampling_cap = args.vllm_importance_sampling_cap self.use_liger_loss = args.use_liger_loss self.loss_type = args.loss_type self.scale_rewards = args.scale_rewards @@ -1158,7 +1160,7 @@ def _generate_and_score_completions( ordered_set_of_images = None with profiling_context(self, "vLLM.generate"): - completion_ids = self.vllm_client.generate( + output = self.vllm_client.generate( prompts=ordered_set_of_prompts, images=ordered_set_of_images, n=self.num_generations, @@ -1171,16 +1173,21 @@ def _generate_and_score_completions( guided_decoding_regex=self.guided_decoding_regex, generation_kwargs=self.args.generation_kwargs, ) + payload = (output["completion_ids"], output["logprobs"]) else: - completion_ids = [None] * len(all_prompts_text) - # Broadcast the completions from the main process to all processes, ensuring each process receives its - # corresponding slice. - completion_ids = broadcast_object_list(completion_ids, from_process=0) + payload = None + + # Broadcast the completions from the main process to all processes, ensuring each process receives its corresponding slice. + obj_list = [payload] + broadcast_object_list(obj_list, from_process=0) + completion_ids, all_logprobs = obj_list[0] + process_slice = slice( self.accelerator.process_index * len(prompts), (self.accelerator.process_index + 1) * len(prompts), ) completion_ids = completion_ids[process_slice] + all_logprobs = all_logprobs[process_slice] # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts elif self.vllm_mode == "colocate": @@ -1198,6 +1205,7 @@ def _generate_and_score_completions( "min_p": 0.0 if self.min_p is None else self.min_p, "max_tokens": self.max_completion_length, "guided_decoding": guided_decoding, + "logprobs": 0, # only return the logprob of the generated token } if self.args.generation_kwargs is not None: generation_kwargs.update(self.args.generation_kwargs) @@ -1235,6 +1243,11 @@ def _generate_and_score_completions( all_outputs = self.llm.generate(vllm_inputs, sampling_params=sampling_params, use_tqdm=False) completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] + all_logprobs = [ + [next(iter(lp.values())).logprob for lp in output.logprobs] + for outputs in all_outputs + for output in outputs.outputs + ] if self.vllm_tensor_parallel_size > 1: # Slice completions for this rank within its TP group. @@ -1242,6 +1255,7 @@ def _generate_and_score_completions( local_rank_in_group = torch.distributed.get_rank(group=self.tp_group) tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) completion_ids = completion_ids[tp_slice] + all_logprobs = all_logprobs[tp_slice] if self.args.vllm_enable_sleep_mode: self.llm.sleep(level=1) @@ -1250,6 +1264,10 @@ def _generate_and_score_completions( completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids] completion_ids = pad(completion_ids, padding_value=self.pad_token_id) prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) + sampling_per_token_logps = [ + torch.tensor(logprobs, device=device, dtype=torch.float32) for logprobs in all_logprobs + ] + sampling_per_token_logps = pad(sampling_per_token_logps, padding_value=0.0) elif self.use_transformers_paged: # Re-process inputs for paged generation if needed @@ -1338,8 +1356,12 @@ def _generate_and_score_completions( # samples may come from an earlier version of the model. In that case, we need to track old_per_token_logps # for importance sampling. If the steps are aligned, importance sampling isn't necessary and we set # old_per_token_logps to None. + # When using vLLM, we always compute old_per_token_logps for importance sampling, it was shown that the + # distribution mismatch between vLLM and the training model can be large and harm the training. generate_every = self.args.steps_per_generation * self.num_iterations # generation frequency - if self.args.gradient_accumulation_steps % generate_every != 0: + if self.args.gradient_accumulation_steps % generate_every != 0 or ( + self.use_vllm and self.vllm_importance_sampling_correction + ): old_per_token_logps, _ = self._get_per_token_logps_and_entropies( self.model, prompt_completion_ids, @@ -1354,6 +1376,13 @@ def _generate_and_score_completions( else: old_per_token_logps = None + # Compute the importance sampling ratio when using vLLM, to correct for potential distribution mismatch + if self.use_vllm and self.vllm_importance_sampling_correction: + importance_sampling_ratio = torch.exp(old_per_token_logps - sampling_per_token_logps) + importance_sampling_ratio = torch.clamp( + importance_sampling_ratio, max=self.vllm_importance_sampling_cap + ) + # Compute the per-token log probabilities for the reference model if self.beta != 0.0: if self.ref_model is not None: @@ -1474,6 +1503,38 @@ def _generate_and_score_completions( if has_images: self._logs["image"].extend(gather_object(images)) + if self.use_vllm and self.vllm_importance_sampling_correction: + delta = torch.abs(old_per_token_logps - sampling_per_token_logps) + delta = delta[completion_mask.bool()] + mean_delta = torch.mean(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) + max_delta = torch.max(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) + self._metrics[mode]["sampling/sampling_logp_difference/mean"].append( + self.accelerator.gather(mean_delta).mean().item() + ) + self._metrics[mode]["sampling/sampling_logp_difference/max"].append( + self.accelerator.gather(max_delta).max().item() + ) + + flat_is_ratio = importance_sampling_ratio[completion_mask.bool()] + min_importance_sampling_ratio = ( + torch.min(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + mean_importance_sampling_ratio = ( + torch.mean(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + max_importance_sampling_ratio = ( + torch.max(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( + nanmin(self.accelerator.gather(min_importance_sampling_ratio)).item() + ) + self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append( + self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item() + ) + self._metrics[mode]["sampling/importance_sampling_ratio/max"].append( + nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item() + ) + output = { "prompt_ids": prompt_ids, "prompt_mask": prompt_mask, @@ -1484,6 +1545,8 @@ def _generate_and_score_completions( } if old_per_token_logps is not None: output["old_per_token_logps"] = old_per_token_logps + if self.use_vllm and self.vllm_importance_sampling_correction: + output["importance_sampling_ratio"] = importance_sampling_ratio if ref_per_token_logps is not None: output["ref_per_token_logps"] = ref_per_token_logps if "pixel_values" in prompt_inputs: @@ -1584,9 +1647,11 @@ def _compute_loss(self, model, inputs): # Compute the loss advantages = inputs["advantages"] - # When using num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps - # old_per_token_logps == per_token_logps, so we can skip it's computation - # (see _generate_and_score_completions) and use per_token_logps.detach() instead. + # When num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps, + # old_per_token_logps == per_token_logps. In this case we can skip its computation + # (see _generate_and_score_completions) and instead use per_token_logps.detach(). + # The exception is when using vLLM, where we always compute old_per_token_logps + # for importance sampling old_per_token_logps = inputs.get("old_per_token_logps") old_per_token_logps = per_token_logps.detach() if old_per_token_logps is None else old_per_token_logps @@ -1616,6 +1681,10 @@ def _compute_loss(self, model, inputs): per_token_loss = -torch.min(per_token_loss1, per_token_loss2) if entropy_mask is not None: per_token_loss = per_token_loss * entropy_mask + + if self.use_vllm and self.vllm_importance_sampling_correction: + per_token_loss = per_token_loss * inputs["importance_sampling_ratio"] + if self.beta != 0.0: per_token_loss = per_token_loss + self.beta * per_token_kl