[rajathpi]

What does this PR do?

AsyncGRPOTrainer gets per-token generator logprobs from the vLLM server over
HTTP (/v1/completions). vLLM's completion response types token_logprobs as
list[float | None] and emits JSON null for a step it has no logprob for, so
the rollout worker can receive None in token_logprobs. Two consequences:

1. Crash: the collator builds torch.tensor([...], dtype=torch.float32)
   over those values, and a None raises TypeError: must be real number, not NoneType, ending a long training run on a single token.
2. Corrupted weights: any NaN that reaches old_log_probs makes the PPO
   ratio and the whole loss NaN, and the optimizer still steps on it. The
   Trainer nan/inf logging filter substitutes the running mean (GRPO: a NaN training step is invisible in log_history, so a corrupted run reports a clean loss curve #6702), so
   the loss curve stays plausible. Make a non-finite training step visible in log_history #6862 makes that condition visible by
   reporting frac_nonfinite_loss; this PR removes one of the causes it would
   report, so the two are complementary.

Scope note, corrected after review: the NaN logprob from #6166 does not
reach the async trainer as null. The worker's request is non-streaming, so
vLLM serializes the response with starlette's json.dumps(allow_nan=False),
which raises on NaN; vLLM's ValueError handler turns that into an HTTP error
and the worker retries. The null this PR guards against is the missing-step
case, and the NaN guard in the loss is defensive. I have not observed either in
a production async run; the behaviour is demonstrated by the tests and the
injected-sample run in the thread.

#6693 hardened the sync GRPOTrainer and grpo_with_replay_buffer against
None logprobs. The async trainer was left behind. This PR completes it:

• async_rollout_worker.py: map None to float("nan") when parsing the
  completions response. Both rollout paths share this method. The existing
  0.0 sentinel is deliberately not reused, since 0.0 is a legitimate
  logprob and these are trained (mask=1) positions.
• async_grpo_trainer.py: in compute_loss, replace NaN old_log_probs with
  log_probs.detach() before the ratio. log_ratio becomes exactly 0, so the
  coefficient is exactly 1 and clipping is inert, while the gradient still
  flows through log_probs. The token trains as plain policy gradient and the
  KL and clip metrics stay finite. This is the same fallback the sync trainer
  already uses when old logprobs are unavailable. The token is treated as
  on-policy even when the rollout is stale, a small deliberate bias.

Three regression tests in TestUnscorableTokenLogprobs pin the ingestion
mapping, the collator handling NaN rows, and the loss neutralization (finite
loss, finite gradients, finite KL). The ingestion and loss tests each fail
without their half of the fix.

#6297 covers the same class for the sync trainer only and predates #6693.
#6862 is complementary rather than overlapping: it surfaces a non-finite loss
after the fact in the same trainers and cites #6166 as one of the causes. This
PR prevents that cause, and additionally fixes the crash, which no other open
PR addresses.

pytest tests/experimental/test_async_grpo_trainer.py -k Unscorable

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

AI writing disclosure

We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

cc @qgallouedec

---

Note

Medium Risk
Changes RL training loss semantics for unscorable tokens (on-policy fallback) and touches the async rollout path; scope is narrow with targeted tests but affects gradient updates when vLLM omits logprobs.

Overview
Async GRPO now survives vLLM returning JSON null for unscorable per-token logprobs (serialized NaN), which previously could crash collation (None in tensors) or poison the PPO loss with NaN weights while logs looked fine.

The rollout worker maps None → float("nan") when parsing token_logprobs. In compute_loss, NaN old_log_probs are replaced with the current policy logprob (detached), matching sync GRPO: importance ratio 1, finite KL/clip metrics, gradients still flow.

TestUnscorableTokenLogprobs adds three regression tests for ingestion, collator tensorization, and loss/gradient finiteness.

^{Reviewed by Cursor Bugbot for commit f134993. Bugbot is set up for automated code reviews on this repo. Configure here.}

[qgallouedec]

thks, can you please provide a minimal repro?

[qgallouedec]

thks, can you please provide a minimal repro?

[rajathpi]

Sure. Only the HTTP call is stubbed, everything else is real TRL code.

import asyncio
import torch

from trl.experimental.async_grpo.async_grpo_trainer import DataCollatorForRollout
from trl.experimental.async_grpo.async_rollout_worker import _AsyncRolloutLoop

# vLLM serializes a NaN logprob as JSON null, so the client receives None
# for a token the engine could not score.
VLLM_RESPONSE = {
    "choices": [{"token_ids": [10, 11, 12], "logprobs": {"token_logprobs": [-0.5, None, -0.7]}}]
}

loop = _AsyncRolloutLoop.__new__(_AsyncRolloutLoop)
loop.model_name, loop.max_tokens, loop.temperature = "m", 8, 1.0
loop.top_p, loop.top_k, loop.repetition_penalty = 1.0, -1, 1.0
loop.min_p, loop.request_timeout = None, 1
loop._post = lambda *a, **k: asyncio.sleep(0, result=VLLM_RESPONSE)

token_ids, logprobs = asyncio.run(loop._generate_one_turn([1, 2]))
print(f"1. worker returns: {logprobs}")

sample = {
    "input_ids": [1, 2] + token_ids,
    "completion_mask": [0, 0, 1, 1, 1],
    "old_log_probs": [0.0, 0.0] + logprobs,
    "advantage": 1.0,
    "group_id": 0,
    "metrics": {"reward": 1.0},
}
try:
    DataCollatorForRollout(pad_token_id=0, num_processes=1).torch_call([[[sample]]])
    print("2. collator: ok")
except TypeError as e:
    print(f"2. collator crashes: {type(e).__name__}: {e}")

# Why the loss also needs handling: had the value arrived as NaN rather than
# None there is no crash, it propagates through the ratio instead.
old = torch.tensor([[-0.5, float("nan"), -0.7]])
new = torch.tensor([[-0.4, -0.6, -0.8]], requires_grad=True)
loss = -(torch.exp(new - old) * torch.ones_like(old)).mean()
print(f"3. loss with one NaN old logprob: {loss.item()}")

On main:

1. worker returns: [-0.5, None, -0.7]
2. collator crashes: TypeError: must be real number, not NoneType
3. loss with one NaN old logprob: nan

On this branch:

1. worker returns: [-0.5, nan, -0.7]
2. collator: ok
3. loss with one NaN old logprob: nan

Step 3 prints nan in both runs on purpose: it is a plain tensor demonstration of
why the loss needs its own guard, not a call into the trainer. The guard itself is
covered by test_compute_loss_neutralizes_nan_old_log_probs, which drives the real
compute_loss and asserts a finite loss, finite gradients and a finite KL.

An end to end repro against a live vLLM server needs a token the engine actually
fails to score, which is intermittent by nature (#6166). Stubbing the response is
the deterministic equivalent.

[qgallouedec]

No but a real training repro

[qgallouedec]

needs a token the engine actually
fails to score

I guess you can recover this from the logs

[rajathpi]

Fair point on the logs. I don't have a long async production run to mine, but that
recovery exists for this bug class already: #6166 is a real 8xA100 colocate run where
entropy collapse (~4e-4) made vLLM emit a NaN processed logprob and the sync trainer
died with Could not infer dtype of NoneType. The async trainer consumes the same
logprobs over HTTP, where NaN cannot survive JSON serialization: it arrives as null,
so the worker hands the trainer a None.

For the training repro, I rented an H100 NVL and ran the real trainer end to end
(torch 2.11.0+cu128, transformers 5.15.1). Following the pattern #6693 used for the
sync twin, the injection point is the rollout buffer via the suite's own
_StubRolloutWorker:

import os

import torch
from datasets import load_dataset
from transformers import AutoTokenizer

from tests.experimental.test_async_grpo_trainer import (
    _StubRolloutWorker,
    _StubWeightTransfer,
    dummy_reward_func,
)
from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer

# INJECT=none -> what main's worker emits for an unscorable token (JSON null -> None)
# INJECT=nan  -> what the post-fix worker emits (null mapped to NaN at ingestion)
BAD = None if os.environ.get("INJECT", "none") == "none" else float("nan")


class OneUnscorableTokenWorker(_StubRolloutWorker):
    def _fill_queue(self):
        super()._fill_queue()
        sample = self.rollout_buffer.get()
        sample.old_log_probs[-1] = BAD
        self.rollout_buffer.put(sample)


model_id = "trl-internal-testing/small-Qwen2ForCausalLM-2.5"
dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_completion", split="train")
args = AsyncGRPOConfig(
    output_dir="async-grpo-unscorable-repro",
    max_steps=2,
    per_device_train_batch_size=3,
    num_generations=3,
    max_completion_length=8,
    token_budget=256,
    vllm_server_timeout=5.0,
    report_to="none",
)
trainer = AsyncGRPOTrainer(
    model=model_id,
    reward_funcs=dummy_reward_func,
    args=args,
    train_dataset=dataset,
    rollout_worker=OneUnscorableTokenWorker(AutoTokenizer.from_pretrained(model_id), dataset, num_generations=3),
    weight_transfer=_StubWeightTransfer(),
)
trainer.train()

losses = [e["loss"] for e in trainer.state.log_history if "loss" in e]
nan_params = any(torch.isnan(p).any().item() for p in trainer.model.parameters())
print("per-step losses:", losses)
print("final train_loss:", trainer.state.log_history[-1].get("train_loss"))
print("model has NaN parameters:", nan_params)

Observed results:

main (6d484ba), INJECT=none: trainer.train() crashes at the first batch fetch:

File "trl/experimental/async_grpo/async_grpo_trainer.py", line 560, in torch_call
  torch.tensor([lp for example in group for lp in example["old_log_probs"]], dtype=torch.float32)
TypeError: must be real number, not NoneType

main, INJECT=nan (this is the part I find most convincing):

per-step losses: [0.5106865167617798, 0.0]
final train_loss: 0.2553432583808899
model has NaN parameters: True

The reported loss curve looks completely healthy while the weights are already
corrupted. That is the failure mode your logging_nan_inf_filter note in #6702
describes, observed live.

this branch, INJECT=nan:

per-step losses: [0.16899284720420837, 0.5174540281295776]
final train_loss: 0.343223437666893
model has NaN parameters: False

On the same H100 I also ran
pytest tests/experimental/test_async_grpo_trainer.py -k "Unscorable or test_train or test_init_minimal":
6 passed, including the GPU-gated test_train.

One caveat for completeness: INJECT=none still crashes on this branch, by
construction. The stub bypasses ingestion, which is where the fix maps null to NaN,
and post-fix the real worker can no longer put a None in the buffer. NaN is the
faithful representation of post-fix behavior.

Happy to add the NaN variant as a regression test next to test_train, behind the
same skipif: train two steps with one poisoned sample, assert finite loss and clean
parameters.

[qgallouedec]

so tldr you never encountered this issue in one of your training?

[rajathpi]

No. Found it reading #6693, which fixed this for GRPOTrainer and
grpo_with_replay_buffer but not the async twin. The field report is #6166,
same trigger, one code path over. Async just has less protection and a worse
failure mode (null over HTTP, so it crashes or silently corrupts, per the
runs above).

[behroozazarkhalili]

I wrote #6693 and the #6862 metric this PR cites, so this is familiar ground. Instead of just reading the diff I checked out the branch and ran your tests both ways.

Everything checks out:

• Your red/green claim holds. On your head the four Unscorable tests pass. If I revert just the two source hunks and keep your tests, the ingestion test and the loss test fail, exactly as the PR text says they should. (The collator test passes either way, since the collator never had a problem with NaN, only with None. Fine as a regression pin, just not part of the red set.)
• You patched the right place. /v1/completions is the only generation endpoint in the rollout worker, and both the single-turn and tool-calling paths go through _generate_one_turn, so the None to NaN mapping there catches every rollout.
• In compute_loss, old_log_probs is used exactly once, in the log_ratio line right below your guard, so nothing else can see the NaN.
• I double-checked your "same fallback the sync trainer uses" line against grpo_trainer.py:3155, and yes, that is exactly what the sync trainer does.
• Good call not reusing the 0.0 sentinel. Those fills are for mask=0 positions; here the token is trained, and 0.0 would quietly mean "probability 1".
• I also like that the loss test asserts the gradients are finite, not just the loss. A NaN slipping into backward would still get caught.

One thought, not a blocker. In the sync trainer this situation is handled by dropping the token from the importance-sampling correction, while your fix trains it as plain on-policy policy gradient with ratio 1. For sync GRPO those are basically the same thing. For AsyncGRPO they are not quite: the rollout can be several policy versions old, so forcing the ratio to 1 on a genuinely off-policy token introduces a small bias. It is bounded, it only hits tokens vLLM itself could not score, and it clearly beats a crash or a NaN loss, so I would keep your approach. I would just add a one-line comment saying the token is treated as on-policy even when the rollout is stale, so the next person changing this code knows it was a deliberate choice. #6945 is related if you want the bigger picture on AsyncGRPO and importance sampling.

Nice work closing this out for the async trainer, and thanks for the real regression coverage.

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[AmineDiro]

Is this a real or theoretical failure mode?

[rajathpi]

Real for the sync trainer: #6166 is a field report, an 8xA100 run where entropy collapse made vLLM emit a NaN logprob and training died. For async I have only reproduced it with a stubbed vLLM response and in the H100 run above. I haven't seen it in a long async production run. The async path has the same trigger and gets null over HTTP, so it fails the same way, it just has less protection. Since the cost is a hard crash or a silently NaN loss, I'd keep the guard even though the trigger is rare.

Also added the one-line comment @behroozazarkhalili suggested, noting the token is deliberately treated as on-policy even when the rollout is stale.

[rajathpi]

Correction to my last comment, after checking the vLLM server path more carefully.

The worker calls /v1/completions non-streaming, and that path returns via JSONResponse(content=response.model_dump()) (completion/api_router.py). Starlette's JSONResponse.render uses json.dumps(..., allow_nan=False), which raises ValueError: Out of range float values are not JSON compliant: nan on a NaN logprob; I confirmed this against a real CompletionResponse on vLLM bca7bea24. vLLM registers a ValueError exception handler, so the client gets an error response and the worker's raise_for_status() triggers a retry, not a null. Only the streaming path (model_dump_json) serializes NaN as null, and the worker does not use it.

So for a NaN logprob the async trainer never receives None. The null that does reach it is the missing-step case (serving.py appends None when a step has no logprobs), which is a different condition from #6166. The failure mode is theoretical for async as I framed it: the guard is still defensive against that null and against any NaN in old_log_probs, but it is not the #6166 crash one code path over.

I have rewritten the PR description and the two code comments to say this (f134993). If you would rather not carry a defensive guard for a condition nobody has reported on the async path, I am fine closing this.

[qgallouedec]

So IIUC: first it was about supporting NaN, that comes from a run that had already collapsed, then a NaN that never reaches the trainer, now a null that, as far as I can tell, only shows up with echo=True, which we don't use? Closing because it seems that this PR never originated from a bug in genuine training.

Please re-open with a clear motivation if I missed something