fix: sanitize non-finite logprobs in vllm async worker

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -16,6 +16,7 @@
 import copy
 import gc
 import logging
+import math
 import threading
 import time
 import uuid
@@ -63,6 +64,34 @@
 from nemo_rl.distributed.refit_watchdog import RefitAborted, is_refit_abort
 
 
+def _replace_non_finite(obj: Any) -> Any:
+    """Replace NaN/Inf floats with 0.0 and warn once per response."""
+    replacement_count = 0
+
+    def replace(value: Any) -> Any:
+        nonlocal replacement_count
+
+        if isinstance(value, float):
+            if not math.isfinite(value):
+                replacement_count += 1
+                return 0.0
+            return value
+        if isinstance(value, dict):
+            return {key: replace(item) for key, item in value.items()}
+        if isinstance(value, list):
+            return [replace(item) for item in value]
+        return value
+
+    sanitized = replace(obj)
+    if replacement_count:
+        LOGGER.warning(
+            "Replaced %d non-finite float values in the vLLM chat completion "
+            "response with 0.0; this may indicate numerical instability.",
+            replacement_count,
+        )
+    return sanitized
+
+
 class _AsyncLLMHTTPClient:
     """Keep HTTP generation on the loop that owns AsyncLLM request state.
 
@@ -1192,7 +1221,7 @@ async def create_chat_completion(
                 content = await asyncio.to_thread(
                     worker_self._finish_request_capture, request, content
                 )
-                return JSONResponse(content=content)
+                return JSONResponse(content=_replace_non_finite(content))
 
             worker_self._abort_request_capture(request, reason="streaming_response")
             return StreamingResponse(content=generator, media_type="text/event-stream")

tests/unit/models/generation/test_vllm_generation.py

@@ -15,6 +15,7 @@
 import asyncio
 import importlib.util
 import json
+import math
 import os
 import sys
 import threading
@@ -47,6 +48,7 @@
 from nemo_rl.models.generation.vllm.vllm_worker_async import (
     VllmAsyncGenerationWorkerImpl,
     _AsyncLLMHTTPClient,
+    _replace_non_finite,
 )
 from nemo_rl.models.policy import LoRAConfig, PolicyConfig
 from nemo_rl.models.policy.lm_policy import Policy
@@ -2479,6 +2481,30 @@ def test_vllm_deferred_model_load(cluster, tokenizer):
     vllm_generation.shutdown()
 
 
+def test_replace_non_finite(caplog):
+    result = _replace_non_finite(
+        {
+            "nan": math.nan,
+            "pos_inf": math.inf,
+            "nested": [-math.inf, 1.5, "unchanged", None],
+        }
+    )
+
+    assert result == {
+        "nan": 0.0,
+        "pos_inf": 0.0,
+        "nested": [0.0, 1.5, "unchanged", None],
+    }
+    assert caplog.messages == [
+        "Replaced 3 non-finite float values in the vLLM chat completion response "
+        "with 0.0; this may indicate numerical instability."
+    ]
+
+    caplog.clear()
+    assert _replace_non_finite({"finite": 1.5}) == {"finite": 1.5}
+    assert not caplog.records
+
+
 def test_VllmAsyncGenerationWorker_replace_prefix_tokens(tokenizer):
     # This test assumes the tokenizer model is for the Qwen 3 family
     eos_token_id = tokenizer.eos_token_id