NVIDIA-NeMo / NVIDIA-NeMo/RL

[Single Controller / Async RL] cleanup tracking issue

Open
#2,625 5 comments 0 reactions 2 assignees Claimed by @mehraakash View on GitHub
enhancement Feature
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.