refactor: unify segment_size + topology setup across grpo / ppo / distillation / single-controller
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
Follow-up from [#3262 review thread](https://github.com/NVIDIA-NeMo/RL/pull/3262#discussion_r3689964419).
Non-colocated PPO in #3262 duplicates ~55 lines of GRPO's `segment_size` + topology setup at `nemo_rl/algorithms/ppo.py:564` (mirrors `nemo_rl/algorithms/grpo.py:770-862`). `nemo_rl/algorithms/distillation.py:423+` and `nemo_rl/algorithms/opd.py:305+` have their own near-copies. Single-controller entrypoints are heading the same way.
The immediate drift (dropped `no-topo_rank` warning + dropped per-backend `gpus_per_instance` branch for trtllm/sglang) was patched into PPO on #3262 to unblock the PR; this issue tracks the shared-helper cleanup so future copies don't diverge again.
## Suggested direction
Extract a helper (e.g. in `nemo_rl/distributed/virtual_cluster.py`, alongside the existing `prepare_segment_topology`) covering the multi-cluster case: topology fetch, alive-node check, train/inference partition, per-backend `gpus_per_instance`, and the `no-topo_rank` warning. Call it from `grpo.py`, `ppo.py`, `distillation.py`, `opd.py`, and single-controller entrypoints.
Contributor guide
Assessment
This issue has not been assessed yet.