NVIDIA-NeMo / NVIDIA-NeMo/RL

refactor: unify segment_size + topology setup across grpo / ppo / distillation / single-controller

Open
#3,464 1 comment 0 reactions 1 assignee Claimed by @bg51717 View on GitHub
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

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.