[Single Controller / Async RL] cleanup tracking issue

[yuki-97]

Tracking issue for loose ends in the Single Controller / Async RL / Streaming.
Items will be checked off or split into separate issues as they land.

Known cleanup items

Drop after new implementation aligned with old implementation.

• Async RL
  • Drop nemo_rl/algorithms/async_utils/trajectory_collector.py .
  • Drop old implementation in nemo_rl/algorithms/async_utils/replay_buffer.py.
• Streaming
  • Drop nemo_rl/experience/rollouts.py, new implementation is in refactor: unify per-prompt rollout to RolloutManager #2567 and we'll no longer maintain sync rollout path.
  • Cleanup sync engine in inference backend, since we'll drop sync rollout path, so sync engine will no longer be used.

Known missing items

✅ Resolved

• Epoch semantics in the Single Controller train loop: _train_pump was bounded only by max_train_steps, with no max_num_epochs/current_epoch notion of full dataset passes. — ✅ Fixed in feat(algorithms): SingleController streaming train_pump (split-API consumer) #2700 (feat(sc): epoch-bounded rollout passes (max_num_epochs parity), commit 4c2bf5b6c): SC now has max_num_epochs + a _current_epoch counter with epoch-bounded rollout dispatch, surfaced in run()/ping().
• Add a retry policy to the Single Controller train loop. _train_pump's per-cycle error path (single_controller.py L437-L573) aborts the worker step and re-raises on any mid-cycle failure (microbatch OOM, NaN-guard trip, prepare_logprobs / begin / finish failure), so a transient error kills the run — there is no retry. Flagged by the TODO(sc): retry policy is a follow-up at L441. (Flagged during feat(algorithms): SingleController streaming train_pump (split-API consumer) #2700 review.) — Not planned.
• Support custom samplers by extracting the TQReplayBuffer implementation. — ✅ Done in feat(sc): train path — StalenessSampler + _train_pump rewrite #3220: PromptGroupSampler protocol + WindowedSampler / WeightFifoSampler / InOrderSampler built-ins in nemo_rl/algorithms/async_utils/staleness_sampler.py, plus CustomSamplerConfig(target="module:ClassName") wired through create_sampler and the AsyncRLConfig.sampler discriminated union.
• Honor recompute_kv_cache_after_weight_updates in async mode. — ✅ Done in feat(sc): setup + entrypoint #3266: SC's _sync_weights (nemo_rl/algorithms/single_controller.py:615) calls self._gen.invalidate_kv_cache() after every weight sync when AsyncRLConfig.recompute_kv_cache_after_weight_updates=True.
• Decouple the SC logprob refresh from the advantage config. — ✅ Done in feat(sc): setup + entrypoint #3266: SC's _policy_logprobs_required / _reference_logprobs_required predicates (single_controller.py:104-110) are loss-driven (force_on_policy_ratio, seq_logprob_error_threshold, skip_reference_policy_logprobs_calculation), replacing the old advantage-config gating.
• Router replay (R3) not supported/validated on the Single Controller (async + TransferQueue) path. — ✅ Done in feat(sc): support router replay with TQ #3378: routed-expert indices are threaded through the async rollout → TQReplayBuffer → training payload, with fail-loud validation when router_replay.enabled=true and routes are missing; adds SC + R3 recipe grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.{yaml,sh} + a nightly correctness test. Closes [Single Controller] Router replay (R3) not supported/validated on the async + TransferQueue path #3327.

• Support multi-reward (GDPO) in AsyncRolloutImpl. Multi-reward aggregation — needed for the GDPO recipe — was dropped when the async rollout impl was introduced in feat: add AsyncRolloutManager for native async per-prompt rollouts #2566; wire it back through the async path for GDPO parity with grpo.

• Support the greedy option in RolloutManager. Dropped during the streaming rewrite in refactor: unify per-prompt rollout to RolloutManager #2567; needed for greedy validation/eval sampling before the sync path can be fully retired.

• Fix env observation chat templating at the env interface level. Templating currently leaks across the env boundary and each backend rolls its own; move it into the env interface so observations arrive already templated. Flagged in feat: add AsyncRolloutManager for native async per-prompt rollouts #2566, context feat: Add Tau bench environment #2479.

• Generation-level backpressure in the Single Controller — cap in-flight generations, not prompt groups. Today _rollout_pump backpressures per prompt group (max_inflight_prompts + the _buffer_capacity semaphore), so a group with one slow generation holds its buffer slot for the whole group → long-tail under-utilization (can't admit the next group even with spare generation capacity). Doing it properly needs the rollout/generation-worker rewrite that streams each generation into the DataPlane as it completes — today generate_and_push pushes a whole group atomically, so SC never sees individual generations finish and a generation-level semaphore has nothing to release against mid-group. Track with that rewrite; until then max_inflight_prompts stays as a conservative bound (under-utilizes on long tails but never over-subscribes DataPlane memory). (Discussed in feat(algorithms): SingleController streaming train_pump (split-API consumer) #2700, thread feat(algorithms): SingleController streaming train_pump (split-API consumer) #2700 (comment), Known Missing Feature in feat(sc): rollout path — TQReplayBuffer, rollout_manager, _rollout_pump #3219)

• Add an over_sampling_ratio (>1) knob to bound rollout over-sampling. The rollout pipeline can over-sample beyond what the trainer will consume; a configurable ratio lets recipes cap wasted generations while keeping enough headroom for slow tails. (Known Missing Feature in feat(sc): rollout path — TQReplayBuffer, rollout_manager, _rollout_pump #3219)

• Support multiple mini-steps inside a single RL step. The SC train loop today does exactly one optimizer.step per RL step; PPO-style multi-mini-step updates (iterating minibatches over the same batch of rollouts) need to be plumbed through the train pump. (Known Missing Feature in feat(sc): train path — StalenessSampler + _train_pump rewrite #3220)

• Add a drain gate to refit. Refit today does not gate in-flight rollouts to fully drain before swapping weights, so straggler generations can still be produced against the pre-refit weights. Needed for correctness on tight staleness bounds. (Known Missing Feature in feat(sc): setup + entrypoint #3266)

• Support (or clearly reject) non-colocated Megatron generation in the Single Controller. setup.py asserts backend != "megatron" on the non-colocated cluster split, whereas legacy grpo.py supports it via a dedicated inference policy (init_megatron_generation non-colocated branch builds MegatronGeneration(cluster=inference_cluster, ...)). All SC recipes use vLLM generation so this is latent; either mirror grpo's non-colocated Megatron-generation path, or keep the fail-loud guard with a message that names the generation backend rather than reading like it's about Megatron training. (Flagged during feat(sc): setup + entrypoint #3266 review)

• FP8 KV-cache scale sync not implemented in the Single Controller. grpo's train loop, when the vLLM backend reports requires_kv_scale_sync (FP8 KV cache), computes FP8 QKV scales via policy.calibrate_qkv_fp8_scales(...) and forwards them through sync_weights(..., kv_scales=...) (feature). SC's _sync_weights calls sync_weights() with no kv_scales and never calls calibrate_qkv_fp8_scales, so an FP8-KV-cache rollout under SC keeps stale quantization scales after weight updates (degraded generation, cf. fix: forward calibrated KV cache scales on colocated ZMQ IPC refit path #3226). The synchronizer interfaces already accept kv_scales; only the driver-side compute+forward is missing. (Flagged during feat(sc): setup + entrypoint #3266 review)

• SC NeMo-Gym non-vLLM guard is unreachable dead code, placed after the expensive build. The if generation_config["backend"] != "vllm": raise NotImplementedError in setup_single_controller (setup.py:380) can't fire through the real entrypoint: sglang+gym trips _should_use_nemo_gym's should_expose_http_server assert and megatron+gym trips _build_generation's ValueError first; the guard only sits after _build_generation/_build_trainer (multi-minute cluster+model build). test_nemo_gym_rejects_non_vllm_backend reaches it only by mocking both earlier gates. Hoist the check to right after use_nemo_gym = _should_use_nemo_gym(...) (setup.py:327) so megatron+gym fails fast with a clear message before any build. (Flagged during feat(sc): NeMo-Gym path #3267 review.)

• Tighten the Single Controller NeMo-Gym functional-test signal once SC validation is wired. tests/functional/grpo_async_gym_single_controller.sh currently checks only max(data["train/reward"]) > 0, which is a liveness bound; replace it with validation accuracy or a trending mean-reward threshold comparable to grpo_async_gym.sh after validation metrics are available. (Flagged during feat(sc): NeMo-Gym path #3267 review.)

• Overlap Single Controller NeMo-Gym spinup with deferred vLLM model loading. SC currently completes _build_generation/finish_generation before spinup_nemo_gym_actor, serializing setup; mirror grpo.py by reserving server URLs with defer_model_load=True and running Gym spinup alongside vLLM weight loading (and policy initialization for non-colocated execution). Tracked by the TODO in single_controller_utils/setup.py. (Flagged during feat(sc): NeMo-Gym path #3267 review.)

---

Above: Jul 28 version.

• val_start_at to delay periodic validation. Skip periodic validation until a configured training step, so early cheap steps don't pay for eval before the model is worth measuring. (Landed in grpo feat(grpo): add val_start_at to delay periodic validation #3400.)

• Validation-only sampling params. Let validation use its own temperature / top_p (typically near-greedy) so accuracy is a stable metric independent of training's exploration setting. (Landed in grpo feat(grpo): validation-only sampling params and grouped pass@k validation #3401.)

• Optional env-flagged sample masking. Gate NeMo-Gym env-flagged sample dropping (default on) behind grpo.mask_env_flagged_samples; today it's always on and makes effective batch composition non-deterministic, which hurts controlled experiments / A/B runs. (Landed in grpo feat(grpo): make env-flagged sample masking optional #3402.)

Add anything else here as it comes up.

[yuki-97]

Priority of the above Known missing items of Jul 28 version + something not list above.

P0:

1. [Foundation] Support validate (not list above) + Tighten the Single Controller NeMo-Gym functional-test signal once SC validation is wired. : feat(sc): add native and NeMo-Gym validation #3203 need merge to main. @haitian-nvidia
   1. num_vali_generations_per_prompt https://github.com/NVIDIA-NeMo/RL/pull/3401/changes#r3718464659 (FYI)
2. [Foundation] Support checkpointing. (not list above): feat(sc): add checkpoint save/restore to SingleController #3138 need merge to main. @haitian-nvidia
3. [Performance] Evict stale in-flight rollouts early. (not list above):
   1. Non-Gym Path feat(sc): abort stale in-flight rollouts early in non Gym path #3263 @RayenTian
   2. Gym Path @RayenTian
4. [Performance] Decouple the SC logprob refresh from the advantage config. ✅ Already done in feat(sc): setup + entrypoint #3266
5. Add an over_sampling_ratio (>1) knob to bound rollout over-sampling. @yuki-97
6. [Performance] Generation-level backpressure in the Single Controller.
7. [Performance] FP8 KV-cache scale sync not implemented in the Single Controller. @mehraakash
8. [Performance] Overlap Single Controller NeMo-Gym spinup with deferred vLLM model loading. @yuki-97
9. [Accuracy] Honor recompute_kv_cache_after_weight_updates in async mode. ✅ Already done in feat(sc): setup + entrypoint #3266
10. [Accuracy] Router replay (R3) not supported/validated on the Single Controller (async + TransferQueue) path. @zyzhou5 ✅ Done in feat(sc): support router replay with TQ #3378
11. [Logging] Log legacy equivalent of rollout_metrics in the SC path

P1:

1. [Streaming 2] Add a drain gate to refit. @RayenTian
2. [Streaming 3] Support streaming 3. (not list above) @mehraakash
3. [Fault Tolerance] Support fault tolerance. (not list above)
4. Support multi-reward (GDPO) in AsyncRolloutImpl.
5. Support / Validate other rollout backend.
   1. Replace "Support (or clearly reject) non-colocated Megatron generation in the Single Controller."
   2. Replace "SC NeMo-Gym non-vLLM guard is unreachable dead code, placed after the expensive build."
6. Support DTensor training backend. (not list above): feat(policy): split-API train-step state machine on DTensor v1/v2 #2692 @mehraakash
7. Support multiple mini-steps inside a single RL step.

P2:

1. Support the greedy option in RolloutManager.
2. Fix env observation chat templating at the env interface level.

[terrykong]

Missing item: the Single Controller train loop drops rollout_metrics entirely.

cc @yuki-97 — flagged during the #3378 (R3 on SC + TQ) review; this is broader than R3 so it belongs here rather than on that PR.

There are two stacked gaps, and the second is easy to miss because fixing the first alone doesn't close it.

Gap A — SC-only, a regression vs the legacy loop. Both legacy drivers log the aggregated rollout_metrics dict:

• sync grpo_train: logger.log_metrics(rollout_metrics, total_steps + 1, prefix="train")
• async async_grpo_train: aggregated via aggregate_rollout_metrics, then metrics.update(rollout_metrics)

SC logs only step_metrics + timing_metrics (single_controller.py:558-563), per the standing TODO at single_controller.py:552-556. So on SC we currently lose mean_gen_tokens_per_sample, the gen/input/total token-length histograms, per_worker_token_counts, and the env timing metrics.

The drop point is concrete: PromptGroupRecord does carry a rollout_metrics field and RolloutManager populates it (rollout_manager.py:97-110), but TQReplayBuffer.commit ships only record_to_train_batch(record, ...), and that function never reads record.rollout_metrics — it returns a tensor-only dict (payload.py:84-95). So the metrics die at the TQ boundary, before the train loop could log them.

Gap B — all paths: the r3/* router-replay fallback counters are never in rollout_metrics to begin with. _add_r3_fallback_metrics writes r3/routed_experts_fallback_token_route_fraction and friends into the local gen_metrics dict at exactly two sites — rollouts.py:370 (legacy sync) and rollouts.py:511 (legacy async). Both callers pop only _response_truncated and gen_leader_worker_idx and let the dict fall out of scope; neither the sync rollout_metrics nor the async sample_metrics ever merges gen_metrics. Nothing repo-wide consumes the r3/* keys. On SC it is worse still: RolloutManager never goes through generate_responses*, so _generate_response builds gen_metrics from scratch and _add_r3_fallback_metrics is never called at all.

This is why #3378's docs change correctly warns "no training loop currently forwards it to the metric logger, so do not rely on it in dashboards or gates" — and why R3 route-capture fallback is currently observable only as a R3 router replay fallback: warning in the run log.

Suggested item for the known-missing list:

• Propagate rollout_metrics through TQ so the Single Controller train loop can log them. PromptGroupRecord.rollout_metrics is populated but dropped at TQReplayBuffer.commit / record_to_train_batch, so SC loses mean_gen_tokens_per_sample, the length histograms, per_worker_token_counts and env timings that both legacy drivers log (grpo.py:2787, grpo.py:4659); tracked by the TODO at single_controller.py:552-556. Sub-item: once that lands, also call _add_r3_fallback_metrics from RolloutManager._generate_response and merge gen_metrics into rollout_metrics — today the r3/* fallback counters are dropped on every path, so router-replay fallback rate cannot be dashboarded or gated anywhere. (Flagged during feat(sc): support router replay with TQ #3378 review.)

@yuki-97 happy to move this into the body's known-missing list if you'd rather have it as a checkbox — didn't want to edit your issue without asking.

[terrykong]

@haitian-nvidia i added a sub-bullet to adding validate to SC as an FYI that this change is going in the legacy path for validation pass@k (and the associated num_val_generations_per_prompt flag)

https://github.com/NVIDIA-NeMo/RL/pull/3401/changes#r3718464659

[tianyi-zhang-02]

@RayenTian @yuki-97 I will take the open P1 item for multiple mini-steps within one RL step.

I am planning to keep the rollout batch and old logprobs fixed across K optimizer updates, retain the TQ/DataPlane payload until the final update, and make RL-step, optimizer-step, and weight-version accounting explicit. I will cover checkpoint resume plus ratio/clipping correctness as part of the same change.

For the first cut, I would start with repeated full-batch updates and leave shuffled minibatches as a follow-up, unless you want minibatching included from the start.

[yuki-97]

PPO on Single Controller (#3773) — gaps left open, tracked here.

1. Streaming critic training. The value workers have no split begin / train_microbatches / finish train API, so one RL step must be one chunk: config.py:808 requires async_rl.min_groups_for_streaming_train == num_prompts_per_step. The policy already has the split API; giving the critic one is the unlock.

2. Rollout drop budgets are rejected under PPO. A drop shortens the step, but _value_train never passes gbs, so tq_value.py:160 shards that short step against the configured value.train_global_batch_size and trips the divisibility assert. Guarded at config.py:823. The resiliency layer (#3665) stays fully available on GRPO — the policy path passes batch_size=None and tolerates a short step. Real fix is to thread the actual step size through, but gbs also drives local_gbs, num_global_batches and scheduler.step(increment=gbs), so it is not a one-liner.

3. Only the in_order sampler. config.py:850 rejects windowed / weight_fifo / custom under PPO — untested during critic warmup in particular.

4. Megatron critic only. Only megatron_value_worker.py mixes in TQWorkerMixin; the DTensor value worker has no data-plane setup to call, so it would die inside Ray with the model already on GPU. Guarded at config.py:793. Note the asymmetry is value-side only — all three policy workers carry the mixin.

5. Reward shaping / filtering knobs. overlong_filtering, reward_shaping, reward_scaling and use_dynamic_sampling are implemented on neither algorithm block on the SC path, so config.py:741 rejects them for GRPO and PPO alike rather than letting a run silently skip the shaping it was configured for.

6. total_flops under-reports under ppo_epochs > 1. Only the last epoch's train result is recorded, and total_flops is additive — see the TODO at single_controller.py:1050. Throughput from an SC PPO run is not comparable to the v1 path until this is fixed.

7. adv_estimator.name: raw_reward under PPO passes the estimator factory but yields no returns column, so the critic train would fetch a column nobody wrote (setup.py:551). ppo.py behaves the same way, so this is not an SC regression — but SC does not reject it either.

8. drop_incomplete_targets_on_restore — ✅ done in #3813. SC builds a TQReplayBuffer, not ReplayBufferImpl, so the legacy ppo.async_ppo.* field never applied — and it is rejected anyway, since the entrypoint requires ppo.async_ppo: null. #3813 adds the SC equivalent, async_rl.drop_incomplete_targets_on_restore (default false, in_order sampler only).

Unrelated to this list: the SC PPO exemplar inherits the exemplar-sizing OOM tracked in #3793, and renaming run_grpo_single_controller.py → run_single_controller.py is a separate PR.