Async GRPO deadlocks at the end of a dataset epoch: the trajectory collector iterates the dataloader only once and is never restarted
- 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
Assessment
This issue has not been assessed yet.