Add sanitize_logprob function for NaN handling in vLLM log probabilities

trl/generation/vllm_generation.py

@@ -15,6 +15,8 @@
 """vLLM-based generation backend for TRL trainers."""
 
 import json
+import logging
+import math
 import os
 from collections.abc import Callable
 from contextlib import nullcontext
@@ -34,6 +36,18 @@
 from .vllm_client import VLLMClient
 
 
+logger = logging.getLogger(__name__)
+
+
+def sanitize_logprob(logprob):
+    value = logprob.logprob
+    if math.isnan(value):
+        logger.warning(f"Generated NaN logprob, token logprob '{logprob}' will be ignored")
+        return None
+
+    return value
+
+
 if TYPE_CHECKING:
     from accelerate import Accelerator
     from peft import PeftModel
@@ -666,7 +680,7 @@ def generate(self, prompts: list, num_generations: int, profiler: ProfilingConte
                 all_prompt_ids = [output.prompt_token_ids for output in all_outputs]
                 all_completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs]
                 all_logprobs = [
-                    [next(iter(lp.values())).logprob for lp in output.logprobs]
+                    [sanitize_logprob(next(iter(lp.values()))) for lp in output.logprobs]
                     for outputs in all_outputs
                     for output in outputs.outputs
                 ]

trl/scripts/vllm_serve.py

@@ -34,6 +34,7 @@
 from transformers import AutoTokenizer, is_torch_xpu_available, is_vision_available
 
 from trl import TrlParser
+from trl.generation.vllm_generation import sanitize_logprob
 from trl.import_utils import (
     is_fastapi_available,
     is_pydantic_available,
@@ -431,17 +432,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_logprob(logprob):
-    import math
-
-    value = logprob.logprob
-    if math.isnan(value):
-        logger.warning(f"Generated NaN logprob, token logprob '{logprob}' will be ignored")
-        return None
-
-    return value
-
-
 def _replace_prefix_tokens(
     tokenizer,
     model_prefix_token_ids: list[int],