From 1c4c6634312312ad08a842db1b60ec1338b7b306 Mon Sep 17 00:00:00 2001 From: George Stykalin Date: Wed, 15 Apr 2026 12:24:37 +0300 Subject: [PATCH] Re-enable InfiniBand for distributed training Now that the dynamic master discovery is working, re-enable IB with proper configuration for better throughput: - NCCL_IB_DISABLE=0 (enable IB) - Regex patterns for IB interfaces (vary by node: ibp6s0, ibp7s0) - RoCE v2 configuration (GID_INDEX=3) - GPUDirect RDMA enabled (NET_GDR_LEVEL=5) Co-Authored-By: Claude --- dags/test-train-pytorch.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py index 8afc74c..3f2f252 100644 --- a/dags/test-train-pytorch.py +++ b/dags/test-train-pytorch.py @@ -113,14 +113,18 @@ def run_training_node_func(rank, world_size): os.environ['WORLD_SIZE'] = str(world_size) os.environ['RANK'] = str(rank) - # NCCL Configuration - Start with basic ethernet config (IB is causing issues) + # NCCL Configuration for InfiniBand os.environ['NCCL_DEBUG'] = 'INFO' os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT) os.environ['NCCL_BLOCKING_WAIT'] = '1' - os.environ['NCCL_IB_DISABLE'] = '1' # Disable IB for now - use ethernet - os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' # Use regular ethernet + os.environ['NCCL_IB_DISABLE'] = '0' # Enable IB + os.environ['NCCL_IB_HCA'] = '^mlx5_[0-9]+$' # Regex for Mellanox HCAs + os.environ['NCCL_SOCKET_IFNAME'] = '^ibp[0-9]+s[0-9]+$' # Regex for IB interfaces os.environ['NCCL_P2P_DISABLE'] = '0' # Enable P2P os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1' # Better compatibility + os.environ['NCCL_NET_GDR_LEVEL'] = '5' # Enable GPUDirect RDMA + os.environ['NCCL_IB_GID_INDEX'] = '3' # Use RoCE v2 (common for modern IB) + os.environ['NCCL_IB_TC'] = '160' # Traffic class for low latency print(f"[{rank}] Environment configured:") print(f" MASTER_ADDR: {MASTER_ADDR}")