From 2ef3af617404e8f919cd931bc33932a4bcc1ac2c Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Thu, 7 Aug 2025 15:33:39 +0200
Subject: [PATCH 01/23] compare vllm rollout probs to actor model

---
 trl/extras/vllm_client.py   |  3 ++-
 trl/scripts/vllm_serve.py   | 11 ++++++++++-
 trl/trainer/grpo_trainer.py |  7 ++++++-
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py
index 94312e53f56..f0c56208cab 100644
--- a/trl/extras/vllm_client.py
+++ b/trl/extras/vllm_client.py
@@ -240,7 +240,8 @@ def pil_to_base64(image):
             },
         )
         if response.status_code == 200:
-            return response.json()["completion_ids"]
+            json_response = response.json()
+            return json_response["completion_ids"], 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 fd320710d01..1b0abb80d92 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -467,6 +467,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):
@@ -525,6 +526,7 @@ async def generate(request: GenerateRequest):
             "min_p": request.min_p,
             "max_tokens": request.max_tokens,
             "guided_decoding": guided_decoding,
+            "logprobs": 0,  # this gives us logps of only the generated token
         }
         generation_kwargs.update(request.generation_kwargs)
         sampling_params = SamplingParams(**generation_kwargs)
@@ -550,8 +552,15 @@ 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]] = [
+            [next(iter(lp.values())).logprob for lp in output.logprobs]
+            for outputs in all_outputs
+            for output in outputs.outputs
+        ]
+        # return {"completion_ids": completion_ids}
+        return {"completion_ids": completion_ids, "logprobs": logprobs}
 
     class InitCommunicatorRequest(BaseModel):
         host: str
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 8cc106e58bb..b98627f5257 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1439,7 +1439,7 @@ def _generate_and_score_completions(
                         ordered_set_of_images = None
 
                     with profiling_context(self, "vLLM.generate"):
-                        completion_ids = self.vllm_client.generate(
+                        completion_ids, vllm_logprobs = self.vllm_client.generate(
                             prompts=ordered_set_of_prompts,
                             images=ordered_set_of_images,
                             n=self.num_generations,
@@ -1742,6 +1742,11 @@ def _generate_and_score_completions(
         if has_images:
             self._logs["image"].extend(gather_object(images))
 
+        vllm_logps = torch.tensor(vllm_logprobs, device=device, dtype=torch.float32)
+        logps_diff = torch.abs(vllm_logps - old_per_token_logps)
+        self._metrics[mode]["Token Probability Difference/max"].append(logps_diff.max().exp().item())
+        self._metrics[mode]["Token Probability Difference/mean"].append(logps_diff.mean().exp().item())
+
         output = {
             "prompt_ids": prompt_ids,
             "prompt_mask": prompt_mask,

From 302b981acec5f36b58184bc38b9eb0ac2228e533 Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Mon, 18 Aug 2025 10:14:44 +0200
Subject: [PATCH 02/23] prompt logprobs

---
 trl/extras/vllm_client.py   |  7 ++++-
 trl/scripts/vllm_serve.py   | 33 +++++++++++++++------
 trl/trainer/grpo_trainer.py | 59 +++++++++++++++++++++++++++++++------
 3 files changed, 80 insertions(+), 19 deletions(-)

diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py
index f0c56208cab..390b41b3ac2 100644
--- a/trl/extras/vllm_client.py
+++ b/trl/extras/vllm_client.py
@@ -178,6 +178,8 @@ def generate(
         max_tokens: int = 16,
         guided_decoding_regex: Optional[str] = None,
         generation_kwargs: Optional[dict] = None,
+        logprobs: Optional[int] = None,
+        prompt_logprobs: Optional[int] = None,  # TODO: DOC STRING
     ) -> list[list[int]]:
         """
         Generates model completions for the provided prompts.
@@ -237,11 +239,14 @@ def pil_to_base64(image):
                 "max_tokens": max_tokens,
                 "guided_decoding_regex": guided_decoding_regex,
                 "generation_kwargs": generation_kwargs or {},
+                "logprobs": logprobs,
+                "prompt_logprobs": prompt_logprobs,
             },
         )
         if response.status_code == 200:
             json_response = response.json()
-            return json_response["completion_ids"], json_response["logprobs"]
+            result = {k: v for k, v in json_response.items() if k in ("completion_ids", "logprobs", "prompt_logprobs")}
+            return result
         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 1b0abb80d92..fc75a83cf45 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -464,9 +464,12 @@ class GenerateRequest(BaseModel):
         max_tokens: int = 16
         guided_decoding_regex: Optional[str] = None
         generation_kwargs: dict = field(default_factory=dict)
+        logprobs: Optional[int] = None
+        prompt_logprobs: Optional[int] = None
 
     class GenerateResponse(BaseModel):
         completion_ids: list[list[int]]
+        prompt_logprobs: list[list[float]]
         logprobs: list[list[float]]
 
     @app.post("/generate/", response_model=GenerateResponse)
@@ -526,14 +529,14 @@ async def generate(request: GenerateRequest):
             "min_p": request.min_p,
             "max_tokens": request.max_tokens,
             "guided_decoding": guided_decoding,
-            "logprobs": 0,  # this gives us logps of only the generated token
+            "logprobs": request.logprobs,
+            "prompt_logprobs": request.prompt_logprobs,
         }
         generation_kwargs.update(request.generation_kwargs)
         sampling_params = SamplingParams(**generation_kwargs)
 
         # Evenly distribute prompts across DP ranks
         chunked_prompts = chunk_list(prompts, script_args.data_parallel_size)
-
         # Send the prompts to each worker
         for connection, prompts in zip(connections, chunked_prompts):
             # When the number of prompts is less than data_parallel_size, some workers will receive empty prompts.
@@ -554,13 +557,25 @@ async def generate(request: GenerateRequest):
         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]
-        logprobs: list[list[float]] = [
-            [next(iter(lp.values())).logprob for lp in output.logprobs]
-            for outputs in all_outputs
-            for output in outputs.outputs
-        ]
-        # return {"completion_ids": completion_ids}
-        return {"completion_ids": completion_ids, "logprobs": logprobs}
+
+        response = {"completion_ids": completion_ids}
+        if request.logprobs is not None:
+            logprobs: list[list[float]] = [
+                [next(iter(lp.values())).logprob for lp in output.logprobs]
+                for outputs in all_outputs
+                for output in outputs.outputs
+            ]
+
+            response["logprobs"] = logprobs
+
+        if request.prompt_logprobs is not None:
+            prompt_logprob: list[list[float]] = [
+                [next(iter(d.values())).logprob for d in output.prompt_logprobs[1:]]  # one float per dict
+                for output in all_outputs
+                # for pl in output.prompt_logprobs                      # one list per PromptLogprobs
+            ]
+            response["prompt_logprobs"] = prompt_logprob
+        return response
 
     class InitCommunicatorRequest(BaseModel):
         host: str
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index b98627f5257..9f3b14b89dd 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1439,7 +1439,7 @@ def _generate_and_score_completions(
                         ordered_set_of_images = None
 
                     with profiling_context(self, "vLLM.generate"):
-                        completion_ids, vllm_logprobs = self.vllm_client.generate(
+                        output = self.vllm_client.generate(
                             prompts=ordered_set_of_prompts,
                             images=ordered_set_of_images,
                             n=self.num_generations,
@@ -1451,17 +1451,29 @@ def _generate_and_score_completions(
                             max_tokens=self.max_completion_length,
                             guided_decoding_regex=self.guided_decoding_regex,
                             generation_kwargs=self.args.generation_kwargs,
+                            prompt_logprobs=0,
+                            logprobs=0,
+                        )
+                        payload = (
+                            output["completion_ids"],
+                            output["prompt_logprobs"],
+                            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, prompt_logprobs, 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]
+                prompt_logprobs = prompt_logprobs[process_slice]
+                logprobs = 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":
@@ -1742,12 +1754,13 @@ def _generate_and_score_completions(
         if has_images:
             self._logs["image"].extend(gather_object(images))
 
-        vllm_logps = torch.tensor(vllm_logprobs, device=device, dtype=torch.float32)
-        logps_diff = torch.abs(vllm_logps - old_per_token_logps)
-        self._metrics[mode]["Token Probability Difference/max"].append(logps_diff.max().exp().item())
-        self._metrics[mode]["Token Probability Difference/mean"].append(logps_diff.mean().exp().item())
+        # vllm_logps = torch.tensor(vllm_logprobs, device=device, dtype=torch.float32)
+        # logps_diff = torch.abs(vllm_logps - old_per_token_logps)
+        # self._metrics[mode]["Token Probability Difference/max"].append(logps_diff.max().exp().item())
+        # self._metrics[mode]["Token Probability Difference/mean"].append(logps_diff.mean().exp().item())
 
         output = {
+            "prompt_texts": prompts_text,
             "prompt_ids": prompt_ids,
             "prompt_mask": prompt_mask,
             "completion_ids": completion_ids,
@@ -1842,6 +1855,34 @@ def _compute_loss(self, model, inputs):
             image_sizes=inputs.get("image_sizes"),
         )
 
+        if self.use_vllm:
+            prompt_texts = inputs["prompts_text"]
+            all_prompts_text = gather_object(prompt_texts)
+            if self.accelerator.is_main_process:
+                ordered_set_of_prompts = all_prompts_text[:: self.num_generations]
+
+                ordered_set_of_images = None
+
+                output = self.vllm_client.generate(
+                    prompts=ordered_set_of_prompts,
+                    images=ordered_set_of_images,
+                    n=self.num_generations,
+                    repetition_penalty=self.repetition_penalty,
+                    temperature=self.temperature,
+                    top_p=self.top_p,
+                    top_k=-1 if self.top_k is None else self.top_k,
+                    min_p=0.0 if self.min_p is None else self.min_p,
+                    max_tokens=0,
+                    guided_decoding_regex=self.guided_decoding_regex,
+                    generation_kwargs=self.args.generation_kwargs,
+                    prompt_logprobs=-1,
+                )
+                completion_ids = output["completion_ids"]
+            else:
+                completion_ids = [None] * len(all_prompts_text)
+
+            # Calculate entropy.
+
         if self.top_entropy_quantile < 1.0:
             entropy_mask = get_high_entropy_mask(entropies, completion_mask, 1 - self.top_entropy_quantile)
         else:

From a57a70f60d899e43d8fc5ee6af4453b803a33e91 Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Mon, 18 Aug 2025 12:29:04 +0200
Subject: [PATCH 03/23] implemented TIS. only for testing, lacks docs and
 cleaning

---
 trl/extras/vllm_client.py   |  4 +---
 trl/scripts/vllm_serve.py   | 10 +-------
 trl/trainer/grpo_trainer.py | 48 +++++++++++++------------------------
 3 files changed, 18 insertions(+), 44 deletions(-)

diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py
index 390b41b3ac2..178e2fc7cac 100644
--- a/trl/extras/vllm_client.py
+++ b/trl/extras/vllm_client.py
@@ -179,7 +179,6 @@ def generate(
         guided_decoding_regex: Optional[str] = None,
         generation_kwargs: Optional[dict] = None,
         logprobs: Optional[int] = None,
-        prompt_logprobs: Optional[int] = None,  # TODO: DOC STRING
     ) -> list[list[int]]:
         """
         Generates model completions for the provided prompts.
@@ -240,12 +239,11 @@ def pil_to_base64(image):
                 "guided_decoding_regex": guided_decoding_regex,
                 "generation_kwargs": generation_kwargs or {},
                 "logprobs": logprobs,
-                "prompt_logprobs": prompt_logprobs,
             },
         )
         if response.status_code == 200:
             json_response = response.json()
-            result = {k: v for k, v in json_response.items() if k in ("completion_ids", "logprobs", "prompt_logprobs")}
+            result = {k: v for k, v in json_response.items() if k in ("completion_ids", "logprobs")}
             return result
         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 fc75a83cf45..671f9fb43e1 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -469,7 +469,6 @@ class GenerateRequest(BaseModel):
 
     class GenerateResponse(BaseModel):
         completion_ids: list[list[int]]
-        prompt_logprobs: list[list[float]]
         logprobs: list[list[float]]
 
     @app.post("/generate/", response_model=GenerateResponse)
@@ -530,7 +529,6 @@ async def generate(request: GenerateRequest):
             "max_tokens": request.max_tokens,
             "guided_decoding": guided_decoding,
             "logprobs": request.logprobs,
-            "prompt_logprobs": request.prompt_logprobs,
         }
         generation_kwargs.update(request.generation_kwargs)
         sampling_params = SamplingParams(**generation_kwargs)
@@ -559,6 +557,7 @@ async def generate(request: GenerateRequest):
         completion_ids = [list(output.token_ids) for outputs in all_outputs for output in outputs.outputs]
 
         response = {"completion_ids": completion_ids}
+
         if request.logprobs is not None:
             logprobs: list[list[float]] = [
                 [next(iter(lp.values())).logprob for lp in output.logprobs]
@@ -568,13 +567,6 @@ async def generate(request: GenerateRequest):
 
             response["logprobs"] = logprobs
 
-        if request.prompt_logprobs is not None:
-            prompt_logprob: list[list[float]] = [
-                [next(iter(d.values())).logprob for d in output.prompt_logprobs[1:]]  # one float per dict
-                for output in all_outputs
-                # for pl in output.prompt_logprobs                      # one list per PromptLogprobs
-            ]
-            response["prompt_logprobs"] = prompt_logprob
         return response
 
     class InitCommunicatorRequest(BaseModel):
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 9f3b14b89dd..e9675b5eb9b 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -654,6 +654,7 @@ def __init__(
         self.reward_processing_classes = reward_processing_classes
 
         # Training arguments
+        self.rollout_importance_sampling_cap = 2
         self.max_prompt_length = args.max_prompt_length
         self.max_completion_length = args.max_completion_length  # = |o_i| in the GRPO paper
         self.num_generations = args.num_generations  # = G in the GRPO paper
@@ -1451,12 +1452,10 @@ def _generate_and_score_completions(
                             max_tokens=self.max_completion_length,
                             guided_decoding_regex=self.guided_decoding_regex,
                             generation_kwargs=self.args.generation_kwargs,
-                            prompt_logprobs=0,
                             logprobs=0,
                         )
                         payload = (
                             output["completion_ids"],
-                            output["prompt_logprobs"],
                             output["logprobs"],
                         )
                 else:
@@ -1465,14 +1464,13 @@ def _generate_and_score_completions(
                 # 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, prompt_logprobs, logprobs = obj_list[0]
+                completion_ids, 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]
-                prompt_logprobs = prompt_logprobs[process_slice]
                 logprobs = logprobs[process_slice]
 
             # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts
@@ -1537,9 +1535,12 @@ def _generate_and_score_completions(
                     completion_ids = completion_ids[tp_slice]
 
             # Pad the completions, and concatenate them with the prompts
+            completion_ids[0].append(24)
             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)
+            vllm_logprobs = [torch.tensor(logp, device=device) for logp in logprobs]
+            vllm_logprobs = pad(vllm_logprobs, padding_value=0.0)
 
         elif self.use_transformers_paged:
             # Re-process inputs for paged generation if needed
@@ -1641,8 +1642,10 @@ def _generate_and_score_completions(
                     pixel_attention_mask=prompt_inputs.get("pixel_attention_mask"),
                     image_sizes=prompt_inputs.get("image_sizes"),
                 )
+                vllm_old_per_token_logps = vllm_logprobs
             else:
                 old_per_token_logps = None
+                vllm_old_per_token_logps = None
 
             # Compute the per-token log probabilities for the reference model
             if self.beta != 0.0:
@@ -1769,6 +1772,8 @@ def _generate_and_score_completions(
         }
         if old_per_token_logps is not None:
             output["old_per_token_logps"] = old_per_token_logps
+        if vllm_old_per_token_logps is not None:
+            output["vllm_old_per_token_logps"] = vllm_old_per_token_logps
         if ref_per_token_logps is not None:
             output["ref_per_token_logps"] = ref_per_token_logps
         if "pixel_values" in prompt_inputs:
@@ -1855,34 +1860,6 @@ def _compute_loss(self, model, inputs):
             image_sizes=inputs.get("image_sizes"),
         )
 
-        if self.use_vllm:
-            prompt_texts = inputs["prompts_text"]
-            all_prompts_text = gather_object(prompt_texts)
-            if self.accelerator.is_main_process:
-                ordered_set_of_prompts = all_prompts_text[:: self.num_generations]
-
-                ordered_set_of_images = None
-
-                output = self.vllm_client.generate(
-                    prompts=ordered_set_of_prompts,
-                    images=ordered_set_of_images,
-                    n=self.num_generations,
-                    repetition_penalty=self.repetition_penalty,
-                    temperature=self.temperature,
-                    top_p=self.top_p,
-                    top_k=-1 if self.top_k is None else self.top_k,
-                    min_p=0.0 if self.min_p is None else self.min_p,
-                    max_tokens=0,
-                    guided_decoding_regex=self.guided_decoding_regex,
-                    generation_kwargs=self.args.generation_kwargs,
-                    prompt_logprobs=-1,
-                )
-                completion_ids = output["completion_ids"]
-            else:
-                completion_ids = [None] * len(all_prompts_text)
-
-            # Calculate entropy.
-
         if self.top_entropy_quantile < 1.0:
             entropy_mask = get_high_entropy_mask(entropies, completion_mask, 1 - self.top_entropy_quantile)
         else:
@@ -1929,6 +1906,13 @@ 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.rollout_importance_sampling_cap > 0 and old_per_token_logps is not None:
+            vllm_old_per_token_logps = inputs.get("vllm_old_per_token_logps")
+            importance_sampling_ratio = torch.exp(old_per_token_logps - vllm_old_per_token_logps)
+            importance_sampling_ratio = torch.min(importance_sampling_ratio, self.rollout_importance_sampling_cap)
+            per_token_loss = per_token_loss * importance_sampling_ratio
+
         if self.beta != 0.0:
             per_token_loss = per_token_loss + self.beta * per_token_kl
 

From a8e27e59f954c92c81c0e6de65edfabd870e1e82 Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Mon, 18 Aug 2025 12:52:16 +0200
Subject: [PATCH 04/23] rename

---
 trl/trainer/grpo_trainer.py | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index e9675b5eb9b..6b954acc11d 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1642,10 +1642,10 @@ def _generate_and_score_completions(
                     pixel_attention_mask=prompt_inputs.get("pixel_attention_mask"),
                     image_sizes=prompt_inputs.get("image_sizes"),
                 )
-                vllm_old_per_token_logps = vllm_logprobs
+                rollout_old_per_token_logps = vllm_logprobs
             else:
                 old_per_token_logps = None
-                vllm_old_per_token_logps = None
+                rollout_old_per_token_logps = None
 
             # Compute the per-token log probabilities for the reference model
             if self.beta != 0.0:
@@ -1757,10 +1757,10 @@ def _generate_and_score_completions(
         if has_images:
             self._logs["image"].extend(gather_object(images))
 
-        # vllm_logps = torch.tensor(vllm_logprobs, device=device, dtype=torch.float32)
-        # logps_diff = torch.abs(vllm_logps - old_per_token_logps)
-        # self._metrics[mode]["Token Probability Difference/max"].append(logps_diff.max().exp().item())
-        # self._metrics[mode]["Token Probability Difference/mean"].append(logps_diff.mean().exp().item())
+        if old_per_token_logps is not None:
+            probs_diff = torch.exp(old_per_token_logps - rollout_old_per_token_logps)
+            self._metrics[mode]["Token Probability Difference/max"].append(probs_diff.max().exp().item())
+            self._metrics[mode]["Token Probability Difference/mean"].append(probs_diff.mean().exp().item())
 
         output = {
             "prompt_texts": prompts_text,
@@ -1772,8 +1772,8 @@ def _generate_and_score_completions(
         }
         if old_per_token_logps is not None:
             output["old_per_token_logps"] = old_per_token_logps
-        if vllm_old_per_token_logps is not None:
-            output["vllm_old_per_token_logps"] = vllm_old_per_token_logps
+        if rollout_old_per_token_logps is not None:
+            output["rollout_old_per_token_logps"] = rollout_old_per_token_logps
         if ref_per_token_logps is not None:
             output["ref_per_token_logps"] = ref_per_token_logps
         if "pixel_values" in prompt_inputs:
@@ -1907,11 +1907,11 @@ def _compute_loss(self, model, inputs):
         if entropy_mask is not None:
             per_token_loss = per_token_loss * entropy_mask
 
-        if self.rollout_importance_sampling_cap > 0 and old_per_token_logps is not None:
-            vllm_old_per_token_logps = inputs.get("vllm_old_per_token_logps")
-            importance_sampling_ratio = torch.exp(old_per_token_logps - vllm_old_per_token_logps)
-            importance_sampling_ratio = torch.min(importance_sampling_ratio, self.rollout_importance_sampling_cap)
-            per_token_loss = per_token_loss * importance_sampling_ratio
+        if self.use_vllm and self.rollout_importance_sampling_cap > 0 and old_per_token_logps is not None:
+            rollout_old_per_token_logps = inputs.get("rollout_old_per_token_logps")
+            rollout_imp_ratio = torch.exp(old_per_token_logps - rollout_old_per_token_logps)
+            rollout_imp_ratio = torch.min(rollout_imp_ratio, self.rollout_importance_sampling_cap)
+            per_token_loss = per_token_loss * rollout_imp_ratio
 
         if self.beta != 0.0:
             per_token_loss = per_token_loss + self.beta * per_token_kl

From caa1ed59e7095a42bcad9dae25b83f1dfa97836a Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Thu, 21 Aug 2025 07:43:17 +0200
Subject: [PATCH 05/23] wip

---
 trl/scripts/vllm_serve.py   | 20 +++++++++++++++++++-
 trl/trainer/grpo_trainer.py | 23 +++++++++++++----------
 2 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/trl/scripts/vllm_serve.py b/trl/scripts/vllm_serve.py
index 671f9fb43e1..dcb7eee5a89 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -376,6 +376,24 @@ 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_float(logprob):
+    import math
+
+    value = logprob.logprob
+    """Replace inf, -inf, and nan with None for JSON compatibility."""
+    if math.isinf(value):
+        # print(f"Warning: Replacing {'positive' if value > 0 else 'negative'} infinity with None.")
+        # print(logprob)
+        exit()
+        return None
+    elif math.isnan(value):
+        # print("Warning: Replacing NaN with None.")
+        # print(logprob)
+        exit()
+        return None
+    return value
+
+
 def main(script_args: ScriptArguments):
     if not is_fastapi_available():
         raise ImportError(
@@ -560,7 +578,7 @@ async def generate(request: GenerateRequest):
 
         if request.logprobs is not None:
             logprobs: list[list[float]] = [
-                [next(iter(lp.values())).logprob for lp in output.logprobs]
+                [sanitize_float(next(iter(lp.values()))) for lp in output.logprobs]
                 for outputs in all_outputs
                 for output in outputs.outputs
             ]
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 6b954acc11d..f08235160dc 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1535,12 +1535,11 @@ def _generate_and_score_completions(
                     completion_ids = completion_ids[tp_slice]
 
             # Pad the completions, and concatenate them with the prompts
-            completion_ids[0].append(24)
             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)
-            vllm_logprobs = [torch.tensor(logp, device=device) for logp in logprobs]
-            vllm_logprobs = pad(vllm_logprobs, padding_value=0.0)
+            rollout_old_per_token_logps = [torch.tensor(logp, device=device) for logp in logprobs]
+            rollout_old_per_token_logps = pad(rollout_old_per_token_logps, padding_value=0.0)
 
         elif self.use_transformers_paged:
             # Re-process inputs for paged generation if needed
@@ -1642,10 +1641,8 @@ def _generate_and_score_completions(
                     pixel_attention_mask=prompt_inputs.get("pixel_attention_mask"),
                     image_sizes=prompt_inputs.get("image_sizes"),
                 )
-                rollout_old_per_token_logps = vllm_logprobs
             else:
                 old_per_token_logps = None
-                rollout_old_per_token_logps = None
 
             # Compute the per-token log probabilities for the reference model
             if self.beta != 0.0:
@@ -1757,13 +1754,15 @@ def _generate_and_score_completions(
         if has_images:
             self._logs["image"].extend(gather_object(images))
 
-        if old_per_token_logps is not None:
+        if self.use_vllm and old_per_token_logps is not None:
             probs_diff = torch.exp(old_per_token_logps - rollout_old_per_token_logps)
+            per_token_kl = probs_diff - (old_per_token_logps - rollout_old_per_token_logps) - 1
+            mean_kl = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
+            self._metrics[mode]["Rollout KL"].append(self.accelerator.gather(mean_kl).nanmean().item())
             self._metrics[mode]["Token Probability Difference/max"].append(probs_diff.max().exp().item())
             self._metrics[mode]["Token Probability Difference/mean"].append(probs_diff.mean().exp().item())
 
         output = {
-            "prompt_texts": prompts_text,
             "prompt_ids": prompt_ids,
             "prompt_mask": prompt_mask,
             "completion_ids": completion_ids,
@@ -1772,7 +1771,7 @@ def _generate_and_score_completions(
         }
         if old_per_token_logps is not None:
             output["old_per_token_logps"] = old_per_token_logps
-        if rollout_old_per_token_logps is not None:
+        if self.use_vllm:
             output["rollout_old_per_token_logps"] = rollout_old_per_token_logps
         if ref_per_token_logps is not None:
             output["ref_per_token_logps"] = ref_per_token_logps
@@ -1910,8 +1909,8 @@ def _compute_loss(self, model, inputs):
         if self.use_vllm and self.rollout_importance_sampling_cap > 0 and old_per_token_logps is not None:
             rollout_old_per_token_logps = inputs.get("rollout_old_per_token_logps")
             rollout_imp_ratio = torch.exp(old_per_token_logps - rollout_old_per_token_logps)
-            rollout_imp_ratio = torch.min(rollout_imp_ratio, self.rollout_importance_sampling_cap)
-            per_token_loss = per_token_loss * rollout_imp_ratio
+            clammped_rollout_imp_ratio = torch.clamp(rollout_imp_ratio, max=self.rollout_importance_sampling_cap)
+            per_token_loss = per_token_loss * clammped_rollout_imp_ratio
 
         if self.beta != 0.0:
             per_token_loss = per_token_loss + self.beta * per_token_kl
@@ -1952,6 +1951,10 @@ def masked_batch_mean(x):
         high_clip = masked_batch_mean(is_high_clipped.float())
         clip_ratio = masked_batch_mean(is_region_clipped.float())
 
+        is_clamped_mask = rollout_imp_ratio > self.rollout_importance_sampling_cap
+        mean_clamped = masked_batch_mean(is_clamped_mask.float())
+        self._metrics[mode]["TIS/clamp_ratio"].append(self.accelerator.gather(mean_clamped).nanmean().item())
+
         gathered_low_clip = self.accelerator.gather(low_clip)
         self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item())
         self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item())

From f38a38ec649cf25545afb75ab0c3f365f9457dae Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Mon, 25 Aug 2025 20:18:01 +0200
Subject: [PATCH 06/23] cleanup v1

---
 trl/scripts/vllm_serve.py   | 1 -
 trl/trainer/grpo_trainer.py | 8 ++++----
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/trl/scripts/vllm_serve.py b/trl/scripts/vllm_serve.py
index dcb7eee5a89..ecb10b5e77f 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -483,7 +483,6 @@ class GenerateRequest(BaseModel):
         guided_decoding_regex: Optional[str] = None
         generation_kwargs: dict = field(default_factory=dict)
         logprobs: Optional[int] = None
-        prompt_logprobs: Optional[int] = None
 
     class GenerateResponse(BaseModel):
         completion_ids: list[list[int]]
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index f08235160dc..09683f72980 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -654,7 +654,7 @@ def __init__(
         self.reward_processing_classes = reward_processing_classes
 
         # Training arguments
-        self.rollout_importance_sampling_cap = 2
+        self.rollout_importance_sampling_cap = -1
         self.max_prompt_length = args.max_prompt_length
         self.max_completion_length = args.max_completion_length  # = |o_i| in the GRPO paper
         self.num_generations = args.num_generations  # = G in the GRPO paper
@@ -1951,9 +1951,9 @@ def masked_batch_mean(x):
         high_clip = masked_batch_mean(is_high_clipped.float())
         clip_ratio = masked_batch_mean(is_region_clipped.float())
 
-        is_clamped_mask = rollout_imp_ratio > self.rollout_importance_sampling_cap
-        mean_clamped = masked_batch_mean(is_clamped_mask.float())
-        self._metrics[mode]["TIS/clamp_ratio"].append(self.accelerator.gather(mean_clamped).nanmean().item())
+        # is_clamped_mask = rollout_imp_ratio > self.rollout_importance_sampling_cap
+        # mean_clamped = masked_batch_mean(is_clamped_mask.float())
+        # self._metrics[mode]["TIS/clamp_ratio"].append(self.accelerator.gather(mean_clamped).nanmean().item())
 
         gathered_low_clip = self.accelerator.gather(low_clip)
         self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item())

From a20b75cb7b478394725fb4ebd98d08957eef423e Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Mon, 25 Aug 2025 20:26:55 +0200
Subject: [PATCH 07/23] cleanup v2

---
 trl/scripts/vllm_serve.py   | 20 +-------------------
 trl/trainer/grpo_trainer.py | 10 ++++------
 2 files changed, 5 insertions(+), 25 deletions(-)

diff --git a/trl/scripts/vllm_serve.py b/trl/scripts/vllm_serve.py
index ecb10b5e77f..f5bf7b4712a 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -376,24 +376,6 @@ 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_float(logprob):
-    import math
-
-    value = logprob.logprob
-    """Replace inf, -inf, and nan with None for JSON compatibility."""
-    if math.isinf(value):
-        # print(f"Warning: Replacing {'positive' if value > 0 else 'negative'} infinity with None.")
-        # print(logprob)
-        exit()
-        return None
-    elif math.isnan(value):
-        # print("Warning: Replacing NaN with None.")
-        # print(logprob)
-        exit()
-        return None
-    return value
-
-
 def main(script_args: ScriptArguments):
     if not is_fastapi_available():
         raise ImportError(
@@ -577,7 +559,7 @@ async def generate(request: GenerateRequest):
 
         if request.logprobs is not None:
             logprobs: list[list[float]] = [
-                [sanitize_float(next(iter(lp.values()))) for lp in output.logprobs]
+                [next(iter(lp.values())) for lp in output.logprobs]
                 for outputs in all_outputs
                 for output in outputs.outputs
             ]
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 09683f72980..88fbf856c8e 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1415,6 +1415,8 @@ def _generate_and_score_completions(
                     re.sub(rf"({re.escape(self.image_token)})+", self.image_token, text) for text in prompts_text
                 ]
 
+        rollout_old_per_token_logps = None
+
         # Generate completions using either vLLM or regular generation
         if self.use_vllm:
             # First, update the vLLM weights if needed
@@ -1754,7 +1756,7 @@ def _generate_and_score_completions(
         if has_images:
             self._logs["image"].extend(gather_object(images))
 
-        if self.use_vllm and old_per_token_logps is not None:
+        if rollout_old_per_token_logps is not None and old_per_token_logps is not None:
             probs_diff = torch.exp(old_per_token_logps - rollout_old_per_token_logps)
             per_token_kl = probs_diff - (old_per_token_logps - rollout_old_per_token_logps) - 1
             mean_kl = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
@@ -1771,7 +1773,7 @@ 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:
+        if rollout_old_per_token_logps is not None:
             output["rollout_old_per_token_logps"] = rollout_old_per_token_logps
         if ref_per_token_logps is not None:
             output["ref_per_token_logps"] = ref_per_token_logps
@@ -1951,10 +1953,6 @@ def masked_batch_mean(x):
         high_clip = masked_batch_mean(is_high_clipped.float())
         clip_ratio = masked_batch_mean(is_region_clipped.float())
 
-        # is_clamped_mask = rollout_imp_ratio > self.rollout_importance_sampling_cap
-        # mean_clamped = masked_batch_mean(is_clamped_mask.float())
-        # self._metrics[mode]["TIS/clamp_ratio"].append(self.accelerator.gather(mean_clamped).nanmean().item())
-
         gathered_low_clip = self.accelerator.gather(low_clip)
         self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item())
         self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item())

From fd35ee739b60f2fe8175bc744c63693af720faa2 Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Mon, 25 Aug 2025 20:51:36 +0200
Subject: [PATCH 08/23] cleanup v3

---
 trl/scripts/vllm_serve.py   |  2 +-
 trl/trainer/grpo_trainer.py | 25 ++++++++++++++++++-------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/trl/scripts/vllm_serve.py b/trl/scripts/vllm_serve.py
index f5bf7b4712a..f7a0b0bc35b 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -534,6 +534,7 @@ async def generate(request: GenerateRequest):
 
         # Evenly distribute prompts across DP ranks
         chunked_prompts = chunk_list(prompts, script_args.data_parallel_size)
+
         # Send the prompts to each worker
         for connection, prompts in zip(connections, chunked_prompts):
             # When the number of prompts is less than data_parallel_size, some workers will receive empty prompts.
@@ -552,7 +553,6 @@ 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]
 
         response = {"completion_ids": completion_ids}
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 88fbf856c8e..c88a9f09013 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1415,9 +1415,8 @@ def _generate_and_score_completions(
                     re.sub(rf"({re.escape(self.image_token)})+", self.image_token, text) for text in prompts_text
                 ]
 
-        rollout_old_per_token_logps = None
-
         # Generate completions using either vLLM or regular generation
+        rollout_old_per_token_logps = None
         if self.use_vllm:
             # First, update the vLLM weights if needed
             if self.state.global_step != self._last_loaded_step:
@@ -1491,6 +1490,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,
                 }
                 if self.args.generation_kwargs is not None:
                     generation_kwargs.update(self.args.generation_kwargs)
@@ -1528,6 +1528,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]
+                logprobs = [
+                    [next(iter(lp.values())) 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.
@@ -1535,6 +1540,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]
+                    logprobs = logprobs[tp_slice]
 
             # Pad the completions, and concatenate them with the prompts
             completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids]
@@ -1757,12 +1763,17 @@ def _generate_and_score_completions(
             self._logs["image"].extend(gather_object(images))
 
         if rollout_old_per_token_logps is not None and old_per_token_logps is not None:
-            probs_diff = torch.exp(old_per_token_logps - rollout_old_per_token_logps)
-            per_token_kl = probs_diff - (old_per_token_logps - rollout_old_per_token_logps) - 1
+            delta = old_per_token_logps - rollout_old_per_token_logps
+            probs_diff = torch.exp(delta)
+            per_token_kl = probs_diff - delta - 1
             mean_kl = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
-            self._metrics[mode]["Rollout KL"].append(self.accelerator.gather(mean_kl).nanmean().item())
-            self._metrics[mode]["Token Probability Difference/max"].append(probs_diff.max().exp().item())
-            self._metrics[mode]["Token Probability Difference/mean"].append(probs_diff.mean().exp().item())
+            self._metrics[mode]["completions/kl_vllm"].append(self.accelerator.gather(mean_kl).nanmean().item())
+            self._metrics[mode]["completions/vllm_token_probability_difference/max"].append(
+                probs_diff.max().exp().item()
+            )
+            self._metrics[mode]["completions/vllm_token_probability_difference/mean"].append(
+                probs_diff.mean().exp().item()
+            )
 
         output = {
             "prompt_ids": prompt_ids,

From d6fde005b059c32ca215a7825ea2df3b1835bb97 Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Tue, 26 Aug 2025 11:42:09 +0200
Subject: [PATCH 09/23] added TIS as configurable option

---
 trl/trainer/grpo_config.py  | 141 +++++++++++++++++++++++++-----------
 trl/trainer/grpo_trainer.py |  13 ++--
 2 files changed, 104 insertions(+), 50 deletions(-)

diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py
index 6a96135aeb5..8e3a68fbe31 100644
--- a/trl/trainer/grpo_config.py
+++ b/trl/trainer/grpo_config.py
@@ -15,8 +15,6 @@
 from dataclasses import dataclass, field
 from typing import Optional, Union
 
-import transformers
-from packaging import version
 from transformers import TrainingArguments
 
 
@@ -96,7 +94,7 @@ class GRPOConfig(TrainingArguments):
         generation_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`):
             Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or `SamplingParams` (if
             using vLLM) when sampling completions. This can be used to further customize the generation behavior, such
-            as setting `supress_tokens`, `num_beams`, etc. If it contains keys that conflict with the other generation
+            as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict with the other generation
             parameters (like `min_p`, `top_p`, etc.), they will override them.
 
         > Parameters that control generation acceleration powered by vLLM
@@ -152,7 +150,7 @@ class GRPOConfig(TrainingArguments):
             Number of iterations per batch (denoted as μ in the algorithm).
         epsilon (`float`, *optional*, defaults to `0.2`):
             Epsilon value for clipping.
-        delta: (`float` or `None`, *optional*, defaults to `None`):
+        delta (`float` or `None`, *optional*, defaults to `None`):
             Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` (default), standard
             GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This method is introduced in
             the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291).
@@ -161,31 +159,39 @@ class GRPOConfig(TrainingArguments):
             specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`.
         importance_sampling_level (`str`, *optional*, defaults to `"token"`):
             Controls whether importance sampling ratios are computed at the `"token"` or `"sequence"` level. `"token"`
-            keeps the raw per-token log-probability ratios (one weight per token).  `"sequence"` averages the
-            log-probability ratios across valid tokens to produce a single ratio per sequence. The
-            [GSPO paper](https://huggingface.co/papers/2507.18071) shows that sequence-level sampling often yields more
-            stable training and better alignment with  sequence-level rewards.
+            keeps the raw per-token log-probability ratios (one weight per token). `"sequence"` averages the
+            log-probability ratios across valid tokens to produce a single ratio per sequence. The [GSPO
+            paper](https://huggingface.co/papers/2507.18071) shows that sequence-level sampling often yields more
+            stable training and better alignment with sequence-level rewards.
         reward_weights (`list[float]` or `None`, *optional*, defaults to `None`):
             Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are
             weighted equally with weight `1.0`.
-        scale_rewards (`bool`, *optional*, defaults to `True`):
-            Whether to scale the rewards by dividing them by their standard deviation. If `True` (default), the rewards
-            are normalized by the standard deviation, ensuring they have unit variance. If `False`, no scaling is
-            applied. The [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) recommends not scaling the rewards,
-            as scaling by the standard deviation introduces a question-level difficulty bias.
-        loss_type (`str`, *optional*, defaults to `"bnpo"`):
+        scale_rewards (`str` or `bool`, *optional*, defaults to `"group"`):
+            Specifies the scaling strategy for rewards. Supported values are:
+
+            - `True` or `"group"` (default): rewards are scaled by the standard deviation within each group, ensuring
+              unit variance within a group.
+            - `"batch"`: rewards are scaled by the standard deviation across the entire batch, as recommended in the
+              [PPO Lite paper](https://huggingface.co/papers/2508.08221).
+            - `False` or `"none"`: no scaling is applied. The [Dr. GRPO
+              paper](https://huggingface.co/papers/2503.20783) recommends not scaling rewards, as scaling by the
+              standard deviation introduces a question-level difficulty bias.
+        loss_type (`str`, *optional*, defaults to `"dapo"`):
             Specifies the loss formulation to use. Supported values are:
 
             - `"grpo"`: Aggregates token-level losses by normalizing over sequence length. Not recommended due to
-                length bias—this approach tends to prefer shorter completions with positive advantages and longer ones
-                with negative advantages.
-            - `"bnpo"`: Aggregates token-level losses by normalizing number of active token in the local batch.
-                Note that normalization is performed over the local batch only, so results may slightly vary depending
-                on the local batch size, despite a constant effective batch size. When using
-                `per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss.
+              length bias—this approach tends to prefer shorter completions with positive advantages and longer ones
+              with negative advantages.
             - `"dr_grpo"`: Aggregates token-level losses by normalizing with a global constant. This method was
-                introduced in the [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) to eliminate length bias.
-                The value of the constant corresponds to `max_completion_length`.
+              introduced in the [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) to eliminate length bias.
+              The value of the constant corresponds to `max_completion_length`.
+            - `"dapo"` (default): Aggregates token-level losses by normalizing with the number of active token in the
+              global accumulated batch. This method was introduced in the [DAPO
+              paper](https://huggingface.co/papers/2503.14476) to eliminate length bias.
+            - `"bnpo"`: Aggregates token-level losses by normalizing with the number of active token in the local
+              batch. Note that normalization is performed over the local batch only, so results may slightly vary
+              depending on the local batch size, despite a constant effective batch size. When using
+              `per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss.
         mask_truncated_completions (`bool`, *optional*, defaults to `False`):
             When enabled, truncated completions are excluded from the loss calculation, preventing them from being
             incorrectly penalized and introducing noise during training. According to the
@@ -207,10 +213,18 @@ class GRPOConfig(TrainingArguments):
             ρ parameter from [Beyond the 80/20 Rule](https://huggingface.co/papers/2506.01939). Keeps in the policy
             loss term only the top-ρ quantile of tokens by entropy of the probability distribution at each sequence
             position, improving results. Range: `[0.0-1.0]`. A value of `0.0` masks all but the highest entropy token;
-            `1.0` keeps all tokens. The paper recommends a value of `0.2`.
-            If used with `mask_truncated_completions=True`, only tokens from non-truncated completions are considered.
+            `1.0` keeps all tokens. The paper recommends a value of `0.2`. If used with
+            `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 `False`):
+            Whether to apply truncated importance sampling 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. Truncated Importance Sampling is proposed as a
+            remedy for this issue.
+        vllm_importance_sampling_cap (`float`, *optional*, defaults to `2.0`):
+            The truncation parameter C for Truncated Importance Sampling. This sets an upper bound on the importance sampling ratio,
+            improving training stability.
 
         > Parameters that control the logging
 
@@ -224,8 +238,7 @@ class GRPOConfig(TrainingArguments):
             are logged.
     """
 
-    if version.parse(transformers.__version__) >= version.parse("4.51.0"):
-        _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"]
+    _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"]
 
     # Parameters whose default values are overridden from TrainingArguments
     learning_rate: float = field(
@@ -239,6 +252,12 @@ class GRPOConfig(TrainingArguments):
             "will be interpreted as ratio of total training steps."
         },
     )
+    gradient_checkpointing: bool = field(
+        default=True,
+        metadata={
+            "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass."
+        },
+    )
     bf16: Optional[bool] = field(
         default=None,
         metadata={
@@ -347,7 +366,7 @@ class GRPOConfig(TrainingArguments):
         metadata={
             "help": "Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or "
             "`SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the "
-            "generation behavior, such as setting `supress_tokens`, `num_beams`, etc. If it contains keys that "
+            "generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that "
             "conflict with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them."
         },
     )
@@ -493,29 +512,35 @@ class GRPOConfig(TrainingArguments):
             "rewards are weighted equally with weight `1.0`."
         },
     )
-    scale_rewards: bool = field(
-        default=True,
+    scale_rewards: str = field(
+        default="group",
         metadata={
-            "help": "Whether to scale the rewards by dividing them by their standard deviation. If `True` (default), "
-            "the rewards are normalized by the standard deviation, ensuring they have unit variance. If `False`, no "
-            "scaling is applied. The Dr. GRPO paper recommends not scaling the rewards, as scaling by the standard "
-            "deviation introduces a question-level difficulty bias."
+            "help": "Specifies the scaling strategy for rewards. Supported values are: "
+            "`True` or `group'` (default): rewards are scaled by the standard deviation within each group, ensuring "
+            "unit variance within a group. "
+            "`'batch'`: rewards are scaled by the standard deviation across the entire batch, as recommended in the "
+            "PPO Lite paper. "
+            "`False` or `'none'`: no scaling is applied. The Dr. GRPO paper recommends not scaling rewards, as "
+            "scaling by the standard deviation introduces a question-level difficulty bias."
         },
     )
     loss_type: str = field(
-        default="bnpo",
-        metadata={
-            "help": "Specifies the loss formulation to use. Supported values are `grpo`, `bnpo`, and `dr_grpo`. "
-            "`'grpo'`: Aggregates token-level losses by normalizing over sequence length. Not recommended due to "
-            "length bias—this approach tends to prefer shorter completions with positive advantages and longer ones "
-            "with negative advantages. "
-            "`'bnpo'`: Aggregates token-level losses by normalizing number of active token in the local batch. "
+        default="dapo",
+        metadata={
+            "help": "Specifies the loss formulation to use. Supported values are 'grpo', 'dapo', 'bnpo', and "
+            "'dr_grpo'. "
+            "'grpo': Aggregates token-level losses by normalizing over sequence length. Not recommended due to length "
+            "bias—this approach tends to prefer shorter completions with positive advantages and longer ones with "
+            "negative advantages. "
+            "'dapo' (default): Aggregates token-level losses by normalizing with the number of active token in the "
+            "global accumulated batch. This method was introduced in the DAPO paper to eliminate length bias. "
+            "'dr_grpo': Aggregates token-level losses by normalizing with a global constant. This method was "
+            "introduced in the Dr. GRPO paper to eliminate length bias. The value of the constant corresponds to "
+            "`max_completion_length`. "
+            "'bnpo': Aggregates token-level losses by normalizing with the number of active token in the local batch. "
             "Note that normalization is performed over the local batch only, so results may slightly vary depending "
             "on the local batch size, despite a constant effective batch size. When using "
-            "`per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss. "
-            "`'dr_grpo'`: Aggregates token-level losses by normalizing with a global constant. This method was "
-            "introduced in the Dr. GRPO paper to eliminate length bias. The value of the constant corresponds to "
-            "`max_completion_length`."
+            "`per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss."
         },
     )
     mask_truncated_completions: bool = field(
@@ -562,6 +587,23 @@ class GRPOConfig(TrainingArguments):
         default=False,
         metadata={"help": "Whether to use the Liger GRPO loss."},
     )
+    vllm_importance_sampling_correction: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to apply truncated importance sampling 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. Truncated Importance Sampling is proposed as a "
+            "remedy for this issue."
+        },
+    )
+    vllm_importance_sampling_cap: float = field(
+        default=False,
+        metadata={
+            "help": "The truncation parameter C for Truncated Importance Sampling. This sets an upper bound on the "
+            "importance sampling ratio, improving training stability."
+        },
+    )
 
     # Parameters that control the logging
     log_completions: bool = field(
@@ -610,6 +652,14 @@ def __post_init__(self):
                 "'generation_batch_size' and 'steps_per_generation' can not be both configured at the same time"
             )
 
+        if self.do_eval and self.eval_strategy != "no":
+            # Just ensure the value is divisible by the global batch size
+            if (self.per_device_eval_batch_size * num_processes) % self.num_generations != 0:
+                raise ValueError(
+                    f"The global eval batch size ({self.per_device_eval_batch_size} * {num_processes}) must be "
+                    f"divisible by num_generations ({self.num_generations})."
+                )
+
         # The generation batch must contain full prompt groups (no partials), so it must be divisible by
         # num_generations.
         if self.generation_batch_size % self.num_generations != 0:
@@ -626,3 +676,6 @@ def __post_init__(self):
 
         if self.delta is not None and self.use_liger_loss:
             raise ValueError("Liger loss does not support two-sided GRPO loss yet.")
+
+        if not self.use_vllm and self.vllm_importance_sampling_correction:
+            raise ValueError("'vllm_importance_sampling_correction' requires vLLM generation. Enable with 'use_vllm'.")
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index c88a9f09013..7f10280509c 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -654,7 +654,6 @@ def __init__(
         self.reward_processing_classes = reward_processing_classes
 
         # Training arguments
-        self.rollout_importance_sampling_cap = -1
         self.max_prompt_length = args.max_prompt_length
         self.max_completion_length = args.max_completion_length  # = |o_i| in the GRPO paper
         self.num_generations = args.num_generations  # = G in the GRPO paper
@@ -674,6 +673,8 @@ def __init__(
         self.importance_sampling_level = args.importance_sampling_level
         self.mask_truncated_completions = args.mask_truncated_completions
         self.top_entropy_quantile = args.top_entropy_quantile
+        self.vllm_importance_sampling_correction = args.vllm_importance_sampling_correction
+        self.vllm_importance_sampling_cap = args.vllm_importance_sampling_cap
         if self.use_liger_loss and self.top_entropy_quantile < 1.0:
             raise NotImplementedError(
                 "Liger Kernels don't currently support masking token positions based on entropy."
@@ -1453,7 +1454,7 @@ def _generate_and_score_completions(
                             max_tokens=self.max_completion_length,
                             guided_decoding_regex=self.guided_decoding_regex,
                             generation_kwargs=self.args.generation_kwargs,
-                            logprobs=0,
+                            logprobs=0 if self.vllm_importance_sampling_correction else None,
                         )
                         payload = (
                             output["completion_ids"],
@@ -1490,7 +1491,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,
+                    "logprobs": 0 if self.vllm_importance_sampling_correction else None,
                 }
                 if self.args.generation_kwargs is not None:
                     generation_kwargs.update(self.args.generation_kwargs)
@@ -1762,7 +1763,7 @@ def _generate_and_score_completions(
         if has_images:
             self._logs["image"].extend(gather_object(images))
 
-        if rollout_old_per_token_logps is not None and old_per_token_logps is not None:
+        if self.vllm_importance_sampling_correction and old_per_token_logps is not None:
             delta = old_per_token_logps - rollout_old_per_token_logps
             probs_diff = torch.exp(delta)
             per_token_kl = probs_diff - delta - 1
@@ -1784,7 +1785,7 @@ def _generate_and_score_completions(
         }
         if old_per_token_logps is not None:
             output["old_per_token_logps"] = old_per_token_logps
-        if rollout_old_per_token_logps is not None:
+        if self.vllm_importance_sampling_correction:
             output["rollout_old_per_token_logps"] = rollout_old_per_token_logps
         if ref_per_token_logps is not None:
             output["ref_per_token_logps"] = ref_per_token_logps
@@ -1919,7 +1920,7 @@ def _compute_loss(self, model, inputs):
         if entropy_mask is not None:
             per_token_loss = per_token_loss * entropy_mask
 
-        if self.use_vllm and self.rollout_importance_sampling_cap > 0 and old_per_token_logps is not None:
+        if self.vllm_importance_sampling_correction and old_per_token_logps is not None:
             rollout_old_per_token_logps = inputs.get("rollout_old_per_token_logps")
             rollout_imp_ratio = torch.exp(old_per_token_logps - rollout_old_per_token_logps)
             clammped_rollout_imp_ratio = torch.clamp(rollout_imp_ratio, max=self.rollout_importance_sampling_cap)

From 49841ea3069a870689be543007fc68852f34e9ce Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Tue, 26 Aug 2025 12:30:15 +0200
Subject: [PATCH 10/23] cleanup

---
 trl/trainer/grpo_config.py | 119 ++++++++++++++-----------------------
 1 file changed, 46 insertions(+), 73 deletions(-)

diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py
index 8e3a68fbe31..c4872237d0e 100644
--- a/trl/trainer/grpo_config.py
+++ b/trl/trainer/grpo_config.py
@@ -15,6 +15,8 @@
 from dataclasses import dataclass, field
 from typing import Optional, Union
 
+import transformers
+from packaging import version
 from transformers import TrainingArguments
 
 
@@ -94,7 +96,7 @@ class GRPOConfig(TrainingArguments):
         generation_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`):
             Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or `SamplingParams` (if
             using vLLM) when sampling completions. This can be used to further customize the generation behavior, such
-            as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict with the other generation
+            as setting `supress_tokens`, `num_beams`, etc. If it contains keys that conflict with the other generation
             parameters (like `min_p`, `top_p`, etc.), they will override them.
 
         > Parameters that control generation acceleration powered by vLLM
@@ -150,7 +152,7 @@ class GRPOConfig(TrainingArguments):
             Number of iterations per batch (denoted as μ in the algorithm).
         epsilon (`float`, *optional*, defaults to `0.2`):
             Epsilon value for clipping.
-        delta (`float` or `None`, *optional*, defaults to `None`):
+        delta: (`float` or `None`, *optional*, defaults to `None`):
             Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` (default), standard
             GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This method is introduced in
             the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291).
@@ -159,39 +161,31 @@ class GRPOConfig(TrainingArguments):
             specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`.
         importance_sampling_level (`str`, *optional*, defaults to `"token"`):
             Controls whether importance sampling ratios are computed at the `"token"` or `"sequence"` level. `"token"`
-            keeps the raw per-token log-probability ratios (one weight per token). `"sequence"` averages the
-            log-probability ratios across valid tokens to produce a single ratio per sequence. The [GSPO
-            paper](https://huggingface.co/papers/2507.18071) shows that sequence-level sampling often yields more
-            stable training and better alignment with sequence-level rewards.
+            keeps the raw per-token log-probability ratios (one weight per token).  `"sequence"` averages the
+            log-probability ratios across valid tokens to produce a single ratio per sequence. The
+            [GSPO paper](https://huggingface.co/papers/2507.18071) shows that sequence-level sampling often yields more
+            stable training and better alignment with  sequence-level rewards.
         reward_weights (`list[float]` or `None`, *optional*, defaults to `None`):
             Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are
             weighted equally with weight `1.0`.
-        scale_rewards (`str` or `bool`, *optional*, defaults to `"group"`):
-            Specifies the scaling strategy for rewards. Supported values are:
-
-            - `True` or `"group"` (default): rewards are scaled by the standard deviation within each group, ensuring
-              unit variance within a group.
-            - `"batch"`: rewards are scaled by the standard deviation across the entire batch, as recommended in the
-              [PPO Lite paper](https://huggingface.co/papers/2508.08221).
-            - `False` or `"none"`: no scaling is applied. The [Dr. GRPO
-              paper](https://huggingface.co/papers/2503.20783) recommends not scaling rewards, as scaling by the
-              standard deviation introduces a question-level difficulty bias.
-        loss_type (`str`, *optional*, defaults to `"dapo"`):
+        scale_rewards (`bool`, *optional*, defaults to `True`):
+            Whether to scale the rewards by dividing them by their standard deviation. If `True` (default), the rewards
+            are normalized by the standard deviation, ensuring they have unit variance. If `False`, no scaling is
+            applied. The [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) recommends not scaling the rewards,
+            as scaling by the standard deviation introduces a question-level difficulty bias.
+        loss_type (`str`, *optional*, defaults to `"bnpo"`):
             Specifies the loss formulation to use. Supported values are:
 
             - `"grpo"`: Aggregates token-level losses by normalizing over sequence length. Not recommended due to
-              length bias—this approach tends to prefer shorter completions with positive advantages and longer ones
-              with negative advantages.
+                length bias—this approach tends to prefer shorter completions with positive advantages and longer ones
+                with negative advantages.
+            - `"bnpo"`: Aggregates token-level losses by normalizing number of active token in the local batch.
+                Note that normalization is performed over the local batch only, so results may slightly vary depending
+                on the local batch size, despite a constant effective batch size. When using
+                `per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss.
             - `"dr_grpo"`: Aggregates token-level losses by normalizing with a global constant. This method was
-              introduced in the [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) to eliminate length bias.
-              The value of the constant corresponds to `max_completion_length`.
-            - `"dapo"` (default): Aggregates token-level losses by normalizing with the number of active token in the
-              global accumulated batch. This method was introduced in the [DAPO
-              paper](https://huggingface.co/papers/2503.14476) to eliminate length bias.
-            - `"bnpo"`: Aggregates token-level losses by normalizing with the number of active token in the local
-              batch. Note that normalization is performed over the local batch only, so results may slightly vary
-              depending on the local batch size, despite a constant effective batch size. When using
-              `per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss.
+                introduced in the [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) to eliminate length bias.
+                The value of the constant corresponds to `max_completion_length`.
         mask_truncated_completions (`bool`, *optional*, defaults to `False`):
             When enabled, truncated completions are excluded from the loss calculation, preventing them from being
             incorrectly penalized and introducing noise during training. According to the
@@ -213,12 +207,13 @@ class GRPOConfig(TrainingArguments):
             ρ parameter from [Beyond the 80/20 Rule](https://huggingface.co/papers/2506.01939). Keeps in the policy
             loss term only the top-ρ quantile of tokens by entropy of the probability distribution at each sequence
             position, improving results. Range: `[0.0-1.0]`. A value of `0.0` masks all but the highest entropy token;
-            `1.0` keeps all tokens. The paper recommends a value of `0.2`. If used with
-            `mask_truncated_completions=True`, only tokens from non-truncated completions are considered.
+            `1.0` keeps all tokens. The paper recommends a value of `0.2`.
+            If used with `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 `False`):
-            Whether to apply truncated importance sampling 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)
+            Whether to apply truncated importance sampling 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. Truncated Importance Sampling is proposed as a
             remedy for this issue.
@@ -238,7 +233,8 @@ class GRPOConfig(TrainingArguments):
             are logged.
     """
 
-    _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"]
+    if version.parse(transformers.__version__) >= version.parse("4.51.0"):
+        _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"]
 
     # Parameters whose default values are overridden from TrainingArguments
     learning_rate: float = field(
@@ -252,12 +248,6 @@ class GRPOConfig(TrainingArguments):
             "will be interpreted as ratio of total training steps."
         },
     )
-    gradient_checkpointing: bool = field(
-        default=True,
-        metadata={
-            "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass."
-        },
-    )
     bf16: Optional[bool] = field(
         default=None,
         metadata={
@@ -366,7 +356,7 @@ class GRPOConfig(TrainingArguments):
         metadata={
             "help": "Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or "
             "`SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the "
-            "generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that "
+            "generation behavior, such as setting `supress_tokens`, `num_beams`, etc. If it contains keys that "
             "conflict with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them."
         },
     )
@@ -512,35 +502,29 @@ class GRPOConfig(TrainingArguments):
             "rewards are weighted equally with weight `1.0`."
         },
     )
-    scale_rewards: str = field(
-        default="group",
+    scale_rewards: bool = field(
+        default=True,
         metadata={
-            "help": "Specifies the scaling strategy for rewards. Supported values are: "
-            "`True` or `group'` (default): rewards are scaled by the standard deviation within each group, ensuring "
-            "unit variance within a group. "
-            "`'batch'`: rewards are scaled by the standard deviation across the entire batch, as recommended in the "
-            "PPO Lite paper. "
-            "`False` or `'none'`: no scaling is applied. The Dr. GRPO paper recommends not scaling rewards, as "
-            "scaling by the standard deviation introduces a question-level difficulty bias."
+            "help": "Whether to scale the rewards by dividing them by their standard deviation. If `True` (default), "
+            "the rewards are normalized by the standard deviation, ensuring they have unit variance. If `False`, no "
+            "scaling is applied. The Dr. GRPO paper recommends not scaling the rewards, as scaling by the standard "
+            "deviation introduces a question-level difficulty bias."
         },
     )
     loss_type: str = field(
-        default="dapo",
-        metadata={
-            "help": "Specifies the loss formulation to use. Supported values are 'grpo', 'dapo', 'bnpo', and "
-            "'dr_grpo'. "
-            "'grpo': Aggregates token-level losses by normalizing over sequence length. Not recommended due to length "
-            "bias—this approach tends to prefer shorter completions with positive advantages and longer ones with "
-            "negative advantages. "
-            "'dapo' (default): Aggregates token-level losses by normalizing with the number of active token in the "
-            "global accumulated batch. This method was introduced in the DAPO paper to eliminate length bias. "
-            "'dr_grpo': Aggregates token-level losses by normalizing with a global constant. This method was "
-            "introduced in the Dr. GRPO paper to eliminate length bias. The value of the constant corresponds to "
-            "`max_completion_length`. "
-            "'bnpo': Aggregates token-level losses by normalizing with the number of active token in the local batch. "
+        default="bnpo",
+        metadata={
+            "help": "Specifies the loss formulation to use. Supported values are `grpo`, `bnpo`, and `dr_grpo`. "
+            "`'grpo'`: Aggregates token-level losses by normalizing over sequence length. Not recommended due to "
+            "length bias—this approach tends to prefer shorter completions with positive advantages and longer ones "
+            "with negative advantages. "
+            "`'bnpo'`: Aggregates token-level losses by normalizing number of active token in the local batch. "
             "Note that normalization is performed over the local batch only, so results may slightly vary depending "
             "on the local batch size, despite a constant effective batch size. When using "
-            "`per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss."
+            "`per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss. "
+            "`'dr_grpo'`: Aggregates token-level losses by normalizing with a global constant. This method was "
+            "introduced in the Dr. GRPO paper to eliminate length bias. The value of the constant corresponds to "
+            "`max_completion_length`."
         },
     )
     mask_truncated_completions: bool = field(
@@ -652,14 +636,6 @@ def __post_init__(self):
                 "'generation_batch_size' and 'steps_per_generation' can not be both configured at the same time"
             )
 
-        if self.do_eval and self.eval_strategy != "no":
-            # Just ensure the value is divisible by the global batch size
-            if (self.per_device_eval_batch_size * num_processes) % self.num_generations != 0:
-                raise ValueError(
-                    f"The global eval batch size ({self.per_device_eval_batch_size} * {num_processes}) must be "
-                    f"divisible by num_generations ({self.num_generations})."
-                )
-
         # The generation batch must contain full prompt groups (no partials), so it must be divisible by
         # num_generations.
         if self.generation_batch_size % self.num_generations != 0:
@@ -676,6 +652,3 @@ def __post_init__(self):
 
         if self.delta is not None and self.use_liger_loss:
             raise ValueError("Liger loss does not support two-sided GRPO loss yet.")
-
-        if not self.use_vllm and self.vllm_importance_sampling_correction:
-            raise ValueError("'vllm_importance_sampling_correction' requires vLLM generation. Enable with 'use_vllm'.")

From e1630eb5e6337b3445dba20afdf03e7e97b29c46 Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Tue, 26 Aug 2025 12:41:26 +0200
Subject: [PATCH 11/23] test

---
 tests/test_grpo_trainer.py | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py
index 449f3e94758..d8e4880c9ef 100644
--- a/tests/test_grpo_trainer.py
+++ b/tests/test_grpo_trainer.py
@@ -1207,6 +1207,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")

From 01a208d069b493ef22de396c641b7a245208bcd6 Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Wed, 27 Aug 2025 10:00:41 +0200
Subject: [PATCH 12/23] fix logs and rename

---
 trl/trainer/grpo_config.py  |  2 +-
 trl/trainer/grpo_trainer.py | 24 +++++++++---------------
 2 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py
index 94e043a9f0f..5c06847abf5 100644
--- a/trl/trainer/grpo_config.py
+++ b/trl/trainer/grpo_config.py
@@ -599,7 +599,7 @@ class GRPOConfig(TrainingArguments):
         },
     )
     vllm_importance_sampling_cap: float = field(
-        default=False,
+        default=2.0,
         metadata={
             "help": "The truncation parameter C for Truncated Importance Sampling. This sets an upper bound on the "
             "importance sampling ratio, improving training stability."
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 1ccc5fd1a95..a9f79f906df 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1427,7 +1427,6 @@ def _generate_and_score_completions(
                         prompts_text = [re.sub(rf"({escaped_img_token})+", "", text) for text in prompts_text]
 
         # Generate completions using either vLLM or regular generation
-        rollout_old_per_token_logps = None
         if self.use_vllm:
             # First, update the vLLM weights if needed
             if self.state.global_step != self._last_loaded_step:
@@ -1557,8 +1556,8 @@ 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)
-            rollout_old_per_token_logps = [torch.tensor(logp, device=device) for logp in logprobs]
-            rollout_old_per_token_logps = pad(rollout_old_per_token_logps, padding_value=0.0)
+            vllm_old_per_token_logps = [torch.tensor(logp, device=device) for logp in logprobs]
+            vllm_old_per_token_logps = pad(vllm_old_per_token_logps, padding_value=0.0)
 
         elif self.use_transformers_paged:
             # Re-process inputs for paged generation if needed
@@ -1784,17 +1783,12 @@ def _generate_and_score_completions(
             self._logs["image"].extend(gather_object(images))
 
         if self.vllm_importance_sampling_correction and old_per_token_logps is not None:
-            delta = old_per_token_logps - rollout_old_per_token_logps
+            delta = old_per_token_logps - vllm_old_per_token_logps
             probs_diff = torch.exp(delta)
             per_token_kl = probs_diff - delta - 1
             mean_kl = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
             self._metrics[mode]["completions/kl_vllm"].append(self.accelerator.gather(mean_kl).nanmean().item())
-            self._metrics[mode]["completions/vllm_token_probability_difference/max"].append(
-                probs_diff.max().exp().item()
-            )
-            self._metrics[mode]["completions/vllm_token_probability_difference/mean"].append(
-                probs_diff.mean().exp().item()
-            )
+            self._metrics[mode]["completions/vllm_token_probability_difference/max"].append(probs_diff.max().item())
 
         output = {
             "prompt_ids": prompt_ids,
@@ -1807,7 +1801,7 @@ def _generate_and_score_completions(
         if old_per_token_logps is not None:
             output["old_per_token_logps"] = old_per_token_logps
         if self.vllm_importance_sampling_correction:
-            output["rollout_old_per_token_logps"] = rollout_old_per_token_logps
+            output["vllm_old_per_token_logps"] = vllm_old_per_token_logps
         if ref_per_token_logps is not None:
             output["ref_per_token_logps"] = ref_per_token_logps
         if "pixel_values" in prompt_inputs:
@@ -1942,10 +1936,10 @@ def _compute_loss(self, model, inputs):
             per_token_loss = per_token_loss * entropy_mask
 
         if self.vllm_importance_sampling_correction and old_per_token_logps is not None:
-            rollout_old_per_token_logps = inputs.get("rollout_old_per_token_logps")
-            rollout_imp_ratio = torch.exp(old_per_token_logps - rollout_old_per_token_logps)
-            clammped_rollout_imp_ratio = torch.clamp(rollout_imp_ratio, max=self.rollout_importance_sampling_cap)
-            per_token_loss = per_token_loss * clammped_rollout_imp_ratio
+            vllm_old_per_token_logps = inputs.get("vllm_old_per_token_logps")
+            rollout_imp_ratio = torch.exp(old_per_token_logps - vllm_old_per_token_logps)
+            clamped_rollout_imp_ratio = torch.clamp(rollout_imp_ratio, max=self.vllm_importance_sampling_cap)
+            per_token_loss = per_token_loss * clamped_rollout_imp_ratio
 
         if self.beta != 0.0:
             per_token_loss = per_token_loss + self.beta * per_token_kl

From 2ad229ff815e1523de95f7e21390ca1a1acc34b8 Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Wed, 27 Aug 2025 18:06:40 +0200
Subject: [PATCH 13/23] comment

---
 trl/trainer/grpo_trainer.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index a9f79f906df..2e35f0fe5c4 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1463,7 +1463,9 @@ def _generate_and_score_completions(
                             max_tokens=self.max_completion_length,
                             guided_decoding_regex=self.guided_decoding_regex,
                             generation_kwargs=self.args.generation_kwargs,
-                            logprobs=0 if self.vllm_importance_sampling_correction else None,
+                            logprobs=0
+                            if self.vllm_importance_sampling_correction
+                            else None,  # 0 represents only returning the selected token logprob
                         )
                         payload = (
                             output["completion_ids"],
@@ -1500,7 +1502,9 @@ 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 if self.vllm_importance_sampling_correction else None,
+                    "logprobs": 0
+                    if self.vllm_importance_sampling_correction
+                    else None,  # 0 represents only returning the selected token logprob
                 }
                 if self.args.generation_kwargs is not None:
                     generation_kwargs.update(self.args.generation_kwargs)

From 96baa2e4c99688f15685abef91d64399a9022baf Mon Sep 17 00:00:00 2001
From: Leon Ericsson <leon.ericsson@icloud.com>
Date: Wed, 27 Aug 2025 19:12:18 +0200
Subject: [PATCH 14/23] always return logprobs from vllm; sanitize logprob

---
 trl/scripts/vllm_serve.py   | 11 ++++++++++-
 trl/trainer/grpo_trainer.py |  8 ++------
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/trl/scripts/vllm_serve.py b/trl/scripts/vllm_serve.py
index 968164319bb..4d1b9032803 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -377,6 +377,15 @@ 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("Generated NaN logprob, token logprob will be ignored")
+        return float("-inf")
+
+
 def main(script_args: ScriptArguments):
     if not is_fastapi_available():
         raise ImportError(
@@ -569,7 +578,7 @@ async def generate(request: GenerateRequest):
 
         if request.logprobs is not None:
             logprobs: list[list[float]] = [
-                [next(iter(lp.values())) for lp in output.logprobs]
+                [sanitize_logprob(next(iter(lp.values()))) for lp in output.logprobs]
                 for outputs in all_outputs
                 for output in outputs.outputs
             ]
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 2e35f0fe5c4..e62b39b2bd1 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1463,9 +1463,7 @@ def _generate_and_score_completions(
                             max_tokens=self.max_completion_length,
                             guided_decoding_regex=self.guided_decoding_regex,
                             generation_kwargs=self.args.generation_kwargs,
-                            logprobs=0
-                            if self.vllm_importance_sampling_correction
-                            else None,  # 0 represents only returning the selected token logprob
+                            logprobs=0,  # only return the logprob of the generated token
                         )
                         payload = (
                             output["completion_ids"],
@@ -1502,9 +1500,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
-                    if self.vllm_importance_sampling_correction
-                    else None,  # 0 represents only returning the selected token logprob
+                    "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)

From 4d10db0730a3d1ee904c1cd9432ac8e2b2165bf5 Mon Sep 17 00:00:00 2001
From: Leon Ericsson <leon.ericsson@icloud.com>
Date: Wed, 27 Aug 2025 19:13:28 +0200
Subject: [PATCH 15/23] nit

---
 trl/scripts/vllm_serve.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/trl/scripts/vllm_serve.py b/trl/scripts/vllm_serve.py
index 4d1b9032803..7d590e45cb1 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -382,7 +382,7 @@ def sanitize_logprob(logprob):
 
     value = logprob.logprob
     if math.isnan(value):
-        logger.warning("Generated NaN logprob, token logprob will be ignored")
+        logger.warning(f"Generated NaN logprob, token logprob '{logprob}' will be ignored")
         return float("-inf")
 
 

From 87c8e8e174336da1b5d2b41b3ce200d04a41c371 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Wed, 27 Aug 2025 20:10:21 +0000
Subject: [PATCH 16/23] Enhance VLLM integration: return logprobs in responses
 and update importance sampling configuration

---
 trl/extras/vllm_client.py   | 12 +++---
 trl/scripts/vllm_serve.py   | 26 +++++--------
 trl/trainer/grpo_config.py  | 30 +++++++--------
 trl/trainer/grpo_trainer.py | 77 +++++++++++++++++++++++++------------
 4 files changed, 83 insertions(+), 62 deletions(-)

diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py
index 1564a96a413..aa2238c0eba 100644
--- a/trl/extras/vllm_client.py
+++ b/trl/extras/vllm_client.py
@@ -178,7 +178,6 @@ def generate(
         max_tokens: int = 16,
         guided_decoding_regex: Optional[str] = None,
         generation_kwargs: Optional[dict] = None,
-        logprobs: Optional[int] = None,
     ) -> list[list[int]]:
         """
         Generates model completions for the provided prompts.
@@ -210,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/"
 
@@ -238,13 +240,11 @@ def pil_to_base64(image):
                 "max_tokens": max_tokens,
                 "guided_decoding_regex": guided_decoding_regex,
                 "generation_kwargs": generation_kwargs or {},
-                "logprobs": logprobs,
             },
         )
         if response.status_code == 200:
             json_response = response.json()
-            result = {k: v for k, v in json_response.items() if k in ("completion_ids", "logprobs")}
-            return result
+            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 7d590e45cb1..604c84b64a9 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -474,7 +474,6 @@ class GenerateRequest(BaseModel):
         max_tokens: int = 16
         guided_decoding_regex: Optional[str] = None
         generation_kwargs: dict = field(default_factory=dict)
-        logprobs: Optional[int] = None
 
     class GenerateResponse(BaseModel):
         completion_ids: list[list[int]]
@@ -511,6 +510,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
@@ -519,7 +520,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)
@@ -546,7 +547,7 @@ async def generate(request: GenerateRequest):
             "min_p": request.min_p,
             "max_tokens": request.max_tokens,
             "guided_decoding": guided_decoding,
-            "logprobs": request.logprobs,
+            "logprobs": 0,
         }
         generation_kwargs.update(request.generation_kwargs)
         sampling_params = SamplingParams(**generation_kwargs)
@@ -573,19 +574,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]
-
-        response = {"completion_ids": completion_ids}
-
-        if request.logprobs is not None:
-            logprobs: list[list[float]] = [
-                [sanitize_logprob(next(iter(lp.values()))) for lp in output.logprobs]
-                for outputs in all_outputs
-                for output in outputs.outputs
-            ]
-
-            response["logprobs"] = logprobs
-
-        return response
+        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 5c06847abf5..d0976619ef4 100644
--- a/trl/trainer/grpo_config.py
+++ b/trl/trainer/grpo_config.py
@@ -217,15 +217,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 `False`):
-            Whether to apply truncated importance sampling 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. Truncated Importance Sampling is proposed as a
-            remedy for this issue.
+        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`):
-            The truncation parameter C for Truncated Importance Sampling. This sets an upper bound on the importance sampling ratio,
-            improving training stability.
+            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,19 +589,19 @@ class GRPOConfig(TrainingArguments):
         metadata={"help": "Whether to use the Liger GRPO loss."},
     )
     vllm_importance_sampling_correction: bool = field(
-        default=False,
+        default=True,
         metadata={
-            "help": "Whether to apply truncated importance sampling 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. Truncated Importance Sampling is proposed as a "
-            "remedy for this issue."
+            "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": "The truncation parameter C for Truncated Importance Sampling. This sets an upper bound on the "
+            "help": "Truncation parameter C for Truncated Importance Sampling (TIS). This sets an upper bound on the "
             "importance sampling ratio, improving training stability."
         },
     )
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index e62b39b2bd1..1be22c4dee7 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -649,14 +649,14 @@ 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 = {True: "group", False: "none"}.get(args.scale_rewards, args.scale_rewards)
         self.importance_sampling_level = args.importance_sampling_level
         self.mask_truncated_completions = args.mask_truncated_completions
         self.top_entropy_quantile = args.top_entropy_quantile
-        self.vllm_importance_sampling_correction = args.vllm_importance_sampling_correction
-        self.vllm_importance_sampling_cap = args.vllm_importance_sampling_cap
         if self.use_liger_loss and self.top_entropy_quantile < 1.0:
             raise NotImplementedError(
                 "Liger Kernels don't currently support masking token positions based on entropy."
@@ -1465,10 +1465,7 @@ def _generate_and_score_completions(
                             generation_kwargs=self.args.generation_kwargs,
                             logprobs=0,  # only return the logprob of the generated token
                         )
-                        payload = (
-                            output["completion_ids"],
-                            output["logprobs"],
-                        )
+                        payload = (output["completion_ids"], output["logprobs"])
                 else:
                     payload = None
 
@@ -1538,8 +1535,8 @@ 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]
-                logprobs = [
-                    [next(iter(lp.values())) for lp in output.logprobs]
+                all_logprobs = [
+                    [next(iter(lp.values())).logprob for lp in output.logprobs]
                     for outputs in all_outputs
                     for output in outputs.outputs
                 ]
@@ -1550,14 +1547,14 @@ 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]
-                    logprobs = logprobs[tp_slice]
+                    all_logprobs = all_logprobs[tp_slice]
 
             # Pad the completions, and concatenate them with the prompts
             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)
-            vllm_old_per_token_logps = [torch.tensor(logp, device=device) for logp in logprobs]
-            vllm_old_per_token_logps = pad(vllm_old_per_token_logps, padding_value=0.0)
+            sampling_per_token_logps = [torch.tensor(logprobs, device=device) 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
@@ -1646,8 +1643,10 @@ 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:
                 old_per_token_logps, _ = self._get_per_token_logps_and_entropies(
                     self.model,
                     prompt_completion_ids,
@@ -1662,6 +1661,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:
@@ -1782,13 +1788,37 @@ def _generate_and_score_completions(
         if has_images:
             self._logs["image"].extend(gather_object(images))
 
-        if self.vllm_importance_sampling_correction and old_per_token_logps is not None:
-            delta = old_per_token_logps - vllm_old_per_token_logps
-            probs_diff = torch.exp(delta)
-            per_token_kl = probs_diff - delta - 1
-            mean_kl = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
-            self._metrics[mode]["completions/kl_vllm"].append(self.accelerator.gather(mean_kl).nanmean().item())
-            self._metrics[mode]["completions/vllm_token_probability_difference/max"].append(probs_diff.max().item())
+        if 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,
@@ -1801,7 +1831,7 @@ def _generate_and_score_completions(
         if old_per_token_logps is not None:
             output["old_per_token_logps"] = old_per_token_logps
         if self.vllm_importance_sampling_correction:
-            output["vllm_old_per_token_logps"] = vllm_old_per_token_logps
+            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:
@@ -1935,11 +1965,8 @@ def _compute_loss(self, model, inputs):
         if entropy_mask is not None:
             per_token_loss = per_token_loss * entropy_mask
 
-        if self.vllm_importance_sampling_correction and old_per_token_logps is not None:
-            vllm_old_per_token_logps = inputs.get("vllm_old_per_token_logps")
-            rollout_imp_ratio = torch.exp(old_per_token_logps - vllm_old_per_token_logps)
-            clamped_rollout_imp_ratio = torch.clamp(rollout_imp_ratio, max=self.vllm_importance_sampling_cap)
-            per_token_loss = per_token_loss * clamped_rollout_imp_ratio
+        if 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

From c3ffa44ea4fcfd7a29d3b85aadb4f6c681b983f8 Mon Sep 17 00:00:00 2001
From: Leon Ericsson <leon.ericsson@icloud.com>
Date: Thu, 28 Aug 2025 13:03:48 +0200
Subject: [PATCH 17/23] nit

---
 trl/trainer/grpo_trainer.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 1be22c4dee7..4a9ea466c3a 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1463,7 +1463,6 @@ def _generate_and_score_completions(
                             max_tokens=self.max_completion_length,
                             guided_decoding_regex=self.guided_decoding_regex,
                             generation_kwargs=self.args.generation_kwargs,
-                            logprobs=0,  # only return the logprob of the generated token
                         )
                         payload = (output["completion_ids"], output["logprobs"])
                 else:

From 5bb5704ec21b3f562c0bdff8f202ff9c6c97ccd9 Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Thu, 28 Aug 2025 13:06:11 +0200
Subject: [PATCH 18/23] nit

---
 trl/trainer/grpo_trainer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 4a9ea466c3a..bd5aea8acc3 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1471,14 +1471,14 @@ def _generate_and_score_completions(
                 # 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, logprobs = obj_list[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]
-                logprobs = logprobs[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":

From d4d8908f0216f2dc0c23d57b8e29a790a1cc1bc6 Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Thu, 28 Aug 2025 14:17:51 +0200
Subject: [PATCH 19/23] fix logprob sanitization

---
 trl/scripts/vllm_serve.py   | 4 +++-
 trl/trainer/grpo_trainer.py | 4 +++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/trl/scripts/vllm_serve.py b/trl/scripts/vllm_serve.py
index 604c84b64a9..12e8dc2dd61 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -383,7 +383,9 @@ def sanitize_logprob(logprob):
     value = logprob.logprob
     if math.isnan(value):
         logger.warning(f"Generated NaN logprob, token logprob '{logprob}' will be ignored")
-        return float("-inf")
+        return None
+
+    return value
 
 
 def main(script_args: ScriptArguments):
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index bd5aea8acc3..e386ba072a5 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1552,7 +1552,9 @@ 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) for logprobs in all_logprobs]
+            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:

From fff82d8f49ce61d419289781a127386dd097d9c1 Mon Sep 17 00:00:00 2001
From: Leon Ericsson <leon.ericsson@icloud.com>
Date: Mon, 1 Sep 2025 18:55:04 +0200
Subject: [PATCH 20/23] nit

---
 trl/trainer/grpo_trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 667a3e52478..367457f92c1 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1501,7 +1501,7 @@ def _generate_and_score_completions(
         if has_images:
             self._logs["image"].extend(gather_object(images))
 
-        if self.vllm_importance_sampling_correction:
+        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)

From dc8972206bcc32d115914374a76ebc42a05db112 Mon Sep 17 00:00:00 2001
From: Leon Ericsson <leon.ericsson@icloud.com>
Date: Mon, 1 Sep 2025 19:27:15 +0200
Subject: [PATCH 21/23] nit

---
 trl/trainer/grpo_trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 367457f92c1..7c5210a6bd7 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1543,7 +1543,7 @@ def _generate_and_score_completions(
         }
         if old_per_token_logps is not None:
             output["old_per_token_logps"] = old_per_token_logps
-        if self.vllm_importance_sampling_correction:
+        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

From dca7e89689725455ec6594b9fbcaebbe74e4763b Mon Sep 17 00:00:00 2001
From: Leon Ericsson <leon.ericsson@icloud.com>
Date: Mon, 1 Sep 2025 19:31:04 +0200
Subject: [PATCH 22/23] nit

---
 trl/trainer/grpo_trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 7c5210a6bd7..bcc2a2d14a8 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1678,7 +1678,7 @@ def _compute_loss(self, model, inputs):
         if entropy_mask is not None:
             per_token_loss = per_token_loss * entropy_mask
 
-        if self.vllm_importance_sampling_correction:
+        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:

From 84fc6263a2eebc0bc73e126abeb4f5eb004aa12c Mon Sep 17 00:00:00 2001
From: Leon <leon.ericsson@foi.se>
Date: Wed, 3 Sep 2025 08:16:13 +0200
Subject: [PATCH 23/23] update comment regarding old_per_token_logps

---
 trl/trainer/grpo_trainer.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 085a3590956..3bb1a0ab4d7 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1359,7 +1359,9 @@ def _generate_and_score_completions(
             # 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 or self.use_vllm:
+            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,
@@ -1645,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