[deps] support GB200 machines - #1936
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for aarch64/GB200 architectures by specifying corresponding wheel URLs in pyproject.toml and updates the NUMA binding logic in worker.py to only bind to NUMA nodes that contain CPUs. Feedback was provided regarding the NUMA binding logic: using torch.cuda.device_count() inside Ray-managed environments can return 1 due to CUDA_VISIBLE_DEVICES isolation, leading to NUMA imbalance. A suggestion was made to query the system device files to determine the total number of physical GPUs instead.
| num_gpu_per_numa_node = max(1, torch.cuda.device_count() // len(cpu_nodes)) | ||
| target_nid = cpu_nodes[min(len(cpu_nodes) - 1, self._local_rank // num_gpu_per_numa_node)] |
There was a problem hiding this comment.
In Ray-managed environments, torch.cuda.device_count() returns 1 inside each worker process because Ray overrides CUDA_VISIBLE_DEVICES to isolate GPUs per actor. Consequently, num_gpu_per_numa_node evaluates to 1, which causes all workers with self._local_rank >= 1 to bind to the last CPU-containing NUMA node (e.g., node 1), leaving node 0 underutilized and causing severe NUMA imbalance and memory bandwidth bottlenecks.
To fix this, we should determine the total number of physical GPUs on the node by querying the system device files (e.g., /dev/nvidia* or /dev/dri/card*), falling back to a sensible default (like 8) if CUDA is available, or 1 if running on CPU.
| num_gpu_per_numa_node = max(1, torch.cuda.device_count() // len(cpu_nodes)) | |
| target_nid = cpu_nodes[min(len(cpu_nodes) - 1, self._local_rank // num_gpu_per_numa_node)] | |
| import glob | |
| total_gpus = len(glob.glob("/dev/nvidia[0-9]*")) or len(glob.glob("/dev/dri/card*")) | |
| if total_gpus == 0: | |
| total_gpus = 8 if torch.cuda.is_available() else 1 | |
| num_gpu_per_numa_node = max(1, total_gpus // len(cpu_nodes)) | |
| target_nid = cpu_nodes[min(len(cpu_nodes) - 1, self._local_rank // num_gpu_per_numa_node)] |
kouroshHakha
left a comment
kouroshHakha
left a comment
There was a problem hiding this comment.
Just one comment. O.w. looks good.
| cpu_nodes = [] | ||
| node_root = "/sys/devices/system/node" | ||
| try: | ||
| node_names = sorted( | ||
| (n for n in os.listdir(node_root) if n.startswith("node") and n[4:].isdigit()), | ||
| key=lambda n: int(n[4:]), | ||
| ) | ||
| except OSError: | ||
| node_names = [] | ||
| for name in node_names: | ||
| try: | ||
| with open(os.path.join(node_root, name, "cpulist")) as f: | ||
| if f.read().strip(): | ||
| cpu_nodes.append(int(name[4:])) | ||
| except OSError: | ||
| continue | ||
| if not cpu_nodes: | ||
| cpu_nodes = [0] |
There was a problem hiding this comment.
does this generalize to non gb machines?
2 participants
Add custom wheels for aarch64/GB200.
Tested gsm8k runs for both megatron and fsdp backends.