[Single Controller / Async RL] cleanup tracking issue
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
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 #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 #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()`.
- [x] ~~**Add a retry policy to the Single Controller train loop.** `_train_pump`'s per-cycle error path ([single_controller.py L437-L573](https://github.com/NVIDIA-NeMo/RL/blob/88e840731149a80138408c09e6c858851b451771/nemo_rl/algorithms/single_controller.py#L437)) 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](https://github.com/NVIDIA-NeMo/RL/blob/88e840731149a80138408c09e6c858851b451771/nemo_rl/algorithms/single_controller.py#L441). (Flagged during #2700 review.)~~ — Not planned.
- [x] ~~Support custom samplers by extracting the TQReplayBuffer implementation.~~ — ✅ **Done in #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.
- [x] ~~Honor `recompute_kv_cache_after_weight_updates` in async mode.~~ — ✅ **Done in #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`.
- [x] ~~Decouple the SC logprob refresh from the advantage config.~~ — ✅ **Done in #3266**: SC's `_policy_logprobs_required` / `_reference_logprobs_required` predicates ([`single_controller.py:104-110`](https://github.com/NVIDIA-NeMo/RL/blob/e5c6f93a4c93720bb4ed663564fb8fbe16da1f34/nemo_rl/algorithms/single_controller.py#L104-L110)) are loss-driven (`force_on_policy_ratio`, `seq_logprob_error_threshold`, `skip_reference_policy_logprobs_calculation`), replacing the old advantage-config gating.
- [x] ~~Router replay (R3) not supported/validated on the Single Controller (async + TransferQueue) path.~~ — ✅ **Done in #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 #3327.
- [x] **Support multi-reward (GDPO) in `AsyncRolloutImpl`.** Multi-reward aggregation — needed for the GDPO recipe — was dropped when the async rollout impl was introduced in #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 #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 #2566, context #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 #2700, thread https://github.com/NVIDIA-NeMo/RL/pull/2700#discussion_r3548128192, Known Missing Feature in #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 #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 #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 #3266)
- [ ] **Support (or clearly reject) non-colocated Megatron generation in the Single Controller.** [`setup.py` asserts `backend != "megatron"`](https://github.com/NVIDIA-NeMo/RL/blob/850b58377eb6ed54716419a90b75d749907ad063/nemo_rl/algorithms/single_controller_utils/setup.py#L101) on the non-colocated cluster split, whereas legacy `grpo.py` supports it via a dedicated inference policy ([`init_megatron_generation` non-colocated branch](https://github.com/NVIDIA-NeMo/RL/blob/850b58377eb6ed54716419a90b75d749907ad063/nemo_rl/algorithms/grpo.py#L1037-L1055) 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 #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`](https://github.com/NVIDIA-NeMo/RL/blob/850b58377eb6ed54716419a90b75d749907ad063/nemo_rl/models/generation/interfaces.py#L329) (FP8 KV cache), computes FP8 QKV scales via `policy.calibrate_qkv_fp8_scales(...)` and forwards them through [`sync_weights(..., kv_scales=...)`](https://github.com/NVIDIA-NeMo/RL/blob/d4a9d0f7c1686d936556441174c6519933aaddea/nemo_rl/algorithms/grpo.py#L2636) ([feature](https://github.com/NVIDIA-NeMo/RL/blob/d4a9d0f7c1686d936556441174c6519933aaddea/nemo_rl/algorithms/grpo.py#L2475)). SC's [`_sync_weights`](https://github.com/NVIDIA-NeMo/RL/blob/850b58377eb6ed54716419a90b75d749907ad063/nemo_rl/algorithms/single_controller.py#L624) 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. #3226). The synchronizer interfaces already accept `kv_scales`; only the driver-side compute+forward is missing. (Flagged during #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](https://github.com/NVIDIA-NeMo/RL/blob/dd3cb538492ec48a0ab96942e1860a6f96dc6119/nemo_rl/algorithms/single_controller_utils/setup.py#L380-L384)) can't fire through the real entrypoint: sglang+gym trips `_should_use_nemo_gym`'s `should_expose_http_server` [assert](https://github.com/NVIDIA-NeMo/RL/blob/dd3cb538492ec48a0ab96942e1860a6f96dc6119/nemo_rl/algorithms/grpo.py#L1983) and megatron+gym trips `_build_generation`'s [`ValueError`](https://github.com/NVIDIA-NeMo/RL/blob/dd3cb538492ec48a0ab96942e1860a6f96dc6119/nemo_rl/algorithms/single_controller_utils/setup.py#L187) 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](https://github.com/NVIDIA-NeMo/RL/blob/dd3cb538492ec48a0ab96942e1860a6f96dc6119/nemo_rl/algorithms/single_controller_utils/setup.py#L327)) so megatron+gym fails fast with a clear message before any build. (Flagged during #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 #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 #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 #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 #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 #3402.)
Add anything else here as it comes up.
Contributor guide
Assessment
This issue has not been assessed yet.