fix: sanitize non-finite logprobs in vllm async worker

[yuchenwang3]

vLLM can emit non-finite logprobs. We hit a NaN consistently around 49K context, and vLLM also uses -inf intentionally when sampling filters remove tokens. Starlette's JSONResponse serializes with allow_nan=False, so any remaining NaN or infinity raises ValueError and turns the chat-completion request into a 500.

This PR sanitizes the dumped response immediately before JSON serialization:

• recursively replaces NaN, +inf, and -inf floats with 0.0;
• emits at most one warning per response with the number of replaced values, so numerical instability is visible rather than silently masked;
• wraps the current model_dump_chat_response_with_dynamic_message_fields result on main.

0.0 is a conservative sentinel for the rollout importance ratio because the train-time logprob is recomputed. The warning/count is intended to preserve the signal needed to investigate genuine model or vLLM numerical failures.

The unit test covers nested dict/list values, all three non-finite forms, finite and non-float passthrough, replacement count, one-warning behavior, and the no-warning finite path. Ruff lint/format, syntax compilation, and DCO pass on the rebased head.

Recent bounded retry and exhausted-rollout handling reduce the blast radius of a 500, but they do not make this response serializable; a deterministic NaN can still fail every retry. This keeps the fix at the source boundary.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuchenwang3]

This has been waiting on the copy-pr-bot gate since Jun 26 — could a maintainer /ok to test and take a review pass? The fix is small and self-contained: non-finite logprobs from the vllm async worker get sanitized at the source instead of flowing into training. Happy to address feedback quickly.

[yuchenwang3]

@yuki-97 sorry for the direct ping — you're the most recent committer on vllm_worker_async.py and merge most PRs here, so you seemed like the right person to ask. This one-file fix sanitizes non-finite logprobs at the source: vLLM can return NaN/Inf logprobs at long context, and JSONResponse (allow_nan=False) turns a single one into an HTTP 500 that aborts GRPO rollout collection. Rebased onto today's main — it now composes with the routed-experts serialization. Could you /ok to test and take a look?

[yuki-97]

thanks @yuchenwang3 , fix LGTM.

but @yfw @ananthsub could you help confirm whether this will affect something?
it looks like:

• before: traj with nan logprobs will return http 500 to gym.
• after: will return traj, and replace logprobs nan with 0.

[yuchenwang3]

Thanks @yuki-97! Unit test added in 3e688cb — covers the float / nested-dict / nested-list branches plus ±inf and non-float passthrough, per your sketch.

On the behavior question for @yfw @ananthsub: the intent is that a single non-finite logprob no longer 500s the whole request (which aborts GRPO rollout collection through the Gym proxy). Replacing with 0.0 makes the affected token look more likely under the rollout policy, so any importance-style correction downweights rather than amplifies it — conservative in that direction — and the train-time recomputed logprobs (which dominate the loss) are unaffected. Happy to switch the sentinel or drop only the affected trajectory instead if you'd prefer different semantics.

[yuchenwang3]

@yfw @ananthsub bump — unit tests went in three weeks ago per @yuki-97's sketch, CI is green, and the one open item is the behavior question in my Jul 15 comment that's addressed to you two. Could you weigh in so this can move to a decision?

[yuki-97]

hi @terrykong , could you help take a look? #2962 (review)

[yfw]

@yuchenwang3 sorry for the delay in looking into this. In the past when we've seen nan logprobs, this was usually indicative of a bug in vllm. My concern is this may mask an issue with the inference engine that needs further investigation. Do you have some cases where the nan logprob is expected from vllm?

[yuchenwang3]

@yfw fair question — I went digging in the vLLM source instead of hand-waving, and it changed how I'd frame this.

vLLM itself already sanitizes non-finite logprobs at the API boundary: the OpenAI serving layer clamps with max(logprob, -9999.0) right before serialization (vllm/entrypoints/openai/chat_completion/serving.py, twice; same in completion/serving.py, where the comment says "JSON-serializable float that OpenAI uses"). So non-finite values reaching the serving layer is normal enough that vLLM defends against it in exactly the layer this PR touches.

There's also a documented source of -inf that has nothing to do with engine bugs: logprobs_mode="processed_logprobs" returns logprobs after temperature/top-k/top-p, so truncated tokens are -inf by construction (vllm/config/model.py). NeMo-RL's own vllm config exposes that mode, so this is inside the supported config space here.

And the reason it bites as NaN specifically: max(nan, -9999.0) is nan in Python, so vLLM's clamp catches -inf but passes NaN straight through — which matches the failure we hit at ~49K context, ValueError: Out of range float values are not JSON compliant: nan.

On masking: agreed that a NaN can mean an upstream numerics problem worth chasing, and I don't want to hide it. But one bad trajectory 500-ing the endpoint and killing the whole GRPO rollout is the wrong failure mode — same call vLLM's clamp already makes. I can add a LOGGER.warning with a per-request count of replaced values so it stays visible without being fatal. Would that work for you?

[yuchenwang3]

Made this runnable to be sure (CPU-only, no GPU needed). The clamp/serialization chain:

vLLM clamp on -inf : max(-inf, -9999.0) = -9999.0
vLLM clamp on NaN  : max(nan,  -9999.0) = nan
serialization of a NaN logprob -> ValueError: Out of range float values are not JSON compliant: nan
after _replace_non_finite -> {"choices": [{"logprobs": {"content": [{"token": "x", "logprob": 0.0}]}}]}

That ValueError is the exact string from our rollout failures, so the NaN-specific story checks out.

For the -inf side I ran vLLM's own apply_top_k_top_p_pytorch (copied as-is from vllm/v1/sample/ops/topk_topp_sampler.py) with top_p=0.9 on a toy 16-token vocab, then log_softmax:

processed logits: -inf -inf -0.3 -inf 0.8 0.7 -0.3 -inf 0.3 -inf 0.3 0.3 0.1 1.2 1.1 -0.2
log_softmax -> 5/16 logprobs are exactly -inf

i.e. under processed_logprobs, -inf in the payload is just what truncation looks like.

The NaN mechanism that fits the ~49K observations — a single overflowed logit is enough:

bf16: 3e38 * 2 = inf
logits [1.0, 2.0, inf, 0.5] -> log_softmax -> [nan, nan, nan, nan]

inf - inf = nan in the max-subtraction, so the whole row goes NaN at once.

Can share the script if useful; the warning + counter offer above still stands.

[yuchenwang3]

@yfw @ananthsub following up on the semantics question: would sanitizing with a single per-request warning and replacement count address the masking concern, or would you prefer the response to fail/drop instead? I have the warning/count version rebased on current main and ready to update once we agree on that behavior.

[yuchenwang3]

Rebased this onto current main, fixed the DCO identity mismatch, and added one warning per response with the number of non-finite values replaced. The call site now wraps the current dynamic-message-fields dump helper. @yfw @ananthsub @terrykong, could one of you run /ok to test e56338ddbf27b6566a42e2ec75a51a5bb8be4e27 and confirm whether sanitize + warning is acceptable versus failing/dropping the response?

[yuchenwang3]

@yfw @ananthsub @terrykong @yuki-97 — I refreshed this onto current main and resolved the only conflict (the new refit-watchdog import and this helper shared the same insertion point; both are preserved). The sanitize + one-warning/count behavior is unchanged.

Current head: 219af62b27d036e591a869e436e6dec4a850d750. Ruff check/format, diff check, and a serialization-chain harness all pass: the raw NaN payload fails strict JSON encoding, while the sanitized payload encodes successfully and emits exactly one warning with the replacement count. Could someone confirm whether the observable sanitize behavior is acceptable and run /ok to test 219af62b27d036e591a869e436e6dec4a850d750? The last maintainer design feedback was on August 7, and the current-main implementation is ready for a decision.

[yuchenwang3]

/ok to test 0a53d46

[yuchenwang3]

@yfw, following up on your concern about masking inference-engine bugs. The PR now makes replacements observable with a per-response warning and count, but the policy choice still needs your judgment: is sanitize-with-warning acceptable, or should these responses fail/be dropped? Could you confirm the preferred direction before we take this further? Thanks.

[yuchenwang3]

Synced with current main and resolved the async-worker conflicts. The upstream engine-loop adapter and token-capture path are preserved; response sanitization now happens after capture, immediately before JSON serialization. This leaves the original capture input untouched.

Changed-file Ruff/format checks pass. Isolated CPU checks cover NaN/Inf serialization, preservation of finite values and capture coordinates, and AST comparison confirming the rest of the upstream async worker is unchanged. These are not full vLLM/Ray integration tests. The sanitize-versus-reject design question remains open; this update does not change that policy.