diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index bb77973fed..fe9d361781 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 @@ -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") diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 1e2885f8a4..8a5a91a647 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/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