fix(async): recover from permanently-failed prompt groups instead of hanging - #3938
Conversation
…hanging
When a prompt group fails permanently (for example, a prompt that a
generation engine deterministically rejects on every retry), its target
weight version is left short in the replay buffer. Two collector behaviors
then combine into a deadlock:
1. _calculate_target_weights starts the target window at
generation_weight_version + 1, so a short target at or below the current
version is invisible to gap-filling: no worker is ever dispatched to
regenerate the missing trajectories.
2. _collection_loop parks on an unbounded _generation_limit_cleared.wait().
The event is only set by a weight update, but training cannot step
without the missing trajectories, so the wait never returns.
The trainer then sits at the buffer wait forever ("Wait iteration N" with a
constant buffer size) while all generation nodes idle. Observed in
production on an async GRPO VLM run where one over-long prompt exceeded
max_model_len on every attempt.
Fix both halves:
- Extend the target window down to the first target training has not
consumed (min(current + 1, last_consumed + 1)). Both callers already skip
targets at or below last_consumed_target, so only genuine holes surface
and steady-state behavior is unchanged.
- Bound the generation-limit pause with a 60s re-check of
_should_pause_for_generation_limits so the collector wakes up and
gap-fills the shortfall with fresh prompts instead of deadlocking.
Signed-off-by: Pulkit Kumar <pulkitk@nvidia.com>
| # Start at the first target training has not consumed instead. Callers | ||
| # still skip targets <= last_consumed_target, so only genuine holes | ||
| # surface. | ||
| last_consumed_target = ray.get( |
There was a problem hiding this comment.
Could we keep _calculate_target_weights() free of Ray RPCs and pass last_consumed_target in from the callers instead? Both _get_next_target_for_generation() and _should_pause_for_generation_limits() already fetch this value immediately afterward, so this currently adds a duplicate actor round-trip on each path.
It also appears to break the existing direct unit tests: test_collector_selects_ppo_config and test_collector_grpo_window_remains_fixed construct the collector with an unconfigured MagicMock replay buffer and call this method for a non-initial weight, causing ray.get() to receive a mock rather than an ObjectRef. Please update/add tests for the widened window, including a short target at or below the current generation version.
| while ( | ||
| not self._generation_limit_cleared.is_set() and self.running | ||
| ): | ||
| self._generation_limit_cleared.wait(timeout=60.0) |
There was a problem hiding this comment.
Could we close the event race with an immediate condition re-check after _generation_limit_cleared.clear() instead of relying on the 60-second timeout?
_run_rollout_batch_worker() already sets this event when a rollout batch fails. A wake can be lost if that set() happens after the first _should_pause_for_generation_limits() check but before the collection loop calls clear(). Rechecking immediately after clear() would detect the released short target without adding up to 60 seconds of idle time. A timeout can remain as a safety net, but it should not be the primary recovery mechanism. Please add a regression test for this exact check/set/clear interleaving.
…ests - _calculate_target_weights no longer performs a Ray RPC: callers fetch last_consumed_target once and pass it in, removing the duplicate actor round-trip and keeping the helper unit-testable against mock buffers. - Extract the generation-limit pause into _pause_for_generation_limits and make the immediate post-clear() condition re-check the primary recovery path: a worker's set() landing between the pause check and the clear() is no longer lost. The timed re-check remains as a safety net only, with the interval lifted to a module constant. - Update the two window tests for the new signature (steady-state last_consumed_target keeps their expectations unchanged) and add regression tests for: the widened window over unconsumed failed targets, the set()-racing-clear() interleaving, the timeout re-check release, and the normal event-set wake path. Signed-off-by: Pulkit Kumar <pulkitk@nvidia.com>
|
@aroshanghias-nvd Addressed in 1f6a0b3 — window calc is now RPC-free with the value passed from callers, the pause is extracted with an immediate post-clear re-check as the primary recovery (timeout kept as safety), and both broken tests are fixed plus four regression tests added for the widened window and the race |
2 participants
What does this PR do?
Fixes a deadlock in the async trajectory collector: when a prompt group fails permanently, training hangs forever at the buffer wait while every generation node idles.
The deadlock
A prompt group that fails on every retry (for example, a prompt that the generation engine deterministically rejects because it exceeds
max_model_len) leaves its target weight version short in the replay buffer. Two collector behaviors then combine:_calculate_target_weightsstarts the target window atgeneration_weight_version + 1. A short target at or below the current weight version is therefore never returned, andlast_consumed_targetis only used by the callers to skip targets — never to widen the window — so no worker is ever dispatched to regenerate the missing trajectories._collection_loopparks on an unbounded_generation_limit_cleared.wait(). That event is set by a weight update — but training cannot step without the missing trajectories, so the update never comes.Observed in production on a 32-node async GRPO VLM run: one over-long prompt failed on every attempt, and the trainer sat at
Wait iteration Nwith a frozen buffer size for the rest of the job while 100+ GPUs idled.The fix
_calculate_target_weights: extend the window down to the first target training has not consumed —target_start = min(current + 1, last_consumed+ 1). Both callers already skip targets<= last_consumed_target, so only genuine holes surface and steady-state behavior is unchanged._collection_loop: bound the generation-limit pause with a 60-second re-check of_should_pause_for_generation_limits(), so the collector wakes up and gap-fills the shortfall with fresh prompts instead of deadlocking. The fast path (event set promptly by a weight update) is unchanged.Validation
Deployed on the failing production run: the previously-hung workload recovered the failed group via gap-fill and trained 50 steps to completion across multiple resumed windows, with no recurrence of the stall. Steady-state runs (no failed groups) show identical scheduling behavior, since the widened window only produces targets the callers would otherwise skip.