feat: Numa aware binding by youngeunkwon0405 · Pull Request #2613 · NVIDIA-NeMo/RL

youngeunkwon0405 · 2026-05-28T20:41:00Z

What does this PR do ?

Adds NUMA-aware CPU affinity + memory binding for GPU workers (Megatron, DTensor/FSDP, vLLM, and the Megatron PPO value worker). Each model-owning worker pins itself to its GPU's NUMA-local CPUs and binds memory to the local NUMA node at init, improving device↔host memory bandwidth on NUMA systems (e.g. GB200 Grace–Grace).

A node-startup probe in ray.sub writes a GPU→cpulist affinity map (from nvidia-smi topo -m's NUMA Affinity column + sysfs), consumed by nemo_rl.distributed.numa_utils.bind_to_gpu_numa(). Disable with NRL_DISABLE_NUMA_BINDING=1 (memory-only opt-out: NRL_DISABLE_NUMA_MEMBIND=1).

Originally cherry-picked from the internal Numa binding branch; rebased onto main after the topology PR (#2612) merged.

Stacked on #2924. The topology bugfix that used to be the base commit of this branch is now its own PR — #2924 (fix(topology): exclude unknown-NVLink-domain nodes from segment selection, base main). This PR stacks on top of it, so its diff shows only the NUMA changes. #2924 fixes a pre-existing select_segment_nodes bug from #2612 (an unschedulable {"unknown": 0.001} placement group on partial / GPU-less-head clusters).

Cross-PR dependency with #2837. The DTensor PPO value worker added in #2837 (feat: Support for dtensor ppo) must also call bind_to_gpu_numa(). Whichever of #2613 / #2837 merges second should add that one line to the new value worker.

Testing

Validated on a GB200 node, 4 GPUs (H100 skipped — down).

Perf — D2H bandwidth (the feature's target)

Committed benchmark tests/unit/distributed/test_numa_utils.py::TestNUMABindingBenchmark::test_d2h_with_numa_binding (256 MB float32, 50 iters, GPU 0) compares unbound (≡ main behavior) vs NUMA-bound (this branch) D2H copy time on the same GPU:

Case	D2H ms/iter
Unbound (memory spread across all NUMA nodes)	2.797
NUMA-bound (GPU-local node, this branch)	1.977
Speedup	1.42× (≈29% faster) — stable across 3 runs

Functional — "nothing broken" (GB200, 1 node × 4 GPU)

SGLang GRPO (grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.yaml, segment_size=1, 5 steps): ✅ reached step 5, clean exit.
Megatron GRPO (grpo-llama3.2-1b-instruct-1n8g-megatron.yaml, 8 steps) with NUMA OFF (NRL_DISABLE_NUMA_BINDING=1) and ON (default): ✅ both reached step 8, clean exit. Mean step time (steps 3–8): OFF 13.47 s, ON 13.51 s — within run-to-run noise. (End-to-end step time on a small single-node run is compute-dominated, not D2H-bound; the bandwidth win above is the relevant signal and would matter more on D2H-heavy / multi-node workloads.)

Unit tests

tests/unit/distributed/test_numa_utils.py — 17 pass / 1 skip (skip = GPU-gated benchmark when CUDA_VISIBLE_DEVICES unset). CPU-binding tests derive the cpulist from os.sched_getaffinity(0) so they are host-portable.
tests/unit/distributed/test_topology_placement.py::TestSelectSegmentNodes — regression tests for the unknown-domain exclusion (pass with the fix, fail without it).

Issues

N/A

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests? (numa_utils host-portability + topology-selection regression tests)
Did you run the unit tests and functional tests locally? (unit tests + GB200 functional/perf — see Testing)
Did you add or update any necessary documentation?

Additional Information

GB200 NUMA Affinity format. nvidia-smi topo -m reports the NUMA Affinity column as a list, e.g. 0,2-17 — the GPU-local CPU NUMA node (0) plus the GPU's HBM NUMA nodes (2-17), which GB200 exposes as additional coherent NUMA nodes. The probe (in ray.sub and the test helper) takes the first entry (the CPU node). Without this, the affinity map is empty on GB200 and binding never engages.

copy-pr-bot · 2026-05-28T20:41:05Z

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 · 2026-05-28T20:42:46Z

/okay to test f371eb0

youngeunkwon0405 · 2026-05-28T20:50:15Z

/okay to test 81d5362

youngeunkwon0405 · 2026-05-28T21:04:49Z

/okay to test 71a42f0

terrykong

@ananthsub to review

copy-pr-bot · 2026-07-07T12:56:29Z

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

ananthsub · 2026-07-07T12:58:09Z

/ok to test 1872748

ananthsub

@yuki-97 could you also take another look as I've made changes based on your feedback

yuki-97

thanks @ananthsub , LGTM, but have a new comment about the newly added dtensor_value_worker_v2.

ananthsub · 2026-07-08T15:56:43Z

/ok to test 5f49787

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

nvidia-smi topo -m's CPU Affinity column is empty on GB200 for GPUs not directly attached to a CPU socket (only GPU0/1 get "0-63"; GPU2/3 get an empty cell that collapses under awk's default whitespace splitting). The NF-2 extraction in the original cherry-pick then grabs a stale "NODE" topology label instead of a CPU range. Use NUMA Affinity instead (always populated, reliably at NF-1 since GPU NUMA ID is the last column) and resolve the CPU list via /sys/devices/system/node/nodeN/cpulist. Produces correct mappings for all 4 GPUs on a GB200 node (e.g. 0:0-71, 1:0-71, 2:72-143, 3:72-143). Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Megatron and vLLM workers already call bind_to_gpu_numa(), but the DTensor/FSDP policy workers (dtensor_policy_worker.py and dtensor_policy_worker_v2.py) did not. FSDP is a first-class training backend and its D2H paths (weight refit, optimizer/checkpoint offload) benefit from NUMA-local CPU/memory binding. Add the bind early in __init__, before CUDA init / model load, mirroring the other workers. Signed-off-by: Terry Kong <terryk@nvidia.com>

The PPO value worker (MegatronValueWorkerImpl) landed in main after this branch was first written and runs a Megatron model on GPU like the policy worker, but did not call bind_to_gpu_numa(). Add the bind early in __init__ for parity with the Megatron/DTensor policy workers and vLLM workers. Note: the DTensor PPO value worker (PR #2837, feat: Support for dtensor ppo) will need the same call; whichever of the two PRs merges second must add it. Signed-off-by: Terry Kong <terryk@nvidia.com>

_write_affinity_file_from_topo parsed the 'CPU Affinity' column (parts[-3]), but production ray.sub deliberately avoids that column (empty/unreliable on GB200 for GPUs not directly attached to a socket) and instead reads the 'NUMA Affinity' column ($(NF-1)) then looks up /sys/devices/system/node/node<N>/cpulist. The GPU-gated benchmark therefore validated a parsing path production abandoned. Mirror ray.sub exactly (NUMA Affinity + sysfs) so the two can't drift. Signed-off-by: Terry Kong <terryk@nvidia.com>

test_successful_cpu_binding and test_multi_gpu_cvd_uses_first hard-coded the cpulist 72-143 and asserted sched_getaffinity == set(range(72,144)). On any host with < 144 CPUs, os.sched_setaffinity raises OSError (swallowed by bind_to_gpu_numa's broad except -> returns False), so the asserts failed off the large CI nodes (e.g. on a 24-CPU dev box). Add a _split_available_cpus() helper that derives two CPU groups from os.sched_getaffinity(0), and drive both tests from it so they exercise a real sched_setaffinity on any host (skipping only when < 2 CPUs are available). Signed-off-by: Terry Kong <terryk@nvidia.com>

On GB200, nvidia-smi topo -m's NUMA Affinity column is a list such as "0,2-17" — the GPU-local CPU NUMA node (0) plus the GPU's HBM NUMA nodes (2-17) — not a single integer. The probe in ray.sub and its mirror in the test helper both required `^[0-9]+$` / str.isdigit(), so they rejected every GPU row on this hardware: ray.sub produced an empty affinity map (NUMA binding never engaged) and the D2H benchmark skipped with 'Could not parse nvidia-smi topo output'. Take the first entry of the NUMA Affinity list (the GPU-local CPU NUMA node) before the digit check, in both the ray.sub awk probe and _write_affinity_file_from_topo. Verified on a real GB200: with this the benchmark runs and NUMA binding gives a ~1.4x (+40%) D2H bandwidth speedup. Signed-off-by: Terry Kong <terryk@nvidia.com>

…EVICES bind_to_gpu_numa() now takes a required gpu_id (the node-global physical GPU index that keys the affinity file) instead of reading CUDA_VISIBLE_DEVICES. Under RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1 (Megatron policy/value and vLLM TP workers) CUDA_VISIBLE_DEVICES lists all node devices, so it cannot identify a worker's own GPU. Call sites pass the physical index they already know: - Megatron policy/value: local_rank (== int(ray.get_gpu_ids()[0])). - DTensor v1/v2 and vLLM TP=1 parent: int(ray.get_gpu_ids()[0]) (does not initialize CUDA, so safe for the pre-CUDA-init binders). - vLLM TP workers (bind_numa via collective_rpc): resolve_visible_gpu_id() maps torch.cuda.current_device() through CUDA_VISIBLE_DEVICES, which is correct for per-instance device subsets (e.g. "4,5"). Add resolve_visible_gpu_id() plus unit tests covering the subset, full-node, isolated, and unresolvable cases; update the CPU-binding tests to pass an explicit gpu_id. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

DTensorValueWorkerV2 (added in #2837) runs a model on GPU like the DTensor policy workers but did not call bind_to_gpu_numa(). Add the bind early in __init__ for parity, passing the physical GPU index int(ray.get_gpu_ids()[0]). Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

ananthsub · 2026-07-08T18:17:28Z

/ok to test 9587146

This reverts commit 9e01af6. Signed-off-by: sna <sna@nvidia.com>

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com> Signed-off-by: Terry Kong <terryk@nvidia.com> Co-authored-by: Ananth Subramaniam <ansubramania@nvidia.com> Co-authored-by: Terry Kong <terryk@nvidia.com> Signed-off-by: Kajal Jain <kajalj@nvidia.com>

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com> Signed-off-by: Terry Kong <terryk@nvidia.com> Co-authored-by: Youngeun Kwon <yukih@nvidia.com> Co-authored-by: Ananth Subramaniam <ansubramania@nvidia.com> Co-authored-by: Terry Kong <terryk@nvidia.com>

youngeunkwon0405 requested review from ananthsub, terrykong and yfw May 28, 2026 20:41

youngeunkwon0405 requested review from a team as code owners May 28, 2026 20:41

youngeunkwon0405 changed the title ~~Feat: Numa aware binding~~ feat: Numa aware binding May 28, 2026

youngeunkwon0405 added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label May 28, 2026

copy-pr-bot Bot temporarily deployed to public May 28, 2026 20:43 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci May 28, 2026 20:43 Error

copy-pr-bot Bot temporarily deployed to public May 28, 2026 20:43 Inactive

youngeunkwon0405 force-pushed the youngeunk/numa-aware-binding branch 2 times, most recently from 3945759 to 81d5362 Compare May 28, 2026 20:49

copy-pr-bot Bot temporarily deployed to public May 28, 2026 20:50 Inactive

copy-pr-bot Bot temporarily deployed to public May 28, 2026 20:51 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci May 28, 2026 20:52 Error

copy-pr-bot Bot temporarily deployed to public May 28, 2026 20:55 Inactive

youngeunkwon0405 force-pushed the youngeunk/numa-aware-binding branch from 81d5362 to 71a42f0 Compare May 28, 2026 20:59

youngeunkwon0405 requested a review from a team as a code owner May 28, 2026 20:59

terrykong reviewed Jun 25, 2026

View reviewed changes

Base automatically changed from youngeunk/fix-topology-unknown-domain to main June 26, 2026 12:33

yuki-97 force-pushed the youngeunk/numa-aware-binding branch from f21c2b7 to 6846a5d Compare June 26, 2026 12:33

yuki-97 reviewed Jun 26, 2026

View reviewed changes

Comment thread nemo_rl/distributed/numa_utils.py Outdated

ananthsub force-pushed the youngeunk/numa-aware-binding branch from 6846a5d to 1872748 Compare July 7, 2026 12:56

ananthsub previously approved these changes Jul 7, 2026

View reviewed changes

yuki-97 reviewed Jul 8, 2026

View reviewed changes

Comment thread nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

ananthsub dismissed their stale review via 5f49787 July 8, 2026 15:56

ananthsub and others added 9 commits July 8, 2026 11:11

add NUMA-aware CPU affinity and memory binding for GPU workers

09a5234

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

ananthsub force-pushed the youngeunk/numa-aware-binding branch from 5f49787 to 9587146 Compare July 8, 2026 18:16

yuki-97 approved these changes Jul 9, 2026

View reviewed changes

yuki-97 merged commit 9e01af6 into main Jul 9, 2026
79 of 80 checks passed

yuki-97 deleted the youngeunk/numa-aware-binding branch July 9, 2026 05:39

seonjinn added a commit to seonjinn/RL that referenced this pull request Jul 14, 2026

Revert "feat: Numa aware binding (NVIDIA-NeMo#2613)"

6e735fa

This reverts commit 9e01af6. Signed-off-by: sna <sna@nvidia.com>

kajalj22 mentioned this pull request Jul 20, 2026

cp: feat: Numa aware binding (#2613) into r0.7.0 #3299

Closed

2 tasks

terrykong mentioned this pull request Aug 15, 2026

perf: support fine-grained activation offloading #2279

Merged

4 tasks

Navigation Menu

Uh oh!

feat: Numa aware binding - #2613

feat: Numa aware binding#2613
yuki-97 merged 9 commits into
mainfrom
youngeunk/numa-aware-binding

youngeunkwon0405 commented May 28, 2026 •

edited by terrykong

Loading

Uh oh!

copy-pr-bot Bot commented May 28, 2026

Uh oh!

youngeunkwon0405 commented May 28, 2026

Uh oh!

youngeunkwon0405 commented May 28, 2026

Uh oh!

youngeunkwon0405 commented May 28, 2026

Uh oh!

terrykong left a comment

Uh oh!

Uh oh!

copy-pr-bot Bot commented Jul 7, 2026

Uh oh!

ananthsub commented Jul 7, 2026

Uh oh!

ananthsub left a comment

Uh oh!

yuki-97 left a comment

Uh oh!

Uh oh!

ananthsub commented Jul 8, 2026

Uh oh!

ananthsub commented Jul 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

youngeunkwon0405 commented May 28, 2026 • edited by terrykong Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Testing

Perf — D2H bandwidth (the feature's target)

Functional — "nothing broken" (GB200, 1 node × 4 GPU)

Unit tests

Issues

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented May 28, 2026

Uh oh!

youngeunkwon0405 commented May 28, 2026

Uh oh!

youngeunkwon0405 commented May 28, 2026

Uh oh!

youngeunkwon0405 commented May 28, 2026

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

copy-pr-bot Bot commented Jul 7, 2026

Uh oh!

ananthsub commented Jul 7, 2026

Uh oh!

ananthsub left a comment

Choose a reason for hiding this comment

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ananthsub commented Jul 8, 2026

Uh oh!

ananthsub commented Jul 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

youngeunkwon0405 commented May 28, 2026 •

edited by terrykong

Loading