From 0a53d4618dbeab2a944e30b870a8a141fe5f36e2 Mon Sep 17 00:00:00 2001
From: Yuchen Wang <93700456+yuchenwang3@users.noreply.github.com>
Date: Mon, 24 Aug 2026 22:45:46 -0700
Subject: [PATCH] fix: sanitize non-finite logprobs in vLLM async worker

vLLM can return NaN or Inf logprobs, which Starlette JSONResponse rejects because it serializes with allow_nan=False. Sanitize the dumped response before serialization so the rollout request remains valid JSON.

Emit one warning with a per-response replacement count so numerical instability remains visible, and cover nested values plus warning behavior with a unit test.

Signed-off-by: Yuchen Wang <93700456+yuchenwang3@users.noreply.github.com>
(cherry picked from commit e56338ddbf27b6566a42e2ec75a51a5bb8be4e27)
---
 .../generation/vllm/vllm_worker_async.py      | 33 +++++++++++++++++--
 .../models/generation/test_vllm_generation.py | 26 +++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py
index 01076d19481..6928da0ff19 100644
--- a/nemo_rl/models/generation/vllm/vllm_worker_async.py
+++ b/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
@@ -62,6 +63,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 VllmAsyncGenerationWorkerImpl(
     VllmAsyncCheckpointEngineRpcMixin, BaseVllmGenerationWorker
 ):
@@ -804,8 +833,8 @@ async def create_chat_completion(
 
             elif isinstance(generator, ChatCompletionResponse):
                 return JSONResponse(
-                    content=model_dump_chat_response_with_dynamic_message_fields(
-                        generator
+                    content=_replace_non_finite(
+                        model_dump_chat_response_with_dynamic_message_fields(generator)
                     )
                 )
 
diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py
index 4429f26446e..5a70a31208a 100644
--- a/tests/unit/models/generation/test_vllm_generation.py
+++ b/tests/unit/models/generation/test_vllm_generation.py
@@ -14,6 +14,7 @@
 
 import importlib.util
 import json
+import math
 import os
 import sys
 import types
@@ -44,6 +45,7 @@
 )
 from nemo_rl.models.generation.vllm.vllm_worker_async import (
     VllmAsyncGenerationWorkerImpl,
+    _replace_non_finite,
 )
 from nemo_rl.models.policy import LoRAConfig, PolicyConfig
 from nemo_rl.models.policy.lm_policy import Policy
@@ -2027,6 +2029,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