Train with Async GRPO#

Async GRPO is an asynchronous training mode that allows trajectory generation and policy training to run concurrently, improving GPU utilization and throughput compared to synchronous GRPO.

Configure Async GRPO#

This section covers how to configure async GRPO by modifying your settings and includes a complete example configuration.

Enable Async GRPO#

To use async GRPO, make these configuration changes:

  1. Enable vLLM async engine:

policy:
  generation:
    backend: "vllm"
    vllm_cfg:
      async_engine: true

  1. Enable importance sampling correction (required for convergence):

loss_fn:
  use_importance_sampling_correction: true

  1. Disable colocated inference (required for async mode with the vLLM backend; the Megatron backend supports colocated async — see examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml):

policy:
  generation:
    colocated:
      enabled: false
      resources:
        num_nodes: 1  # or more
        gpus_per_node: 2  # adjust based on your setup

  1. Add async GRPO configuration:

grpo:
  async_grpo:
    enabled: true
    max_trajectory_age_steps: 1  # Maximum age, in training steps, for trajectories
    max_generation_failures: 0  # Consecutive worker failures to tolerate
    in_flight_weight_updates: false  # Enable for faster weight synchronization
    recompute_kv_cache_after_weight_updates: false # Invalidates kv cache after weight-updates

Complete Example Config#

policy:
  generation:
    backend: "vllm"
    colocated:
      enabled: false
      resources:
        num_nodes: 1
        gpus_per_node: 2
    vllm_cfg:
      async_engine: true

loss_fn:
  use_importance_sampling_correction: true

grpo:
  num_prompts_per_step: 32
  num_generations_per_prompt: 4
  async_grpo:
    enabled: true
    max_trajectory_age_steps: 1
    max_generation_failures: 0  # Consecutive worker failures to tolerate
    in_flight_weight_updates: false  # Enable for faster weight synchronization
    recompute_kv_cache_after_weight_updates: false # Invalidates kv cache after weight-updates

cluster:
  num_nodes: 2
  gpus_per_node: 4

Implementation Structure#

This section covers the internal architecture of async GRPO and includes detailed explanations of how the core components interact.

Core Components#

The async GRPO implementation consists of three main components:

1. Main Training Loop (async_grpo_train in grpo.py)#

  • Coordinates overall training process

  • Samples trajectories from replay buffer

  • Runs policy training steps

  • Handles validation and checkpointing

  • Manages weight synchronization between training and generation

2. Async Trajectory Collector (AsyncTrajectoryCollector in async_utils/trajectory_collector.py)#

  • Runs in background Ray actor

  • Continuously generates trajectories using current policy weights

  • Manages generation scheduling and weight version tracking

  • Handles pause/resume for weight updates and validation

  • Coordinates with replay buffer for trajectory storage

3. Replay Buffer (ReplayBuffer in async_utils/replay_buffer.py)#

  • Stores generated trajectories with metadata

  • Tracks weight versions for both generation and intended training use

  • Implements age-based filtering to prevent stale trajectories

  • Provides sampling interface for training steps

Weight Version Tracking#

Async GRPO uses a weight versioning system:

  • Generation Weight Version: The policy weights used to generate a trajectory

  • Target Weight Version: The training step where the trajectory will be used

  • Max Trajectory Age: How many steps old a trajectory can be before being discarded

Example with max_trajectory_age_steps: 1:

  • Trajectory generated with weights v10 can be used for training steps v10 or v11

  • At training step v12, trajectories from v10 are too old and discarded

Coordination Flow#

  1. Startup: Trajectory collector starts generating trajectories in background

  2. Buffer Fill: Training waits until buffer has sufficient trajectories

  3. Training Step:

    • Sample trajectories from buffer

    • Run policy training

    • Update weights and notify collector

  4. Weight Sync: Collector pauses, waits for weight refit, then resumes

  5. Repeat: Process continues with updated weights

Architecture Diagram#

The following sequence diagram illustrates the interactions between the three main components:

        sequenceDiagram
    participant Training as Training Loop
    participant Collector as Trajectory Collector
    participant Buffer as Replay Buffer
    
    Note over Training, Buffer: Startup
    Training->>Collector: Start generation
    Training->>Buffer: Initialize
    
    Note over Training, Buffer: Main Loop
    loop Async Training
        par Background Generation
            Collector->>Buffer: Store trajectories
        and Training Steps
            Training->>Buffer: Sample trajectories
            Buffer-->>Training: Return valid data
            Training->>Training: Update policy weights
            Training->>Collector: Sync new weights
        end
    end

Checkpointing#

Async GRPO checkpoints the replay buffer alongside the rest of training state so that in-progress trajectory generation is not lost across restarts.

What is saved#

On each checkpoint, a replay_buffer.pt file is written next to the other checkpoint artifacts. It contains all trajectories currently in the buffer together with their weight and target versions, and the last_target_weight_already_generated watermark.

The dataloader state saved alongside it is frontier-aligned: it captures the position of the first prompt that training has not yet consumed, rather than the collector’s live cursor (which has already advanced past prompts that are still generating or sitting in the buffer). Every yielded prompt carries a monotonic ordinal, and the saved replay buffer records which ordinals its groups cover.

Restore behaviour#

On resume, the buffer is restored before the trajectory collector starts, then cleaned up as follows:

  1. Past targets dropped — trajectories whose target step is earlier than the resume step are removed.

  2. Stale trajectories evicted — if max_trajectory_age_steps is set, trajectories too old for their target step are removed.

  3. Incomplete targets kept — target steps that still lack a full batch are kept in the buffer. The collector will gap-fill only the missing trajectories for those targets before moving on.

  4. Buffer truncated — if the restored count exceeds max_size, the buffer is truncated, prioritising entries closest to the resume step.

The dataloader then re-yields the window between the trained frontier and the old cursor. Prompts already trained or retained in the restored buffer are dropped; everything else — in-flight work lost at the save, groups evicted during cleanup — is regenerated in order, so no prompt is skipped and none is duplicated as a result of the save/restore itself (prompts trained above a conservatively-cut checkpoint’s threshold — see the target-interleaving note below — are carried in the checkpoint and stay covered on resume).

The guarantee is scoped to frontier-aligned checkpoints. A run falls back to the previous live-cursor checkpoints — which can skip prompts that were in flight at the save — under any of these conditions, each of which is logged when it occurs:

  • Legacy checkpoints. Resuming from a checkpoint written before frontier alignment disables it for the run, and because that run’s checkpoints then carry no frontier metadata, for every run descended from them. A run only regains frontier alignment by starting fresh.

  • Unstampable batches. Prompts whose extra_env_info rows are not dicts cannot carry stream ordinals; the first such batch permanently disables frontier alignment for the run.

  • Snapshot-ring eviction. If no retained dataloader snapshot sits at or below the trained frontier, that checkpoint alone falls back to the live cursor.

Independently of checkpointing, prompt groups whose generation fails a tolerated number of times (max_generation_failures > 0) are dropped during normal operation and are not recovered by a resume.

Target interleaving and the conservative cut. A target refilled from later prompts — after a tolerated failure, or routinely when a resume gap-fills an incomplete restored target — interleaves targets’ ordinals: training the refilled target can advance the frontier past another target’s still-generating groups. Checkpoints written in that window conservatively cut below the trained frontier, at the lowest ordinal still in flight (logged when it happens), and persist the ordinals already trained at or above the cut. A resume from such a checkpoint regenerates the in-flight prompts and drops the persisted trained ones — no skip, no re-training.

Gap-filling after restore#

After a restore, last_target_weight_already_generated is reset to current_training_step - 1 so the collector re-evaluates every target from the resume step onward. For each target it queries get_trajectories_needed and spawns only the workers required to complete the batch — previously buffered trajectories are reused and the collector does not regenerate them.

Disabling replay-buffer restore#

Set checkpointing.load_replay_buffer: false to skip the restore: the buffer starts empty and, on a frontier-aligned checkpoint, the whole buffered window is regenerated fresh from the rewound dataloader. This trades resume compute for an unbiased step composition — a prompt group completes only when its longest rollout finishes, so the groups sitting in the buffer at a save boundary systematically skew toward short rollouts; regenerating the window removes that bias. Restoration remains enabled by default and for legacy configs that omit the field. The conservative cut (target interleaving, above) also accounts for buffered-but-untrained groups — the checkpoint cut never sits above an ordinal whose only record is the buffer — so discarding the buffer is safe even on a cut-lowered checkpoint: the resume regenerates those groups from the rewound stream.

If no replay_buffer.pt file is found in the latest checkpoint directory, training likewise starts with an empty buffer and waits for the collector to fill it before the first training step.

Usage Tips#

  1. Buffer Sizing: The replay buffer size is automatically calculated as:

    buffer_size = num_prompts_per_step × max_trajectory_age_steps × 2

  2. Age Limits: Start with max_trajectory_age_steps: 1 and increase if needed for higher throughput

  3. Resource Allocation: Ensure sufficient GPU memory for both the training and generation clusters

  4. In-Flight Weight Updates: Enable in_flight_weight_updates: true to refit without waiting for the longest in-flight generation to finish. Except for managed Dynamo, the collector requests a generation pause and resume from every async backend around the weight transfer. Async vLLM implements this contract while preserving request state. A backend that does not implement the hook emits a warning once per backend type per process and refits without a collector-side pause or drain; SGLang is in this group today and instead relies on the pause its own weight synchronizer performs around the transfer. Managed Dynamo always drains active trajectories before refit. vLLM requires async_engine: true; the Megatron backend is always async-engine.

  5. Recompute KV Cache After Weight Updates: Set recompute_kv_cache_after_weight_updates: true to invalidate reusable KV/prefix caches when weights change. On the native async vLLM in-flight path, caches are cleared while generation is paused, so preserved requests recompute their KV after resuming. Other refit paths keep their existing post-update invalidation behavior. When false, in-flight requests retain their pre-update KV cache. On the Megatron generation backend, this must agree with policy.generation.mcore_generation_config.kv_cache_management_mode; setup errors on a mismatch.

Why Importance Sampling Correction Is Required for Async#

The GRPO Objective#

The standard GRPO loss function (without KL penalty) is:

\[ L(\theta) = E_{x \sim \pi_{\theta_{\text{old}}}} \Big[ \min \Big(\frac{\pi_\theta(x)}{\pi_{\theta_{\text{old}}}(x)}A_t, \text{clip} \big( \frac{\pi_\theta(x)}{\pi_{\theta_{\text{old}}}(x)}, 1 - \varepsilon, 1 + \varepsilon \big) A_t \Big) \Big] \]

where:

  • \(\pi_\theta\) is the policy model we are currently optimizing

  • \(\pi_{\theta_{\text{old}}}\) is the previous policy model (from the beginning of this step)

  • \(A_t\) is the advantage estimate

  • \(\varepsilon\) is a clipping hyperparameter

In standard GRPO, we assume trajectories are sampled from \(\pi_{\theta_{\text{old}}}\). However, in async GRPO, trajectories are actually sampled from \(\pi_{\theta_{\text{generator}}}\), which is the policy weights from N training steps ago (where N ≥ 1 depending on max_trajectory_age_steps).

Without importance sampling correction, the GRPO objective becomes fundamentally incorrect:

  1. Incorrect probability ratios: The ratio \(\frac{\pi_\theta(x)}{\pi_{\theta_{\text{old}}}(x)}\) uses \(\pi_{\theta_{\text{old}}}\) probabilities that were never actually used to generate the trajectories.

  2. Biased gradient estimates: Since we’re computing gradients based on samples from the wrong distribution, the policy updates become biased and can lead to instability.

When we enable importance sampling correction (use_importance_sampling_correction: true), we introduce the corrective term:

\[ \frac{\pi_{\text{training}}(x)}{\pi_{\text{generator}}(x)} \]

This transforms our loss function to properly account for the distribution mismatch. The corrected objective becomes:

\[ L(\theta) = E_{x \sim \pi_{\theta_{\text{generator}}}} \Big[ \frac{\pi_{\text{training}}(x)}{\pi_{\text{generator}}(x)} \min \Big(\frac{\pi_\theta(x)}{\pi_{\theta_{\text{old}}}(x)}A_t, \text{clip} \big( \frac{\pi_\theta(x)}{\pi_{\theta_{\text{old}}}(x)}, 1 - \varepsilon, 1 + \varepsilon \big) A_t \Big) \Big] \]

The importance sampling ratio \(\frac{\pi_{\text{training}}(x)}{\pi_{\text{generator}}(x)}\) is effectively \(\frac{\pi_{\theta_{\text{old}}}(x)}{\pi_{\theta_{\text{generator}}}(x)}\), which corrects for the N-step gap between the generator policy and the policy we assume we’re sampling from.

This correction ensures that we have unbiased gradient estimates and stable convergence.