NVIDIA-NeMo / NVIDIA-NeMo/RL

Async GRPO deadlocks at the end of a dataset epoch: the trajectory collector iterates the dataloader only once and is never restarted

Open
#2,899 0 comments 1 reaction 1 assignee Claimed by @terrykong View on GitHub
bug community-request waiting-on-maintainers
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Bug

In async GRPO (`grpo.async_grpo.enabled=true`), generation runs in a single background thread, `AsyncTrajectoryCollector._collection_loop`, which iterates the dataloader **exactly once** (`for batch in self.dataloader`). When the dataset's single pass is exhausted, the loop ends, the `finally` sets `self.running = False`, prints `🛑 Trajectory collection stopped`, and the thread dies. Nothing ever restarts it (no watchdog; `start_collection` is called once and returns immediately).

But training is bounded by `max_num_steps`, which is normally far larger than one epoch. So the collector dies, the replay buffer stops being fed, and the training driver spins forever in its "waiting for samples" stall — a silent, permanent deadlock with all GPUs idle and no error raised.

This diverges from the **synchronous** path, which wraps iteration in an outer epoch loop (`grpo.py:1674`) and re-iterates the dataset across epochs.

## Why it's a hard deadlock

The final partial step is unrecoverable: `replay_buffer.sample()` (`replay_buffer.py:163-181`) requires **exactly** `num_prompt_groups` trajectories with `target_weight_version == current`. When the collector dies mid-step, its in-flight rollouts are dropped (the enqueue guard `while self.running:` at `trajectory_collector.py:479` is now `False`), and the missing trajectories can never be regenerated because weights have already advanced. `sample()` returns `None` forever.

## Key locations (verified on `main`)

- `nemo_rl/algorithms/async_utils/trajectory_collector.py:189` — `_collection_loop`: single `for batch in self.dataloader`, no outer epoch loop; `finally` kills the thread (`:243-244`).
- `nemo_rl/algorithms/grpo.py:330` — dataloader is a single-pass `StatefulDataLoader(..., drop_last=True)`, no cycling.
- `nemo_rl/algorithms/grpo.py:1674` — synchronous path's outer epoch loop (the behavior async is missing).

## Reproduce

Run async GRPO (`max_trajectory_age_steps=1`) with a finite dataset and `max_num_steps > len(dataloader)`; let it pass the end of epoch 1. (Resuming from a checkpoint near the epoch boundary triggers it within a few steps.) Result: `🛑 Trajectory collection stopped` prints once, then the `STALLING` log repeats indefinitely.

## Observed logs

```
🎯 Found 49 trajectories intended for current step 57
⏸️ STALLING: Need 64 trajectories for step 57, but only 49 are ready
(AsyncTrajectoryCollector) 🛑 Trajectory collection stopped <-- thread exits, never restarts
========================= Step 58/1000000 ========================= <-- STALLING now repeats forever
```

No traceback precedes the stop → the loop ended normally (dataset exhausted), not via an error.

## Expected / proposed fix

The collector should keep supplying trajectories until `running` is cleared or `max_num_steps` is reached. Wrap the loop in an outer `while self.running:` that re-iterates (re-shuffles) the dataset on exhaustion, mirroring the sync path. Optional hardening: drain in-flight rollouts before stopping, and have training detect a dead collector )`) instead of hanging silently.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.