Train with Single-Controller (Async GRPO and PPO)#

Warning

The Single-Controller path is a beta feature and still under active development. The API and configuration surface are not yet stable and may change without notice. Issues and feedback are welcome — please file them at github.com/NVIDIA-NeMo/RL/issues.

The Single-Controller (SC) path is an alternative async GRPO and PPO runtime that runs rollout generation and policy training as two independent pumps coordinated by a single Ray actor (SingleControllerActor) sitting over a shared TransferQueue (TQ) data plane. Compared to the legacy async GRPO in async-grpo.md, SC decouples per-prompt rollouts from the per-step batch boundary: producers push finished rollouts into TQReplayBuffer at group granularity, and a pluggable StalenessSampler decides which groups the trainer consumes on each step.

Configure the Single-Controller Path#

The SC path is launched via a dedicated entrypoint:

uv run examples/run_grpo_single_controller.py --config <your-sc.yaml>

run_grpo_single_controller.py mirrors run_grpo.py for config loading — the same YAML files apply — but requires a few settings the legacy path does not. The default exemplar lives at examples/configs/grpo_math_1B_megatron_single_controller.yaml; the PPO one at examples/configs/ppo_math_1B_megatron_single_controller.yaml.

Mandatory settings#

Enable the TransferQueue data plane (required — the entrypoint refuses to start otherwise):
```
data_plane:
  enabled: true
```
Pick a generation backend. With vLLM, disable colocated inference (setup rejects colocated.enabled: true for every backend except Megatron) and enable the async engine (SC drives rollout via RolloutManager.generate_and_push, which is only supported on the disaggregated async engine):
```
policy:
  generation:
    backend: "vllm"
    vllm_cfg:
      async_engine: true
    colocated:
      enabled: false
      resources:
        num_nodes: 1
        gpus_per_node: 4  # inference GPUs; remainder go to training
```
Megatron generation is also supported, colocated or non-colocated. It requires the Megatron trainer (policy.megatron_cfg.enabled: true) and NeMo-Gym rollouts additionally require policy.generation.mcore_generation_config.expose_http_server: true. Colocated (colocated.enabled: true) additionally requires async_rl.min_groups_for_streaming_train == num_prompts_per_step: the engine stands down for the whole train step, so each step must be assembled as one full batch before training takes the GPUs. Per-request deadlines (generation_timeout_s, rollout_timeout_s) exclude the time the engine is stood down: in-flight requests freeze with their clocks suspended, and the clocks resume on the post-step wake. The non-colocated exemplar — a NeMo-Gym run with the OpenAI server exposed — lives at examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml; the colocated exemplar at examples/configs/grpo_math_1B_megatron_generation_colocated_single_controller.yaml:
```
policy:
  megatron_cfg:
    enabled: true
  generation:
    backend: "megatron"
    mcore_generation_config:
      expose_http_server: true  # required for NeMo-Gym rollouts
    colocated:
      enabled: false  # set true for colocated (no resources split needed)
      resources:
        num_nodes: 1
        gpus_per_node: 1  # inference GPUs; remainder go to training
```
One RL step = one training batch. The batch a step trains on is the whole step (see validate_single_controller_config in nemo_rl/algorithms/single_controller_utils/config.py). A GRPO step is also one optimizer step. A PPO step applies ppo.ppo_epochs actor updates and ppo.critic_ppo_epochs critic updates over that same batch. Both counts must be at least 1 and can be configured independently; the exemplar defaults the critic count to ${ppo.ppo_epochs}.
```
num_prompts_per_step * num_generations_per_prompt == policy.train_global_batch_size
num_prompts_per_step * num_generations_per_prompt == value.train_global_batch_size  # PPO
```
Enable importance sampling correction whenever the sampler admits off-policy data (any max_staleness_versions > 0 on the windowed/weight_fifo samplers, or max_lookahead_versions > 0 on in_order). The correction and its derivation are the same as for legacy async GRPO — see Why Importance Sampling Correction Is Required for Async:
```
loss_fn:
  use_importance_sampling_correction: true
```
Save the data plane for replay recovery. When Single-Controller checkpointing is enabled, all built-in samplers require checkpointing.save_data_plane: true so completed, unconsumed rollout groups survive a restart. Native TQ checkpointing currently supports only the simple storage backend. For multi-node runs, checkpoint_dir must be on a durable filesystem visible at the same path from every node.
```
checkpointing:
  enabled: true
  checkpoint_dir: /shared/checkpoints/my-run
  save_data_plane: true

data_plane:
  enabled: true
  backend: "simple"
```
(PPO) Set ppo: instead of grpo: — the two algorithm blocks are mutually exclusive, and SC reads every step setting from whichever one is present. A PPO run also needs value:, value_loss_fn: and ppo.adv_estimator.name: gae (same schemas as legacy PPO), a Megatron critic, and policy.offload_optimizer_for_logprob: true, which is what keeps the policy optimizer off the GPU while the critic runs. ppo.policy_training_start_step: N gives the usual critic warmup: for the first N steps the policy is neither trained nor refit, while the critic trains every step. ppo.warm_start_value_checkpoint seeds that critic from another run’s checkpoint instead, so a fresh run can skip the online warmup entirely — see Warm-Starting the Critic.

Checkpointing and Replay Recovery#

With checkpointing.save_data_plane: true, each Single-Controller checkpoint contains:

The normal model, dataloader, and controller state, plus optimizer state when configured.
A native TQ snapshot containing rollout tensor payloads and TQ state.
A metadata-only replay index describing the completed rollout groups stored in TQ.
A rollout_recovery.pt ownership ledger describing unfinished prompt groups that must be redispatched after a restart.
A replacement_reserve.pt sidecar containing prompts held for dropped-rollout replacement, when applicable.
The sampler dispatch position needed to continue scheduling from the correct point.

The TQ snapshot and replay index are captured under the same checkpoint barrier. Generation may continue while the snapshot is written, but completed-group commits and destructive TQ clears wait at the barrier. This ensures that the TQ snapshot and replay index describe the same set of groups.

On resume, Single-Controller validates the TQ snapshot against the trainer checkpoint, restores the replay index, and makes completed, committed, unconsumed groups available to the sampler before training resumes.

Replay recovery is supported by all built-in samplers: in_order, weight_fifo, ready_first, and windowed. Custom samplers must explicitly declare supports_buffer_checkpoint = True. Otherwise, setup emits a warning and completed buffered groups are not restored.

Periodic rollout snapshots#

Normal trainer checkpoints are written at step boundaries. Periodic rollout snapshots preserve newer rollout progress between those trainer checkpoints, including while the train pump is accumulating a streamed step:

checkpointing:
  enabled: true
  checkpoint_dir: /shared/checkpoints/my-run
  save_data_plane: true
  save_period: 1

rollout_checkpointing:
  snapshot_attempt_interval_s: 120
  keep_latest_k: 2
  restore_mode: latest
  extra_fingerprint_excluded_paths: []

token_capture:
  enabled: true

snapshot_attempt_interval_s is the cadence at which Single-Controller attempts a rollout snapshot. It is not a guarantee that a snapshot is written at every interval. An attempt after step N succeeds only when the immutable trainer checkpoint step_N is already durable. Consequently, save_period: 1 is recommended for continuous post-step coverage; with a larger value, attempts are skipped until the matching trainer checkpoint exists. Before the first training step, snapshots are anchored to the initial model and a fingerprint of the rollout-semantic configuration.

The bootstrap fingerprint is fail-closed: every configuration value affects compatibility unless NeMo-RL’s built-in denylist identifies it as operational, such as logging, cluster placement, checkpoint location, runtime ports, or credentials. This means configuration added by an external algorithm is safe by default—a change prevents bootstrap recovery instead of silently mixing incompatible rollout state. The bootstrap manifest also stores this credential-redacted compatibility identity so a rejected restart can report the exact changed dotpaths instead of showing only two opaque digests.

An integration may use extra_fingerprint_excluded_paths for additional runtime-only values that are not part of NeMo-RL’s built-in configuration:

rollout_checkpointing:
  extra_fingerprint_excluded_paths:
    - custom_algo.observability
    - env.private_agent.runtime_endpoint
    - env.private_agent.workers.*.log_dir

Each dotpath removes that value and its children from the compatibility identity. * matches one mapping or list level and ** matches any number of levels. Only exclude values that cannot affect prompts, generation, rewards, lineage, or the interpretation of persisted rollout data. These exclusions must be set on the original run as well as its restart.

Periodic snapshots currently require all of the following:

checkpointing.enabled: true and checkpointing.save_data_plane: true.
data_plane.backend: simple, because native TQ save/load is required.
token_capture.enabled: true.
A replay-recoverable sampler with training-claim ownership. All built-in samplers qualify. A custom sampler must explicitly declare both supports_buffer_checkpoint = True and supports_training_claims = True.

Each trainer or bootstrap anchor has a rollout_snapshots/ directory. A published snapshot_NNNNNN/ contains the native TQ snapshot and matching replay, dataloader, controller, replacement-reserve, and unfinished-rollout metadata. keep_latest_k retains recent committed snapshots as fallbacks; temporary or interrupted directories are never selected for recovery.

With restore_mode: latest, startup selects the newest compatible committed snapshot under the latest trainer anchor. With trainer_checkpoint, it ignores newer periodic rollout progress and resumes from the trainer checkpoint bundle. Checkpoint selection is read-only: neither mode removes snapshots. If no trainer checkpoint exists, trainer_checkpoint cannot safely reuse an existing bootstrap namespace, so startup fails without modifying it. Recover that state with latest or choose a new checkpoint_dir to start a fresh bootstrap lineage. Obsolete bootstrap snapshots are removed only by retention after a durable trainer checkpoint exists.

Bootstrap-only trainer_checkpoint behavior: A bootstrap rollout snapshot has no corresponding model or optimizer checkpoint. Therefore restore_mode: trainer_checkpoint deliberately fails when bootstrap state exists but no trainer checkpoint does. It does not ignore or delete that state. Use latest to recover it, or select a new checkpoint_dir to start from scratch.

Note

Completed groups are restored directly from the TQ snapshot. For unfinished token-capture groups, rollout_recovery.default_granularity controls both live failure and restart behavior:

sibling preserves each sealed sibling and redispatches only unfinished ones.
prompt_group retries every sibling in the group when any sibling is unfinished.

sibling is the default and avoids regenerating completed work. Use prompt_group when every generation in a recovered group must come from the policy weights live at redispatch.

task_source_granularity_overrides can select the policy using the Gym task_source embedded in the raw rollout row. Unlike agent_ref, this identity is available before Gym resolves the concrete agent and SC reserves the recovery group. When a row already carries an agent_ref, a matching agent_granularity_overrides entry wins over a matching task-source entry, mirroring Gym’s concrete-route precedence. Otherwise the task-source override, then the global default, applies. The agent map also keeps datasets collated before Gym recorded task_source working, although re-collating them is recommended. Non-default policies require token_capture.enabled: true. The task source and resolved policy are persisted in rollout_recovery.pt, so recovery does not reinterpret an existing group using changed configuration. A generation that already finished keeps its tokens in the token-capture staging area, so sibling reuses them unchanged; a redispatched sibling produces a new sample from the same prompt. Neither becomes a training row until every generation in the group has finished.

When a sampler does not support replay recovery, a requested data-plane checkpoint is written in shadow mode. The TQ snapshot is retained, but no authoritative replay index is written and its rows are not restored into the training replay buffer.

Native TQ save/load currently requires data_plane.backend: "simple". Mooncake-backed storage is not recoverable through this mechanism. A failure while saving or validating the TQ snapshot prevents the incomplete checkpoint bundle from becoming the latest resumable checkpoint.

Async-RL Knobs and Sampler Modes#

All SC async-RL runtime knobs live under async_rl: in the master config. The most important choice is the sampler, which sets the staleness policy shared by the rollout pump (how far it may run ahead) and the train pump (which groups it may consume).

Sampler modes#

Pick one of five modes with sampler.name. Each mode takes its own knobs, listed below — a knob from one mode has no effect under another:

Sampler modes: same buffer, four different training batches

Same buffer under trainer weight 2 (num_prompts_per_step=2, staleness window [0, 2]). windowed and weight_fifo select on start_weight (stamped at dispatch); in_order selects on target_step (stamped at admit). The three usually pick the same groups and diverge only when rollouts finish out of order, as drawn here.

`sampler.name`	Rollout gating	Train selection	Typical use
`in_order`	Dispatch may lead the trainer by up to `max_lookahead_versions` batches. Each dispatch is stamped with a `target_step`.	Consume the group whose `target_step == current_train_weight`.	Sync mode (`max_lookahead_versions=0`) and legacy-async exact-batch semantics (`max_lookahead_versions>=1`). The only mode supported on a PPO run.
`weight_fifo`	Same gate as `in_order` (`max_staleness_versions` of lookahead).	Drain the oldest in-window `start_weight` first, waiting for that weight’s batch to fill.	Strict weight-version FIFO under a bounded lookahead.
`ready_first`	Same gate as `weight_fifo` (`max_staleness_versions` of lookahead).	Take any ready group generated by a policy version no newer than the trainer, including late stragglers.	Completion-order streaming without stale-group eviction.
`windowed`	Ungated — rollout keeps producing until the buffer fills.	Take any ready group with `start_weight` in `[train - max_staleness_versions, train]`, optionally freshest-first.	Over-sampled streaming; aged groups outside the window are evicted (wasted compute).
`custom`	Determined by the imported class.	Determined by the imported class.	`target: "module:ClassName"` — bring your own `PromptGroupSampler`.

Config → behavior map#

The shipped exemplars cover three of the five modes:

Mode	`sampler.name`	Sampler knob	`min_groups_for_streaming_train`	`max_buffered_rollouts`	Exemplar
Sync / on-policy	`in_order`	`max_lookahead_versions: 0`	`${grpo.num_prompts_per_step}`	`num_prompts_per_step × 1`	`grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml`
Async, exact batch→step matching	`in_order`	`max_lookahead_versions: >= 1`	`x <= num_prompts_per_step`	`num_prompts_per_step × (max_lookahead_versions + 1)`	`grpo_math_1B_megatron_single_controller.yaml`
Streaming, gated dispatch	`weight_fifo`	`max_staleness_versions: >= 1`	`x <= num_prompts_per_step`	`num_prompts_per_step × (max_staleness_versions + 1)`	— (none shipped)
Streaming, ready-first	`ready_first`	`max_staleness_versions: >= 1`	`x <= num_prompts_per_step`	`num_prompts_per_step × (max_staleness_versions + 1)`	— (none shipped)
Streaming, over-sampled	`windowed`	`max_staleness_versions: >= 1`	`x <= num_prompts_per_step`	Larger than the gated capacity (dispatch is ungated)	`grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml`

Field definitions:

max_buffered_rollouts — hard cap on unconsumed rollout groups buffered in the data plane. Validated at setup against the gated sampler’s required capacity; a value too small deadlocks the rollout pump, so setup raises instead of silently blocking. Sized from the widest window the run ever uses, so warmup_lookahead_versions rather than max_lookahead_versions when it is set.
min_groups_for_streaming_train — minimum ready groups the trainer waits for before dispatching a batch. Set to num_prompts_per_step for sync/legacy semantics; lower for streaming. (PPO) Must equal num_prompts_per_step — the critic has no split train API, so each critic epoch calls the full-step train_from_meta once per chunk. Splitting an RL step across chunks would multiply both models’ configured optimizer updates by the number of chunks.
sampler.warmup_lookahead_versions (PPO) — lookahead used while ppo.policy_training_start_step critic warmup is in progress, shrinking back to max_lookahead_versions afterwards. The SC equivalent of ppo.async_ppo.warmup_generation_lead_steps.

Implementation Structure#

The SC path splits the async-GRPO loop across a rollout pump and a train pump that share a TQReplayBuffer and are orchestrated by the SingleControllerActor.

SC ownership and call model: driver, SingleControllerActor, and remote worker groups

The driver builds every heavy object and cloudpickles it into SingleControllerActor (a CPU-only Ray actor). TQReplayBuffer, RolloutManager, and the sampler live inside that actor — reaching them is a direct call, not RPC. Only the generation worker group, TQPolicy, and the TransferQueue data plane are separate processes; the dashed arrows are the only hops that cross a process boundary.

Core components#

1. `SingleControllerActor` (`nemo_rl/algorithms/single_controller.py`)#

Single Ray actor that runs _rollout_pump and _train_pump concurrently as asyncio tasks.
Receives a fully-constructed SingleControllerActorArgs (cloudpickled by the driver) — the actor does no construction work of its own, because running setup inside a nested Ray actor breaks runtime_env resolution. Exception: Logger is built inside the actor, because wandb/TB backends hold a _thread.lock that cloudpickle can’t serialize.
On startup, rebinds self._rollout_manager._tq_buffer = self._buffer. rollout_manager and tq_buffer are separate fields on the args dataclass, so Ray deserializes them as two independent buffers; without the rebind, the writer and the sampler would see different copies and the sampler would never observe committed groups.
Pump crashes propagate to the driver; in-flight rollouts drain on exit.

2. `TQReplayBuffer` (`nemo_rl/algorithms/async_utils/replay_buffer.py`)#

Group-granular replay buffer with reserve/commit slot accounting.
start_weight is stamped on the slot at reserve (dispatch); commit tensorizes the group into N training-shaped rows, records end_weight, and flips the slot ready. Because slots are appended at reserve, buffer index order equals dispatch order and start_weight only ever increases down the buffer — which is why windowed and weight_fifo usually pick the same groups and diverge only when rollouts finish out of order.
Tracks target_step per group when the sampler assigns one at admit time (used by in_order).

3. `RolloutManager.generate_and_push` (`nemo_rl/experience/rollout_manager.py`)#

One entry point per prompt group: reserve a buffer slot, drive the rollout via AsyncRolloutImpl or AsyncNemoGymRolloutImpl, and commit with the observed weight versions.
env_handles provide per-task environments to the rollout implementations.

4. Samplers (`nemo_rl/algorithms/async_utils/staleness_sampler.py`)#

Filter-only prompt-group selector over TQReplayBuffer. The base PromptGroupSampler protocol defines admit, select, and evict.
WindowedSampler, ReadyFirstSampler, WeightFifoSampler, and InOrderSampler are the built-in policies (one per row in the Sampler modes table). The custom mode (CustomSamplerConfig.target) makes create_sampler import a user-supplied class by FQN and type-check it against PromptGroupSampler.

5. `_rollout_pump` and `_train_pump`#

_rollout_pump: pulls prompts from the dataloader, calls sampler.admit, dispatches RolloutManager.generate_and_push, and honours max_inflight_prompts as a backpressure cap.
_train_pump: sampler.evict → sampler.select → _value_stage (PPO only) → _advantage_stage → _value_train_epochs (PPO only) → TQPolicy split API (begin_train_step / train_microbatches_from_meta / finish_train_step) → dp_client.clear_samples.

Coordination Flow#

Driver setup: setup_single_controller builds the worker groups, virtual cluster, dp client, dataloader, TQReplayBuffer, RolloutManager, and weight synchronizer, and packs them into a SingleControllerActorArgs that the entrypoint cloudpickles into the actor.
Actor startup: SingleControllerActor launches _rollout_pump and _train_pump concurrently as asyncio tasks; both share the same TQReplayBuffer and StalenessSampler.
Rollout pump loop: sampler.admit gates dispatch against the current trainer version (returning a target_step for in_order); the pump then reserves a buffer slot, drives RolloutManager.generate_and_push, and commits with the observed start_weight / end_weight.
Train pump loop: sampler.evict drops out-of-window groups and sampler.select picks the next batch. On PPO, _value_stage runs the critic forward, _advantage_stage computes advantages, and _value_train_epochs runs ppo.critic_ppo_epochs critic updates. The TQPolicy split API then runs one optimizer step per RL step on GRPO, or ppo.ppo_epochs policy updates on PPO.
Weight sync: after each optimizer step the pump bumps the trainer version, clears rollout permission, calls the weight synchronizer, and re-opens the rollout pump for the next version.

Relation to Legacy Async GRPO#

The legacy async GRPO (grpo.async_grpo.enabled: true under run_grpo.py) and the SC path both target the same async training problem but partition responsibilities differently:

	Legacy async GRPO	Single-Controller
Entrypoint	`run_grpo.py`	`run_grpo_single_controller.py`
Data-plane	Direct actor RPC	TransferQueue (`data_plane.enabled: true` required)
Rollout batching	Full-batch `AsyncTrajectoryCollector`	Per-prompt `RolloutManager.generate_and_push` into a group-granular `TQReplayBuffer`
Staleness policy	Single knob (`max_trajectory_age_steps`)	Pluggable `StalenessSampler` (`in_order` / `weight_fifo` / `ready_first` / `windowed` / `custom`)
Batch boundary	Sampled by target weight	Sampler-defined; can decouple rollout dispatch from train batch (streaming)

Migrating a legacy async config#

SC reads its async knobs from async_rl: and requires grpo.async_grpo: null (or ppo.async_ppo: null on a PPO run) — run_grpo_single_controller.py raises if a legacy block is still present, so null it out when porting rather than leaving it in place.

Do not carry max_num_epochs: -1 across either. ppo.md requires that value for legacy async PPO, but SC has no -1 convention: the rollout pump gates on _current_epoch < max_num_epochs, so any non-positive value trains zero steps and exits successfully. Setup rejects it — set a positive max_num_epochs and bound the run with max_num_steps.

Legacy `grpo.async_grpo.` / `ppo.async_ppo.`	SC equivalent `async_rl.*`
`enabled: true`	Implicit — SC is always async; use `sampler.max_lookahead_versions: 0` for sync semantics, `>= 1` for async
`max_trajectory_age_steps: N`	`sampler.name: in_order` with `sampler.max_lookahead_versions: N`
`warmup_generation_lead_steps` (PPO)	`sampler.warmup_lookahead_versions` — the lookahead to use while critic warmup is in progress
`recompute_kv_cache_after_weight_updates`	`recompute_kv_cache_after_weight_updates` (same)
`in_flight_weight_updates`	Always effectively true; `false`-equivalent behavior is not yet supported (drain-gate tracked in issue #2625)
(no legacy equivalent — matches legacy full-batch train semantics)	`min_groups_for_streaming_train: ${grpo.num_prompts_per_step}`, or `${ppo.num_prompts_per_step}` on a PPO run
(no legacy equivalent — matches legacy `max_trajectory_age + 1` batches in flight)	`max_inflight_prompts: num_prompts_per_step × (max_lookahead_versions + 1)`
(no legacy equivalent — legacy sizes its buffer to `num_prompts_per_step × max_trajectory_age_steps × 2`)	`max_buffered_rollouts: num_prompts_per_step × (max_lookahead_versions + 1)` (tight; see the Config → behavior map for per-sampler values)

Known Missing Features#

The SC path is still under active development. Feature gaps are tracked in issue #2625. Notable items:

Multimodal/VLM GRPO is supported with Megatron generation. Set policy.is_vlm: true; see the CLEVR Single-Controller recipe.
Multi-Teacher On-Policy Distillation (MOPD) is supported for text-only NeMo Gym rollouts; multimodal/VLM MOPD is not yet supported. See Multi-Teacher On-Policy Distillation.
Train backend: only Megatron is supported and validated; the AutoModel training path has not been tested on SC.
Generation backend: vLLM and Megatron generation are supported (Megatron in both non-colocated and colocated modes); SGLang and TRT-LLM have not been tested on SC.
Validation is not yet supported (setup raises on val_period > 0, val_at_start, or val_at_end); checkpointing is.
(PPO) Rollout drop budgets — async_rl.rollout_failure.max_skipped_prompts and max_consecutive_dropped_prompts must both be 0. A drop shortens the step, and the critic shards it against the configured value.train_global_batch_size rather than its actual size, so setup rejects a non-zero budget. The resiliency layer stays available on GRPO.
Reward shaping and sample filtering — reward_shaping, reward_scaling, and use_dynamic_sampling are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. Environment-flagged sample masking and overlong_filtering are supported; truncated completions are excluded from the loss through sample_mask, and a step in which every completion is filtered is rejected rather than skipped.
The windowed sampler has no over_sampling_ratio cap — over-produced groups aged past the window are evicted, wasting rollout compute.
The drain gate in refit is not yet supported.