Pythorch pinning example - rFronteddu/general_wiki GitHub Wiki
This script binds the main training process and each DataLoader worker process to the GPU’s local NUMA node to prevent cross-NUMA memory access. In the DataLoader, we pass a closure-based worker_init_fn that reapplies the precomputed NUMA binding inside each worker. And we do this without touching any CUDA APIs in the worker.
Because some launchers, container runtimes, or kernels do not reliably propagate NUMA policy to children, we explicitly reapply and verify the binding in every forked worker. It’s not safe to rely on inheritance alone.
Set pin_memory=True and use non_blocking=True on H2D copies so that page-locked host buffers stay on the correct NUMA node. Prefer persistent_workers=True to avoid re-forking workers and losing their affinity between epochs. Do not call torch.cuda.* in worker_init_fn. Instead, pass the GPU index using a closure or environment variable. Data preparation and batch loading can then happen entirely in local memory. This way, your GPUs stay busy and never need to pause for a remote‐NUMA hop. With this code, you get robust, topology‐aware affinity on any Linux server with libnuma and numactl installed.
import os
import re
import glob
import subprocess
import psutil
import ctypes
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset
from functools import partial
# Optional: NVML is preferred for GPU↔NUMA mapping
try:
import pynvml as nvml # pip install nvidia-ml-py3
_HAS_NVML = True
except Exception:
_HAS_NVML = False
# --- libnuma for memory binding
_libnuma = ctypes.CDLL("libnuma.so")
if _libnuma.numa_available() < 0:
raise RuntimeError("NUMA not available on this system")
_libnuma.numa_run_on_node.argtypes = [ctypes.c_int]
_libnuma.numa_set_preferred.argtypes = [ctypes.c_int]
def parse_physical_cpu_list(phys_str: str):
"""Parse '0-3,8-11' -> [0,1,2,3,8,9,10,11]."""
cpus = []
if not phys_str:
return cpus
for part in phys_str.split(','):
part = part.strip()
if not part:
continue
if '-' in part:
start, end = map(int, part.split('-'))
cpus.extend(range(start, end + 1))
else:
cpus.append(int(part))
return cpus
def get_numa_cpus_for_node(node: int):
"""Read /sys/devices/system/node/node{node}/cpulist."""
path = f"/sys/devices/system/node/node{node}/cpulist"
with open(path, "r") as f:
return parse_physical_cpu_list(f.read().strip())
def get_numa_cpus_and_memory():
"""Return (current_cpu_mask, preferred_node) from numactl --show."""
out = subprocess.run(
["numactl", "--show"],
capture_output=True,
text=True
).stdout
phys = re.search(r"physcpubind:\s*([\d,\-\s]+)", out).group(1)
cpus = parse_physical_cpu_list(phys)
node = int(
re.search(r"preferred node:\s*(-?\d+)", out).group(1)
)
return cpus, node
def get_gpu_numa_node(device: int) -> int:
"""
Determine NUMA node for a GPU (prefer NVML; fall back to sysfs;
final fallback to current preferred node).
"""
# NVML path (preferred)
if _HAS_NVML:
try:
nvml.nvmlInit()
props = torch.cuda.get_device_properties(device)
pci = props.pci_bus_id
# Normalize to 8-hex-digit domain if needed for NVML
try:
domain, bus, devfn = pci.split(':')
if len(domain) < 8:
domain = domain.rjust(8, '0')
pci8 = f"{domain}:{bus}:{devfn}"
except ValueError:
pci8 = pci
try:
handle = nvml.nvmlDeviceGetHandleByPciBusId_v2(pci8)
except AttributeError:
handle = nvml.nvmlDeviceGetHandleByPciBusId(pci8)
# Direct NUMA ID if driver exposes it
try:
numa_id = nvml.nvmlDeviceGetNUMANodeId(handle)
if isinstance(numa_id, int) and numa_id >= 0:
return numa_id
except Exception:
pass
# Derive from NVML CPU affinity
cpu_count = psutil.cpu_count(logical=True)
elems = (cpu_count + 63) // 64
mask = nvml.nvmlDeviceGetCpuAffinity(handle, elems)
cpus = []
for i, m in enumerate(mask):
m = int(m)
for b in range(64):
if m & (1 << b):
cpu_id = i * 64 + b
if cpu_id < cpu_count:
cpus.append(cpu_id)
# Build CPU→NUMA map from sysfs and choose majority node
cpu2node = {}
for node_path in sorted(
glob.glob("/sys/devices/system/node/node*")
):
node_id = int(
os.path.basename(node_path).replace("node", "")
)
with open(os.path.join(node_path, "cpulist"), "r") as f:
for c in parse_physical_cpu_list(f.read().strip()):
cpu2node[c] = node_id
counts = {}
for c in cpus:
n = cpu2node.get(c)
if n is not None:
counts[n] = counts.get(n, 0) + 1
if counts:
return max(
counts.items(),
key=lambda kv: kv[1]
)[0]
except Exception:
pass
# sysfs fallback
try:
props = torch.cuda.get_device_properties(device)
pci = props.pci_bus_id
sysfs_path = f"/sys/bus/pci/devices/{pci}/numa_node"
with open(sysfs_path, "r") as f:
val = int(f.read().strip())
return val if val >= 0 else 0
except Exception:
pass
# Last resort: current preferred node
_, node = get_numa_cpus_and_memory()
return node if node >= 0 else 0
def set_numa_affinity(node: int):
"""Bind current process to CPUs and memory of the given NUMA node."""
cpus = get_numa_cpus_for_node(node)
# IMPORTANT: CPUs of target node
psutil.Process(os.getpid()).cpu_affinity(cpus)
_libnuma.numa_run_on_node(node)
_libnuma.numa_set_preferred(node)
print(
f"PID={os.getpid()} bound to NUMA node {node} "
f"(CPUs={cpus})"
)
return cpus
def _worker_init_fn(worker_id: int, node: int, cpus: list):
"""Reapply binding in each DataLoader worker (no CUDA calls here)."""
psutil.Process(os.getpid()).cpu_affinity(cpus)
_libnuma.numa_run_on_node(node)
_libnuma.numa_set_preferred(node)
print(
f"Worker {worker_id} "
f"(PID={os.getpid()}) "
f"bound to NUMA node {node}"
)
# ----- Example usage below -----
class MyDataset(Dataset):
def __len__(self):
return 1024
def __getitem__(self, idx):
return torch.randn(224 * 224 * 3, device="cpu")
def main():
# DDP setup
dist.init_process_group(
backend="nccl",
init_method="env://"
)
device = torch.cuda.current_device()
# Determine GPU's NUMA node and bind this process
gpu_node = get_gpu_numa_node(device)
cpus = set_numa_affinity(gpu_node)
# Build dataloader with closure-based worker_init_fn
dataset = MyDataset()
init_fn = partial(
_worker_init_fn,
node=gpu_node,
cpus=cpus
)
dataloader = DataLoader(
dataset,
batch_size=32,
num_workers=4,
pin_memory=True,
persistent_workers=True,
worker_init_fn=init_fn,
prefetch_factor=2,
)
# Model and DDP
model = torch.nn.Linear(
224 * 224 * 3,
10,
bias=True
).to("cuda")
ddp_model = DDP(
model,
device_ids=[device],
static_graph=True
)
for batch in dataloader:
batch = batch.to(
"cuda",
non_blocking=True
)
out = ddp_model(batch)
# ... loss, backward, optimizer ...
if __name__ == "__main__":
main()