Blocking await put() can hang on shutdown

Medium Severity

The old put_nowait() retry loop that checked stop_event between attempts was replaced with a bare await self._groups_to_score.put(group). If the scoring queue is full when stop() is called, the score loop exits (stops consuming), and this await put() blocks forever. Since both loops run inside asyncio.gather, the generate loop never finishes, preventing the worker thread from shutting down cleanly.

 

What the motivation to use a callback instead of simply calling the weight sync in the trainer?

_sync_weight() was called at the top of _inner_training_loop, before super()._inner_training_loop() ran. it's super()._inner_training_loop() that performs FSDP wrapping and accelerator.prepare() (which moves the model to GPU). So weights were still on CPU when NCCL tried to broadcast them.

So I had:

  ┃  [rank0]:   File "/fsx/amine_dirhoussi/trl/.venv/lib/python3.11/site-packages/vllm/distributed/device_communicators/pynccl.py", line
  ┃  363, in broadcast
  ┃  [rank0]:     assert tensor.device == self.device, (
  ┃  [rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  ┃  [rank0]: AssertionError: this nccl communicator is created to work on cuda:0, but the input tensor is on cpu

IMO we can remove this

it's safer because it will fail if we try to step before the reset

I think test_list has to be init to [] if I am not mistaken

because it's missing the docstring, it won't be exposed as a tool to the model. If it's not meant to be exposed, then I'd recommend def _is_done. If it is meant to be exposed as a tool, it requires a docstring.

it is not a tool function. the _is_done means it it a private function but we cant to signal that it needs to be implemented.
The issue we currently have is that user can't define a stop function to exist the generation loop. For example, when in code environment, we want to exist if the tests success but the current setup keeps rolling until we reach max_turns or token exhaustion

what if one sequence in the group reaches the limit? The whole group is discarded?

max_seq_length is only used in _generate_one function to truncate generation or to skip prompt. It is also skipped in scoring loop:

  if self.max_seq_length is not None and seq_len > self.max_seq_length:
      logger.warning(f"Dropping overlong sample (seq_len={seq_len}, max_seq_length={self.max_seq_length})")

you don't need this change, apply_chat_template will get the json schema for you

I added this because I kept getting error in my async_grpo_mbpp.py env. The isfunction check fails as it is a method.

To be precise: transformers.utils.chat_template_utils.render_jinja_template uses inspect.isfunction() to validate tools, but isfunction() returns False for bound methods. The environment's execute_python_code is collected as a bound method via inspect.getmembers so I was getting

  ┃    File "/fsx/amine_dirhoussi/trl/trl/experimental/async_grpo/async_rollout_worker.py", line 287, in _run_loops
  ┃      await asyncio.gather(
  ┃    File "/fsx/amine_dirhoussi/trl/trl/experimental/async_grpo/async_rollout_worker.py", line 392, in _generate_loop
  ┃      prompt_ids = self.tokenizer.apply_chat_template(
  ┃                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  ┃    File "/fsx/amine_dirhoussi/trl/.venv/lib/python3.11/site-packages/transformers/tokenization_utils_base.py", line 1667, in
  ┃  apply_chat_template
  ┃      rendered_chat, generation_indices = render_jinja_template(
  ┃                                          ^^^^^^^^^^^^^^^^^^^^^^
  ┃    File "/fsx/amine_dirhoussi/trl/.venv/lib/python3.11/site-packages/transformers/utils/chat_template_utils.py", line 493, in
  ┃  render_jinja_template
  ┃      raise ValueError(
  ┃  ValueError: Tools should either be a JSON schema, or a callable function with type hints and a docstring suitable for auto-conversion to a schema.

Maybe I am missing a quick fix or I defined the Env wrongly?

you must upgrade transformers to at least 5.2.0

for ref huggingface/transformers#43968

Weird, I am sure I have the correct transformers version in env:

(trl) amine_dirhoussi@login-node-1 trl ±|rollout-generation ✗|→ uv pip list | rg transformers
transformers                             5.2.0

I'll retry this specific parsing function is isolation 👍🏼

are you sure it's needed? if you set max_inflight_tasks = max_staleness * per_device_train_batch_size * gradient_accumulation_steps * num_processes (default), then you shouldn't need this in my understanding

Yes we still need this to cancel current inflight. max_inflight_tasks will limit fresh requests. If for example there is a generation that takes a long time that gets stale, we need to cancel it. If we don't we'll be wasting vllm's compute resources on a stale requests that will get dropped.

ah yes good point

I strongly think that we shouldn't add any done signal, see #5235 (comment)

Requiring extra methods unnecessarily narrows the set of valid implementations and reduces substitutability. We should depend on the minimal protocol the class actually needs. There is nothing we can't do without is_done, so I'd recommend simply not having it. Plus at inference, the model generation won't have access to this flag anyway, so it's useless

The done was added for this case (cf earlier comment):

The issue we currently have is that user can't define a stop function to exit the generation loop. For example, when in code environment, we want to exist if the tests succeed but the current setup keeps rolling until we reach max_turns or token exhaustion

yes I know, but ideally, the model should learn no to continue calling the tool if it gets an error each time. IMO the right way to handle done would be:

class MyEnv:
    def reset(self, **kwargs):
        self._done = False

    def my_func(self):
        """Some nice documentation"""
        if self._done:
            raise Exception("Session expired, you can't use this tool anymore!")
        if some_condition:
            self._done = True
        return

If the generation stops when the environment is done, then the model never learns to know that the [environment is done] implies that [it can't use it anymore].

Sorry, maybe I didn't explain well the usage of the is_done in this portion codebase. The current impl does the following:

async def _generate_one(...):
    while True:
        # .... 
        if tool_calls is None or (max_iterations is not None and iteration_num >= max_iterations):
            return ...

Now we can have done() raise an Exception but I am not sure it provides a clean signal to know that this rollout had exception because of some env specific errors vs it genuinely reached end of turn. If we go with Exception route, I would image that we need to define a library side Exception like StopGeneration(Exception) to cleanly catch it and know when env has stopped. But then this is equivalent to defining a done function IMO.

Also having a done env defined tool extends tools context for the LLM which can be a good choice in some specific cases but overall maybe unnecesary ??

So if we don't use the done flag to stop the generation, what's the point of having it? I can see it's used here, but the reason is not clear:

https://github.com/huggingface/trl/pull/5299/changes#diff-175873f8c2363ec8322c4a0af91437e83da991bb5d388f2d01342b2d31eaf8a3R822

it seems like a proxy the a new is_truncated flag, which I don't quite understand either. The truncation is equivalent to "the last token isn't an EOS", why adding this flag here?

But more generally, don't you think it's out of the scope of this PR?

So if we don't use the done flag to stop the generation, what's the point of having it? I can see it's used here, but the reason is not clear:

To be precise, I put the previous generation loop. The changed one does use it to break : line

 if is_done is not None and is_done():
     return _build_completion(turns, truncated=False, total_duration=time.monotonic() - t_start)

it seems like a proxy the a new is_truncated flag, which I don't quite understand either. The truncation is equivalent to "the last token isn't an EOS", why adding this flag here?

Truncation is separate from done and is mainly for metrics and debugging purposes. It can be computed from "isn't last EOS" but its just a bool and recomputing something that small seems unnecessary ?

But more generally, don't you think it's out of the scope of this PR?

I think there was some feature creep, but this PR was originally intended to change the async-grpo branch directly. All the commits were added to debug the mbpp example: metrics, trajectories, the done flag, etc.

I believe the first async-grpo PR iteration had metrics, so I viewed this as a logical extension, especially since there wasn’t a strictly defined feature scope for this specific PR

Advantage normalization precedes sample dropping, breaking GRPO

Medium Severity

In _score_group, advantages are normalized across all completions in the group, then individual samples exceeding max_seq_length are dropped. The surviving samples retain advantage values computed relative to the full group (including dropped samples). This breaks the zero-mean property of GRPO normalization and can bias training—e.g., if the highest-reward sample is dropped, remaining samples all train with distorted advantages.

 

Stale cancellation of partially-dispatched group causes permanent leak

Medium Severity

When _cancel_stale_tasks drops a group that has only been partially dispatched (fewer than num_generations tasks created so far), the remaining yields from _repeat_iterator for that same group_id will re-create the group via the if group_id not in pending_groups branch, but only the leftover yields produce tasks. The re-created group's pending_completed can never reach num_generations, so it permanently leaks in pending_groups and pending_completed, occupying slots that complete but never trigger group finalization.