feat(ppo): support async ppo

.agents/contributor-skills/config-conventions/SKILL.md

@@ -33,7 +33,7 @@ Use the right tool for the job. **v2 (the new convention):**
 
 **v1 (legacy, being migrated away):**
 
-- **`typing.TypedDict` — v1, legacy / not-yet-migrated user-facing config.** Most nested sub-configs (e.g. `GRPOConfig`, `RewardScalingConfig`, `AsyncGRPOConfig`) are still `TypedDict`. Continue to maintain them with the same defaults rules below until they are migrated to `BaseModel`. Use `typing.NotRequired` to mark optional attributes. **Do not add new `TypedDict`-based config classes.**
+- **`typing.TypedDict` — v1, legacy / not-yet-migrated user-facing config.** Some nested sub-configs are still `TypedDict`. Continue to maintain them with the same defaults rules below until they are migrated to `BaseModel`. Use `typing.NotRequired` to mark optional attributes. **Do not add new `TypedDict`-based config classes.**
 
 When in doubt: *is this class populated from a user-edited YAML?* If yes → `BaseModel` (or legacy `TypedDict`). If no → `@dataclass`.
 

docs/guides/ppo.md

@@ -57,6 +57,18 @@ policy:
 
 When only one node remains for policy and generation after other resources are reserved, `gpus_per_node` reserves that many GPUs for generation and `num_nodes` must be `null` or `1`. When more than one node remains for training and generation, generation uses complete nodes: set `num_nodes` to the number of inference nodes and `gpus_per_node` equal to `cluster.gpus_per_node`. Non-colocated SGLang generation is not currently supported by PPO.
 
+### Asynchronous PPO
+
+Set `ppo.async_ppo.enabled: true` to overlap rollout generation with training. A background collector fills a replay buffer on the non-colocated vLLM GPUs while the policy and value model train on their shared cluster. Values and policy/reference log probabilities are recomputed when a trajectory is sampled, then PPO runs GAE and its normal `ppo_epochs` updates before publishing one new policy version to vLLM.
+
+Async PPO reuses the trajectory collector, replay buffer, and weight-versioning infrastructure described in the [Async GRPO guide](async-grpo.md); this section focuses on PPO-specific behavior and constraints.
+
+Async PPO requires non-colocated vLLM generation with `vllm_cfg.async_engine: true`, `loss_fn.use_importance_sampling_correction: true`, and `loss_fn.force_on_policy_ratio: false`. Dynamic sampling, reward scaling, reward shaping, multiple dataloaders, NeMo Gym, colocated generation, and FP8 KV-scale synchronization are not supported yet.
+
+`max_trajectory_age_steps` is the normal policy-training age limit. The recommended value is `1`; larger values improve overlap but increase off-policy bias in GAE. When `policy_training_start_step > 0`, set `warmup_generation_lead_steps` to a larger value to bank additional rollout batches while the policy is frozen for critic warmup. The collector caps frozen-policy targets at `policy_training_start_step + max_trajectory_age_steps`, so their actual policy-update age remains within the normal limit. The buffer keeps these batches valid through that frontier and then restores the normal age limit. `null` uses `max_trajectory_age_steps` as the generation lead throughout.
+
+Async training stops at `max_num_steps`; the collector cycles the training dataloader as needed. `max_num_epochs` is not supported yet and must be set to `-1`; use `max_num_steps` to control training length. Async checkpoints save the collector dataloader and replay-buffer state together with policy and value state. By default, incomplete restored targets are retained and gap-filled. Setting `drop_incomplete_targets_on_restore: true` discards their restored rows and fills the target from subsequent dataloader prompts; it does not regenerate the original prompts.
+
 ### Value Model Configuration
 
 ```yaml
@@ -235,6 +247,14 @@ ppo:
   # null logs mismatch metrics without masking; set a threshold to mask sequences.
   seq_logprob_error_threshold: null
 
+  async_ppo:
+    enabled: false
+    max_trajectory_age_steps: 1
+    warmup_generation_lead_steps: null
+    in_flight_weight_updates: false
+    recompute_kv_cache_after_weight_updates: false
+    drop_incomplete_targets_on_restore: false
+
   adv_estimator:
     name: "gae"
     gae_lambda: 0.95
@@ -275,14 +295,15 @@ value_loss_fn:
 - **`ppo.ppo_epochs`**: Number of training updates per rollout batch
 - **`ppo.policy_training_start_step`**: Number of critic-only warmup steps before policy training begins
 - **`ppo.seq_logprob_error_threshold`**: Nullable sequence-level multiplicative probability-error threshold. PPO always logs sequence-level train/generation mismatch metrics; when this is set, sequences above the threshold are excluded from advantage and loss computation.
+- **`ppo.async_ppo`**: Enables replay-buffer-based asynchronous PPO. See [Asynchronous PPO](#asynchronous-ppo) for requirements and staleness controls.
 - **`ppo.adv_estimator.name`**: Set to `"gae"` for GAE advantage estimation (PPO default)
 - **`ppo.adv_estimator.gae_lambda`**: GAE $\lambda$ parameter (bias-variance tradeoff, typically 0.95)
 - **`ppo.adv_estimator.gae_gamma`**: Discount factor $\gamma$ (typically 1.0 for outcome-supervised tasks)
 - **`value_loss_fn.scale`**: Scaling factor for the value loss
 - **`value_loss_fn.cliprange`**: Clip range for value function predictions
 - **`loss_fn.positive_example_nll_weight`**: VAPO NLL auxiliary loss weight on correct samples (0 = disabled)
 
-All other parameters (clipping, KL, importance sampling, dynamic sampling, reward shaping, reward scaling) work identically to GRPO. See the [GRPO Guide](grpo.md) for details.
+For synchronous PPO, the remaining clipping, KL, sampling, and reward options work as documented in the [GRPO Guide](grpo.md). Async PPO has the limitations listed above.
 
 ## Metrics
 

examples/configs/ppo_math_1B.yaml

@@ -22,6 +22,17 @@ ppo:
   batch_multiplier: 1
   skip_reference_policy_logprobs_calculation: true  # No KL, so skip ref logprobs
 
+  # Async PPO requires non-colocated vLLM generation and importance correction.
+  async_ppo:
+    enabled: false
+    max_trajectory_age_steps: 1
+    # Requires policy_training_start_step > 0. null uses max_trajectory_age_steps.
+    warmup_generation_lead_steps: null
+    in_flight_weight_updates: false
+    recompute_kv_cache_after_weight_updates: false
+    # true discards partial restored rows and fills from subsequent prompts.
+    drop_incomplete_targets_on_restore: false
+
   reward_shaping:
     enabled: true
     overlong_buffer_length: 2048
@@ -443,6 +454,7 @@ data:
   max_input_seq_length: 2048
   shuffle: true
   num_workers: 1
+  use_multiple_dataloader: false
   train:
     dataset_name: DAPOMath17K
   validation:

examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated-async.yaml

@@ -0,0 +1,67 @@
+defaults: ../../ppo_math_1B.yaml
+ppo:
+  num_prompts_per_step: 1024
+  num_generations_per_prompt: 1
+  # Async PPO cycles the dataloader and stops via max_num_steps.
+  max_num_epochs: -1
+  ppo_epochs: 1
+  policy_training_start_step: 5
+  val_period: 1
+  overlong_filtering: true
+  async_ppo:
+    enabled: true
+    warmup_generation_lead_steps: 2
+    in_flight_weight_updates: true
+  reward_shaping:
+    enabled: false
+  adv_estimator:
+    gae_lambda_value: 1.0
+    gae_lambda_policy: 1
+  reward_scaling:
+    enabled: false
+loss_fn:
+  ratio_clip_max: 0.2
+  ratio_clip_c: 3
+  use_importance_sampling_correction: true
+value_loss_fn:
+  scale: 1.0
+  cliprange: 0.5
+checkpointing:
+  checkpoint_dir: results/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated-async
+policy:
+  model_name: Qwen/Qwen2.5-1.5B-Instruct
+  train_global_batch_size: 256
+  max_total_sequence_length: 1024
+  generation:
+    max_new_tokens: 512
+    vllm_cfg:
+      async_engine: true
+      gpu_memory_utilization: 0.4
+      max_model_len: 1024
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+value:
+  model_name: Qwen/Qwen2.5-1.5B-Instruct
+  train_micro_batch_size: 4
+data:
+  max_input_seq_length: 512
+  train:
+    dataset_name: gsm8k
+    split: train
+  validation:
+    dataset_name: gsm8k
+    split: test
+  default:
+    system_prompt_file: examples/prompts/gsm8k.txt
+env:
+  math:
+    math_verify_impl: hf_math_verify
+logger:
+  log_dir: logs/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated-async
+  wandb:
+    project: nemo-rl
+    name: ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated-async
+cluster:
+  gpus_per_node: 8

examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async.yaml

@@ -0,0 +1,89 @@
+defaults: ../../ppo_math_1B_megatron.yaml
+ppo:
+  num_prompts_per_step: 1024
+  num_generations_per_prompt: 1
+  # Async PPO cycles the dataloader and stops via max_num_steps.
+  max_num_epochs: -1
+  ppo_epochs: 1
+  val_period: 1
+  overlong_filtering: true
+  async_ppo:
+    enabled: true
+    in_flight_weight_updates: true
+  reward_shaping:
+    enabled: false
+  adv_estimator:
+    gae_lambda_value: 1.0
+    gae_lambda_policy: 1
+  reward_scaling:
+    enabled: false
+loss_fn:
+  ratio_clip_max: 0.2
+  ratio_clip_c: 3
+  use_importance_sampling_correction: true
+value_loss_fn:
+  scale: 1.0
+  cliprange: 0.5
+checkpointing:
+  checkpoint_dir: results/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async
+policy:
+  model_name: Qwen/Qwen2.5-1.5B-Instruct
+  train_global_batch_size: 256
+  max_total_sequence_length: 1024
+  megatron_cfg:
+    tensor_model_parallel_size: 2
+    context_parallel_size: 2
+    sequence_parallel: true
+    optimizer:
+      weight_decay: 0.01
+    scheduler:
+      start_weight_decay: 0.01
+      end_weight_decay: 0.01
+      lr_warmup_iters: 0
+      lr_warmup_init: 0
+  make_sequence_length_divisible_by: 8
+  generation:
+    max_new_tokens: 512
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 8
+        num_nodes: 1
+    vllm_cfg:
+      async_engine: true
+      gpu_memory_utilization: 0.4
+      max_model_len: 1024
+value:
+  model_name: Qwen/Qwen2.5-1.5B-Instruct
+  train_micro_batch_size: 4
+  megatron_cfg:
+    tensor_model_parallel_size: 2
+    sequence_parallel: true
+    optimizer:
+      lr: 1.0e-05
+      weight_decay: 0.01
+    scheduler:
+      lr_warmup_iters: 0
+  dynamic_batching:
+    enabled: true
+data:
+  max_input_seq_length: 512
+  train:
+    dataset_name: gsm8k
+    split: train
+  validation:
+    dataset_name: gsm8k
+    split: test
+  default:
+    system_prompt_file: examples/prompts/gsm8k.txt
+env:
+  math:
+    math_verify_impl: hf_math_verify
+logger:
+  log_dir: logs/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async
+  wandb:
+    project: nemo-rl
+    name: ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async
+cluster:
+  num_nodes: 2
+  gpus_per_node: 8

examples/nemo_gym/run_distillation_nemo_gym.py

@@ -28,11 +28,13 @@
     distillation_train,
     setup,
 )
-from nemo_rl.algorithms.grpo import _should_use_nemo_gym
 from nemo_rl.algorithms.utils import get_tokenizer
 from nemo_rl.data.utils import setup_response_data
 from nemo_rl.distributed.virtual_cluster import init_ray
-from nemo_rl.environments.nemo_gym import setup_nemo_gym_config
+from nemo_rl.environments.nemo_gym import (
+    setup_nemo_gym_config,
+    should_use_nemo_gym,
+)
 from nemo_rl.models.generation import configure_generation_config
 from nemo_rl.utils.config import (
     load_config,
@@ -105,7 +107,7 @@ def main() -> None:
     setup_nemo_gym_config(config, tokenizer)
 
     # We assert here since this is right after the final config has been materialized.
-    assert _should_use_nemo_gym(config)
+    assert should_use_nemo_gym(config)
 
     # NeMo-Gym environment needs to get dp_openai_server_base_urls from
     # student_generation, so we don't setup env here.

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -33,7 +33,6 @@
     MasterConfig,
     StatefulDataLoader,
     TokenizerType,
-    _should_use_nemo_gym,
     grpo_train,
     refit_policy_generation,
     setup,
@@ -42,7 +41,10 @@
 from nemo_rl.algorithms.utils import get_tokenizer
 from nemo_rl.data.utils import setup_response_data
 from nemo_rl.distributed.virtual_cluster import init_ray
-from nemo_rl.environments.nemo_gym import setup_nemo_gym_config
+from nemo_rl.environments.nemo_gym import (
+    setup_nemo_gym_config,
+    should_use_nemo_gym,
+)
 from nemo_rl.experience.rollouts import run_nemo_gym_rollout_sync
 from nemo_rl.models.generation import configure_generation_config
 from nemo_rl.utils.config import (
@@ -188,7 +190,7 @@ def main() -> None:
         setup_nemo_gym_config(config, tokenizer)
 
     # We assert here since this is right after the final config has been materialized.
-    assert _should_use_nemo_gym(config)
+    assert should_use_nemo_gym(config)
 
     # NeMo-Gym environment needs to get dp_openai_server_base_urls from policy_generation, so we don't setup env here.
     with rl_init_timer.time("data"):