feat: Topology aware placement

[youngeunkwon0405]

What does this PR do?

Adds topology-aware NVLink-domain placement for all NeMo-RL training algorithms (GRPO, SFT, DPO, RM, distillation). On multi-rack GB200 NVL72 (or any cluster where nodes are grouped into NVLink switch fabrics), this ensures that tensor-parallel and pipeline-parallel groups stay within the same NVLink domain, so collective operations use NVLink instead of InfiniBand.

How it works

1. Cluster setup — ray.sub

ray.sub runs a topology probe on each node at cluster start. The probe reads two pieces of information and registers them as Ray custom resources:

Resource key	Value	Source
nvlink_domain_<ClusterUUID>	1.0	nvidia-smi -q ClusterUUID — all nodes sharing a NVLink switch fabric get the same UUID
topo_rank	integer	Priority: SLURM_TOPOLOGY_ADDR (block.node) → SLURM_PROCID → hostname digits

No changes to ray.sub invocation are required — the probe runs unconditionally and falls back silently when topology info is unavailable (e.g. DGX/HGX, non-SLURM environments).

2. YAML configuration

Set cluster.segment_size to the number of training nodes per NVLink domain segment:

cluster:
  gpus_per_node: 4
  num_nodes: 36          # total nodes (train + inference)
  segment_size: 18       # nodes per NVLink domain (one NVL72 rack = 18 nodes w/ 4 GPUs each)

segment_size: null (the default in all exemplar configs) disables topology-aware placement and falls back to standard Ray scheduling.

The value to use depends on your hardware:

• GB200 NVL72: one rack = 18 compute nodes × 4 GPUs = 72 GPUs. Set segment_size: 18.
• General rule: segment_size = number of nodes that share one NVLink switch fabric.

3. How training node selection works

When segment_size is set, NeMo-RL:

1. Queries ray.nodes() for each node's nvlink_domain_* and topo_rank custom resources.
2. Groups nodes by NVLink domain and sorts domains by their minimum topo_rank.
3. Greedily selects complete segments (segment_size nodes each) from domains in topological order until num_nodes training nodes are claimed.
4. Pins those nodes to the selected domains via Ray placement group resource constraints, so Ray cannot schedule workers elsewhere.
5. Passes segment_size to RayVirtualCluster so that rank assignment within each placement group also follows the topological order (domain_min_topo_rank → topo_rank → gpu_id).

If no NVLink domain info is found (nodes have no nvlink_domain_* resource), the feature degrades gracefully to unordered placement with a warning.

4. Inference segment size (non-colocated GRPO)

For non-colocated GRPO, the inference cluster segment size is derived automatically from the generation config — you do not need to set it manually:

gpus_per_instance = TP × PP          (vLLM)  or  gpus_per_server  (SGLang)
nodes_per_instance = ceil(gpus_per_instance / inference_gpus_per_node)

Inference topology constraints are applied only when nodes_per_instance > 1 (cross-node model parallelism) and inference_nodes is divisible by nodes_per_instance. Inference nodes are allocated from the nodes not claimed by training, so the two pools never overlap.

Example (40 nodes, 5 NVLink domains × 8 nodes each, segment_size=8, training on 24 nodes, vLLM TP=32, gpus_per_node=8):

• Training claims domains A, B, C (24 nodes, 3 segments of 8).
• gpus_per_instance = 32, nodes_per_instance = 4.
• Inference claims 16 remaining nodes (domains D, E) in groups of 4.

Issues

N/A

Testing

Validated on the current branch HEAD on two clusters — GB200 (hsg) and H100 (cw-dfw) — across both SGLang and Megatron backends, with cluster.segment_size set.

GB200 (hsg) — NVL72 rack = 18 nodes × 4 GPU; one rack = one NVLink domain

Performance — Megatron MoE, cross-rack expert parallelism (the feature's target case):

• Recipe examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n4g-megatron.yaml (Qwen3-30B-A3B GRPO), expert_model_parallel_size=8 (each EP group spans 2 nodes), TP=PP=1.
• 4 nodes allocated across 2 NVLink racks; compared segment_size=2 (ON) vs null (OFF) in the same allocation (null ≡ main behavior, feature gated off). Mean of steps 6–12:

Stage	ON (s)	OFF (s)	Improvement
policy_and_reference_logprobs (EP all-to-all fwd pass)	4.60	6.29	~27%
Non-generation (training) work	32.41	34.08	~5%
policy_training (fwd+bwd+optim)	8.18	8.30	~1.5%
generation (vLLM, placement-insensitive)	31.24	30.88	~0%
Total step time	63.66	64.97	~2%

With the EP group kept inside one NVLink domain (ON), the expert all-to-all runs over NVLink instead of InfiniBand (OFF). The end-to-end step gain (~2%) is diluted by generation (which the feature does not affect); the training collectives — especially the logprob forward pass — improve ~5–27%.

Functional — SGLang generation with the feature:

• Recipe grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.yaml, 1 node × 4 GPU, segment_size=1, 5 steps: ✅ completed.

H100 (cw-dfw) — 8 GPU/node, no multi-node NVLink fabric

Functional "nothing broken" with the feature enabled (1 node × 8 GPU, segment_size=1, 5 steps each):

• SGLang GRPO — grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.yaml: ✅ completed.
• Megatron GRPO — grpo-llama3.2-1b-instruct-1n8g-megatron.yaml: ✅ completed.

On H100 the cross-rack perf benefit is not expected (NVLink domains are per-node); this confirms the feature degrades to a correct no-op and breaks nothing on that architecture.

Unit tests: tests/unit/distributed/test_topology_placement.py (30 tests) pass.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? (unit tests added in tests/unit/distributed/test_topology_placement.py)
• Did you run the unit tests and functional tests locally? (functional + perf runs on GB200/hsg — see Testing section)
• Did you add or update any necessary documentation?

Additional Information

• segment_size: null is added to all exemplar configs under examples/configs/*.yaml as the documented default.
• All five algorithms (GRPO, SFT, DPO, RM, distillation) support topology-aware placement via the same cluster.segment_size key.
• A prepare_segment_topology helper in virtual_cluster.py centralises the probe-and-select logic so algorithms don't duplicate it.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[youngeunkwon0405]

/okay to test b9bed0a

[youngeunkwon0405]

/okay to test 37452f1

[terrykong]

/ok to test 6874bf5

[terrykong]

/ok to test 4b82af4