NVIDIA-NeMo / NVIDIA-NeMo/RL

async PPO: startup barrier still requires a complete lookahead, serializing every resume

Open
#3,828 0 comments 0 reactions 1 assignee Claimed by @terrykong View on GitHub
Feature
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Summary

Async GRPO's checkpoint-resume startup barrier was relaxed in #3435 so that training can begin once the lookahead target is **claimed** by the collector, rather than waiting for it to be **complete**. Async PPO carries a structurally identical barrier that was not changed, so every async PPO resume still serializes a full rollout round before training can start.

## Background

#3066 introduced the barrier to fix a resume deadlock. After restoring a checkpoint where target `N` is complete but `N+1` is absent, startup would consume `N`, train, and refit — at which point the collector's target window advances to `[N+2, ...]` and `N+1` is left with no producer, stalling the run forever. The barrier fixed that by requiring `N+1` to be complete before training begins.

#3435 observed that the barrier only needs to know `N+1` has a *producer*, not that it is finished, and relaxed the GRPO condition to `has_complete_batch(N+1) or N+1 in collector_status["generating_targets"]`.

## Where async PPO stands today

`async_ppo_train` still has the pre-#3435 form:

https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/algorithms/ppo.py#L2343-L2367

```python
if current_step_ready:
# The initial collector is the only window that can generate
# both `step` and `step + 1`. Fill both before the first refit.
need_lookahead = step + 1 < max_training_steps
if need_lookahead:
lookahead_step_ready = ray.get(
replay_buffer.has_complete_batch.remote(
step + 1, num_prompts_per_step, initial_buffer_max_age
)
)
if not lookahead_step_ready:
...
time.sleep(1.0)
continue
break
```

This is not dead code on the resume path: async PPO restores its replay buffer from the checkpoint at [`ppo.py:2232-2249`](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/algorithms/ppo.py#L2232-L2249) (via `load_from_path`), so the "target `N` complete, `N+1` absent" state that motivated the barrier is reachable there too.

Both algorithms share `AsyncTrajectoryCollector`, so the `generating_targets` key added to `get_status()` in #3435 is already available to PPO — only the driver-side condition differs.

## Why this is not a drop-in port

`_startup_pipeline_ready` cannot simply be called from `async_ppo_train` as-is:

1. **Different age parameter.** PPO passes `initial_buffer_max_age`, computed by `_async_ppo_buffer_max_age` from the critic-warmup schedule, rather than a flat `max_trajectory_age_steps`.
2. **Different guard.** PPO's `need_lookahead` omits the `max_trajectory_age_steps > 0` term entirely.
3. **Warmup lead.** PPO can configure `warmup_generation_lead_steps` independently of `max_trajectory_age_steps`. The "a claim implies a producer" argument rests on the reserved worker delivering to its reserved `target_weight_version` across the refit boundary; that needs re-deriving under a warmup generation lead before it can be assumed safe.

Also note PPO's version already calls `_raise_if_collector_stopped(...)` *before* its `continue`, so it never had the GRPO bug where the `continue` skipped the exhausted/errored check and hung — that half of #3435 does not apply here.

## Proposed work

- Decide whether async PPO should adopt the claim-aware condition, or deliberately stay serialized (the barrier costs one rollout round per resume, which may be an acceptable price for the warmup path).
- If adopting: generalize `_startup_pipeline_ready` to take the effective max-age and lookahead predicate as parameters, and verify the claim argument holds under `warmup_generation_lead_steps`.
- If not adopting: add a short comment at `ppo.py:2343` recording the deliberate divergence, so the two barriers do not silently drift.
- Either way, PPO's barrier currently has no unit test; GRPO's is covered by `test_startup_pipeline_ready` and `test_startup_barrier_clears_when_collector_reserves_lookahead`.

Follow-up to #3435. Related: #3066.

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.