feat(ppo): support async ppo

[bg51717]

What does this PR do?

Based on #3185.

Adds asynchronous PPO on top of non-colocated PPO, overlapping rollout generation on dedicated vLLM GPUs with policy and value-model training on the training cluster.

This is the asynchronous follow-up to #3262. The existing synchronous colocated and non-colocated PPO paths remain supported.

Design

The main changes are:

• Extend the shared asynchronous rollout collector and replay buffer to support Async PPO while preserving the existing Async GRPO behavior.
• Generate trajectories continuously on the inference cluster while the policy and value models train on the separate training cluster.
• Recompute values and policy/reference logprobs when a trajectory is sampled, then run GAE and the configured PPO epochs against the current learner state.
• Track the policy version used for generation and bound off-policy staleness with max_trajectory_age_steps.
• Allow the collector to bank additional rollouts during critic warmup with warmup_max_trajectory_age_steps, while keeping the effective policy-update age within the normal window once policy training starts.
• Publish updated policy weights to vLLM after each policy-training step, with support for in-flight weight updates and explicit synchronization before checkpointing or shutdown.
• Save and restore the replay buffer, collector dataloader state, policy, and value model so asynchronous training can resume without losing completed trajectories.
• Regenerate incomplete rollout targets after restore by default instead of treating partially generated data as complete.
• Cycle the training dataloader until max_num_steps is reached so asynchronous progress is controlled by learner steps rather than dataloader exhaustion.
• Migrate PPOConfig from the legacy TypedDict representation to a Pydantic BaseModel, preserving existing defaults and extra-field compatibility while aligning PPO configuration access with GRPO.
• Fail fast for unsupported combinations, including colocated generation, non-vLLM generation, synchronous vLLM engines, dynamic sampling, reward scaling or shaping, multiple dataloaders, NeMo Gym, and FP8 KV-cache scale synchronization.
• Add unit coverage, AutoModel and Megatron functional checkpoint/resume tests, Async GRPO compatibility coverage, 40-step nightly recipes, and PPO documentation.

Usage

Enable asynchronous PPO together with non-colocated asynchronous vLLM generation and importance-sampling correction:

ppo:
  max_num_epochs: -1
  async_ppo:
    enabled: true
    max_trajectory_age_steps: 1
    warmup_max_trajectory_age_steps: null
    in_flight_weight_updates: true
    recompute_kv_cache_after_weight_updates: false
    drop_incomplete_targets_on_restore: true
  reward_scaling:
    enabled: false
  reward_shaping:
    enabled: false

loss_fn:
  use_importance_sampling_correction: true

policy:
  generation:
    backend: vllm
    vllm_cfg:
      async_engine: true
    colocated:
      enabled: false
      resources:
        gpus_per_node: 4
        num_nodes: null

data:
  use_multiple_dataloader: false

max_trajectory_age_steps=1 is the recommended default. When critic warmup is enabled, warmup_max_trajectory_age_steps can be increased to let the frozen policy bank more rollouts before policy updates begin.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests.
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Test

The following test matrix passed:

• Unit tests covering PPO asynchronous configuration and validation, trajectory-age windows, critic-warmup rollout banking, replay-buffer restore, dataloader cycling, weight-update coordination, failure propagation, and cleanup.

• Async GRPO replay-buffer checkpoint compatibility tests passed after the shared asynchronous utility changes.

• Synchronous colocated and non-colocated PPO functional regression tests passed on both AutoModel and Megatron.

• AutoModel and Megatron asynchronous PPO functional tests passed, covering critic warmup, replay-buffer population, consistent critic-loss sanity checks, policy/value checkpointing, restore, and continued training after resume.

• AutoModel (green, with policy_training_start_step=5 and warmup_max_trajectory_age_steps=2) and Megatron (red, without critic warmup) 40-step Async PPO nightly runs both passed. correctness gates for reward, validation accuracy and train/generation logprob mismatch.

• A matched 40-step comparison between synchronous non-colocated Megatron PPO(blue) and Async Megatron PPO(yellow) used the same two-node training/generation split and showed no obvious regression in reward, validation accuracy, or train/generation logprob correctness.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[bg51717]

/ok to test b8776c9

[yuki-97]

Thanks @bg51717 — a lot of careful work here, and the warmup lead/age design in particular is subtler than it looks and holds up. Picking up async_ppo_train and async_utils/, the part I deferred on the earlier pass.

Two things I traced end to end, both clean:

• Critic-warmup → normal-training transition. Walked the version/lead/window schedule across the frontier with a warmup config where lead and buffer age actually diverge. No data discard and no deadlock: consumption age never exceeds buffer age (exactly equal at the two frontier targets), the wakeup path can't miss a version bump, peak occupancy (warmup_age+1)·N stays under the 2·warmup_age·N capacity, and the one dangerous shape — evicting rows for a target the collector can no longer regenerate, since past the frontier its window is only [v+1] — never triggers. The <= in step <= warmup_rollout_frontier is what guarantees that last one, and test_async_ppo_consumes_frozen_policy_rollout_at_safe_warmup_frontier already pins it.

• Policy / value GPU residency. Matches sync ppo_train stage for stage — value forward, logprobs, the per-epoch value and policy training pairs, the refit offload, and both checkpoint saves. Every load is paired, both models are on CPU at step boundaries, and the conditional skip of the post-refit offload is correct because the policy is never loaded during warmup.

[bg51717]

/ok to test bac4a8c

[yuki-97]

thanks for the updates, overall LGTM, just some nits and one nightly coverage.

[bg51717]

/ok to test 6e0362d

[terrykong]

Deep review by a team of 6 agents (RL/config, bug-finding, tests, design, comment audit) plus an adversarial pass that dropped most candidate findings. Everything below survived that pass and was re-verified against 6e0362d after your latest push.

First — this is a large, careful piece of work, and the review it has already had is excellent. I checked every thread marked fixed in <sha> and they all genuinely hold. Your three recent commits also resolved several things I had queued up, so I dropped them: the #3611 carry-over in should_use_async_rollouts (correct now — asserts and returns True), the TODO(#2625) rationale, the in_flight_weight_updates exemplar/docs default, the max_trajectory_age_steps > 1 warning, and the avg_trajectory_age nightly bound — where you also added the companion ["40"] <= 1.0 assertion and pointed the AutoModel nightly at a real warmup config. That last one closes the coverage gap that mattered most.

Two items are worth your attention before merge.

1. drop_incomplete_targets_on_restore — suggest removing it. Making the parameter required was the right call, but I think the option itself should go. true deletes restored rollouts that false keeps, and both branches then train on the same number of prompt groups, so the deletion only costs dataset coverage and regeneration compute — what it keeps is a strict subset of what the other branch keeps. It also cannot be a safety measure, because it leaves complete restored steps untouched. main and async GRPO both behave as false, and nothing sets the key, so every async PPO run today takes the deleting branch. Walkthrough with figures: https://terrykong.github.io/gh-pages-poc/terryk/pr-3410-resume-drops-batch.html — and if you have a case where deleting beats gap-filling, please say so and I will withdraw it.

2. The name of warmup_max_trajectory_age_steps. It reads as a staleness bound, but I verified it cannot affect staleness at all — it is a concurrency knob. Renaming is free right now (0 references on main) and a breaking change after merge.

⚠️ Still shows CONFLICTING against main, so one more rebase is needed — but the specific #3611 hazard I was going to flag is already handled.

Verification notes: ruff, ruff-format, taplo and minimize-check pass. pyrefly could not run here (the uv lockfile is linux-only; this is macOS) — an environment limit, not a finding. No test was executed (no GPU, no runnable env), so please treat the test suggestions as reviewed-by-reading only. GPU-hour figures were measured with tools/launch's own formula, not estimated.

Generated by Claude Code

[bg51717]

/ok to test 9182c74