feat: cherry-pick numa aware binding (#2613) to r0.7.0

nemo_rl/distributed/numa_utils.py

@@ -0,0 +1,181 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""NUMA-aware CPU affinity and memory binding for GPU workers.
+
+Uses a GPU→cpulist mapping file written by topology_probe.sh (in ray.sub)
+at node startup. The file path is communicated via the NRL_GPU_CPU_AFFINITY_FILE
+environment variable. See ray.sub for the writer side.
+
+Disable all binding with NRL_DISABLE_NUMA_BINDING=1.
+Disable only memory policy with NRL_DISABLE_NUMA_MEMBIND=1.
+"""
+
+import ctypes
+import ctypes.util
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+# IMPORTANT: This default path must stay in sync with topology_probe.sh in ray.sub.
+# The canonical path is set via the NRL_GPU_CPU_AFFINITY_FILE env var exported by ray.sub.
+GPU_CPU_AFFINITY_PATH = os.environ.get(
+    "NRL_GPU_CPU_AFFINITY_FILE", "/tmp/nrl_gpu_cpu_affinity"
+)
+
+
+def bind_to_gpu_numa(gpu_id: int) -> bool:
+    """Pin the current process to the NUMA-local CPUs and memory of the given GPU.
+
+    Reads the GPU→cpulist mapping written by topology_probe.sh at node
+    startup, then calls os.sched_setaffinity() for CPU pinning and
+    numa_set_membind() for memory policy. Best-effort: failures are
+    logged, never raised.
+
+    Args:
+        gpu_id: Node-global physical GPU index (``nvidia-smi`` numbering), which
+            is how the affinity file is keyed. Passed explicitly because
+            ``CUDA_VISIBLE_DEVICES`` lists all devices on the node under
+            ``RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1`` and so does not
+            identify a single worker's GPU. In a Ray actor this is
+            ``int(ray.get_gpu_ids()[0])``.
+
+    Returns True if CPU binding succeeded, False if skipped or failed.
+    Memory binding is attempted independently and logged separately.
+    """
+    if os.environ.get("NRL_DISABLE_NUMA_BINDING") == "1":
+        return False
+
+    gpu = str(gpu_id)
+    try:
+        with open(GPU_CPU_AFFINITY_PATH) as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                idx, cpulist = line.split(":", 1)
+                if idx == gpu:
+                    cpus = _parse_cpulist(cpulist)
+                    os.sched_setaffinity(0, cpus)
+                    logger.info("NUMA CPU binding: GPU %s → CPUs %s", gpu, cpulist)
+                    _set_numa_membind(cpus)
+                    return True
+        logger.debug("NUMA binding: GPU %s not found in %s", gpu, GPU_CPU_AFFINITY_PATH)
+    except FileNotFoundError:
+        logger.debug("NUMA binding skipped: %s not found", GPU_CPU_AFFINITY_PATH)
+    except Exception as exc:
+        logger.debug("NUMA binding skipped: %s", exc)
+    return False
+
+
+def resolve_visible_gpu_id(local_index: int) -> int | None:
+    """Map a process-local CUDA device index to its node-global physical GPU id.
+
+    ``CUDA_VISIBLE_DEVICES`` lists the physical GPU ids visible to this process
+    in device-index order, and ``local_index`` (e.g.
+    ``torch.cuda.current_device()``) indexes into that list. The affinity file is
+    keyed by the physical id, so return ``CUDA_VISIBLE_DEVICES[local_index]``.
+
+    ``CUDA_VISIBLE_DEVICES`` contents depend on the worker:
+      - vLLM TP>1 (``RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1``): the
+        per-instance device subset, e.g. ``"4,5"``.
+      - vLLM TP=1: a single isolated device, so ``local_index`` is 0.
+
+    Returns the physical GPU id, or None if it cannot be resolved (unset CVD,
+    index out of range, or non-integer entries such as MIG UUIDs).
+    """
+    cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+    if not cvd:
+        return None
+    devices = cvd.split(",")
+    if local_index < 0 or local_index >= len(devices):
+        return None
+    try:
+        return int(devices[local_index])
+    except ValueError:
+        return None
+
+
+def _load_libnuma() -> ctypes.CDLL | None:
+    """Load libnuma, returning None if unavailable."""
+    try:
+        return ctypes.CDLL("libnuma.so.1")
+    except OSError:
+        return None
+
+
+def _get_numa_node(libnuma: ctypes.CDLL, cpus: set[int]) -> int:
+    """Return the NUMA node for the given CPU set, or -1 on failure."""
+    libnuma.numa_node_of_cpu.restype = ctypes.c_int
+    return libnuma.numa_node_of_cpu(min(cpus))
+
+
+def _set_numa_membind(cpus: set[int]) -> bool:
+    """Hard-bind memory allocations to the NUMA node of the given CPUs."""
+    if os.environ.get("NRL_DISABLE_NUMA_MEMBIND") == "1":
+        return False
+
+    libnuma = _load_libnuma()
+    if libnuma is None:
+        logger.debug("NUMA membind skipped: libnuma.so.1 not available")
+        return False
+
+    try:
+        numa_node = _get_numa_node(libnuma, cpus)
+        if numa_node < 0:
+            logger.debug(
+                "NUMA membind skipped: numa_node_of_cpu(%d) returned %d",
+                min(cpus),
+                numa_node,
+            )
+            return False
+
+        libnuma.numa_allocate_nodemask.restype = ctypes.c_void_p
+        libnuma.numa_bitmask_setbit.argtypes = [ctypes.c_void_p, ctypes.c_uint]
+        libnuma.numa_bitmask_setbit.restype = ctypes.c_void_p
+        libnuma.numa_set_membind.argtypes = [ctypes.c_void_p]
+        libnuma.numa_bitmask_free.argtypes = [ctypes.c_void_p]
+
+        nodemask = libnuma.numa_allocate_nodemask()
+        if not nodemask:
+            logger.debug("NUMA membind skipped: numa_allocate_nodemask returned NULL")
+            return False
+
+        try:
+            libnuma.numa_bitmask_setbit(nodemask, numa_node)
+            libnuma.numa_set_membind(nodemask)
+        finally:
+            libnuma.numa_bitmask_free(nodemask)
+
+        logger.info(
+            "NUMA membind: hard-bound to node %d (from CPU %d)", numa_node, min(cpus)
+        )
+        return True
+    except Exception as exc:
+        logger.debug("NUMA membind skipped: %s", exc)
+        return False
+
+
+def _parse_cpulist(cpulist: str) -> set[int]:
+    """Parse a Linux cpulist string like '0-71' into a set of ints."""
+    cpus: set[int] = set()
+    for part in cpulist.split(","):
+        part = part.strip()
+        if "-" in part:
+            lo, hi = part.split("-", 1)
+            cpus.update(range(int(lo), int(hi) + 1))
+        else:
+            cpus.add(int(part))
+    return cpus

nemo_rl/models/generation/vllm/vllm_backend.py

@@ -92,6 +92,25 @@ def _read_mtp_layer_weights_from_checkpoint(
 
 
 class VllmInternalWorkerExtension:
+    def bind_numa(self) -> bool:
+        """Pin this TP worker to its GPU's NUMA-local CPUs/memory.
+
+        Invoked via ``collective_rpc`` on each vLLM TP worker once the engine
+        (and CUDA) is up, so the worker's physical GPU id is resolved from its
+        local device index (see ``resolve_visible_gpu_id``).
+        """
+        import torch
+
+        from nemo_rl.distributed.numa_utils import (
+            bind_to_gpu_numa,
+            resolve_visible_gpu_id,
+        )
+
+        gpu_id = resolve_visible_gpu_id(torch.cuda.current_device())
+        if gpu_id is None:
+            return False
+        return bind_to_gpu_numa(gpu_id)
+
     def init_collective(
         self,
         rank_prefix: int,

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -182,6 +182,19 @@ def __init__(
                 _load_model() later to perform the heavy model loading. This
                 enables overlapping vLLM model loading with NeMo Gym init.
         """
+        from nemo_rl.distributed.numa_utils import bind_to_gpu_numa
+
+        # Only bind single-GPU workers to their GPU's NUMA node.
+        # For TP>1 workers, the parent process spans multiple NUMA nodes;
+        # binding it would incorrectly constrain the EngineCore subprocess
+        # (which inherits sched_setaffinity + numa_set_membind via fork).
+        # Individual TP workers get their own NUMA binding via collective_rpc
+        # in post_init / post_init_async.
+        # ray.get_gpu_ids()[0] is this worker's physical GPU index, which keys
+        # the affinity file.
+        if bundle_indices is not None and len(bundle_indices) == 1:
+            bind_to_gpu_numa(int(ray.get_gpu_ids()[0]))
+
         self._init_config(
             config, bundle_indices, fraction_of_gpus, seed, extra_env_vars
         )
@@ -589,6 +602,8 @@ def _create_engine(self, llm_kwargs: dict[str, Any]) -> None:
         self.llm = vllm.LLM(**llm_kwargs)
 
     def post_init(self):
+        if self.llm is not None:
+            self.llm.collective_rpc("bind_numa", args=tuple())
         self.vllm_device_ids = self.report_device_id()
         if self._mtp_load_from_disk:
             self.llm.collective_rpc(

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -422,6 +422,8 @@ def clear_vllm_logger_metrics(self) -> None:
             self.generation_tokens = []
 
     async def post_init_async(self):
+        if self.llm is not None:
+            await self.llm.collective_rpc("bind_numa", args=tuple())
         self.vllm_device_ids = await self.report_device_id_async()
         if self._mtp_load_from_disk:
             await self.llm.collective_rpc(

nemo_rl/models/policy/workers/dtensor_policy_worker.py

@@ -200,6 +200,14 @@ def __init__(
         **kwargs: Any,
     ):
         """Initialize the DTensorPolicyWorker."""
+        from nemo_rl.distributed.numa_utils import bind_to_gpu_numa
+
+        # Pin to this worker's GPU-local CPUs/memory before CUDA init or model
+        # load; FSDP's D2H paths (weight refit, optimizer/checkpoint offload)
+        # benefit. ray.get_gpu_ids()[0] is the physical GPU index that keys the
+        # affinity file, and reading it does not initialize CUDA.
+        bind_to_gpu_numa(int(ray.get_gpu_ids()[0]))
+
         self.tokenizer = tokenizer
         self.processor = processor
         self.is_vlm = processor is not None

nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

@@ -231,6 +231,14 @@ def __init__(
         # Apply TE patch until TE is upgraded to 2.10.0
         apply_transformer_engine_patch()
 
+        from nemo_rl.distributed.numa_utils import bind_to_gpu_numa
+
+        # Pin to this worker's GPU-local CPUs/memory before model load; FSDP's
+        # D2H paths (weight refit, optimizer/checkpoint offload) benefit.
+        # ray.get_gpu_ids()[0] is the physical GPU index that keys the affinity
+        # file, and reading it does not initialize CUDA.
+        bind_to_gpu_numa(int(ray.get_gpu_ids()[0]))
+
         # Store configuration
         self.cfg = config
 

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -268,6 +268,14 @@ def __init__(
         # Apply patch from https://github.com/NVIDIA/TransformerEngine/pull/2286/files
         apply_transformer_engine_patch()
 
+        from nemo_rl.distributed.numa_utils import bind_to_gpu_numa
+
+        # local_rank (== ray.get_gpu_ids()[0]) is the physical GPU index that
+        # keys the affinity file. Pass it explicitly: CUDA_VISIBLE_DEVICES lists
+        # all node devices here (RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1,
+        # set by configure_worker), so it can't identify this worker's GPU.
+        bind_to_gpu_numa(local_rank)
+
         self.cfg = config
         self._router_replay_enabled = router_replay_enabled(config)
 

nemo_rl/models/value/workers/dtensor_value_worker_v2.py

@@ -148,6 +148,14 @@ def __init__(
         # Apply patches
         apply_transformer_engine_patch()
 
+        from nemo_rl.distributed.numa_utils import bind_to_gpu_numa
+
+        # Pin to this worker's GPU-local CPUs/memory before model load; FSDP's
+        # D2H paths (weight refit, optimizer/checkpoint offload) benefit.
+        # ray.get_gpu_ids()[0] is the physical GPU index that keys the affinity
+        # file, and reading it does not initialize CUDA.
+        bind_to_gpu_numa(int(ray.get_gpu_ids()[0]))
+
         # Store configuration and tokenizer
         self.cfg = config
         self.tokenizer = tokenizer

nemo_rl/models/value/workers/megatron_value_worker.py

@@ -395,6 +395,15 @@ def __init__(
 
         apply_transformer_engine_patch()
 
+        from nemo_rl.distributed.numa_utils import bind_to_gpu_numa
+
+        # Pin to this worker's GPU-local CPUs/memory before model load, matching
+        # the policy workers. local_rank (== ray.get_gpu_ids()[0]) is the physical
+        # GPU index that keys the affinity file. Pass it explicitly:
+        # CUDA_VISIBLE_DEVICES lists all node devices here
+        # (RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1).
+        bind_to_gpu_numa(local_rank)
+
         self.cfg = config
         self.rank = get_rank_safe()
 

pyrefly.toml

@@ -115,6 +115,7 @@ project-includes = [
   "nemo_rl/distributed/__init__.py",
   "nemo_rl/distributed/collectives.py",
   "nemo_rl/distributed/named_sharding.py",
+  "nemo_rl/distributed/numa_utils.py",
   "nemo_rl/distributed/ray_actor_environment_registry.py",
   "nemo_rl/distributed/virtual_cluster.py",
   "nemo_rl/distributed/worker_group_utils.py",

ray.sub

@@ -387,6 +387,34 @@ else
   fi
 fi
 
+# Write GPU→cpulist mapping for NUMA binding.
+# IMPORTANT: The env var name NRL_GPU_CPU_AFFINITY_FILE and default path must stay
+# in sync with GPU_CPU_AFFINITY_PATH in nemo_rl/distributed/numa_utils.py.
+export NRL_GPU_CPU_AFFINITY_FILE="/tmp/nrl_gpu_cpu_affinity"
+# nvidia-smi topo -m's CPU Affinity column is unreliable on GB200 (empty for
+# GPUs not directly attached to the socket). Use NUMA Affinity (always
+# populated, at field NF-1 since GPU NUMA ID is last) and look up the
+# node-local CPU list from sysfs. On GB200 the NUMA Affinity column can be a
+# list like "0,2-17" (the GPU-local CPU NUMA node plus the GPU's HBM NUMA
+# nodes); take the first entry, which is the local CPU NUMA node.
+nvidia-smi topo -m 2>/dev/null | awk '/^GPU[0-9]/ {
+  gpu = \$1; sub(/GPU/, "", gpu)
+  numa = \$(NF-1)
+  sub(/[,-].*/, "", numa)
+  if (numa ~ /^[0-9]+\$/) print gpu, numa
+}' | while read -r _gpu _numa; do
+  _cpulist=\$(cat "/sys/devices/system/node/node\${_numa}/cpulist" 2>/dev/null)
+  if [[ -n "\$_cpulist" ]]; then
+    echo "\${_gpu}:\${_cpulist}"
+  fi
+done > "\$NRL_GPU_CPU_AFFINITY_FILE" || true
+if [[ -s "\$NRL_GPU_CPU_AFFINITY_FILE" ]]; then
+  echo "NUMA affinity map written to \$NRL_GPU_CPU_AFFINITY_FILE:"
+  cat "\$NRL_GPU_CPU_AFFINITY_FILE"
+else
+  echo "WARNING: Could not generate NUMA affinity map (nvidia-smi topo unavailable)"
+fi
+
 # Use \\\" so that when --resources="\$RAY_RESOURCES" expands, we pass valid JSON to ray
 # IMPORTANT: The key names "nvlink_domain_" and "topo_rank" below must stay in sync
 # with the constants NVLINK_DOMAIN_PREFIX and TOPO_RANK_KEY defined in