diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py index 853bd55..8393b4a 100644 --- a/dags/test-train-pytorch.py +++ b/dags/test-train-pytorch.py @@ -40,17 +40,13 @@ def run_training_node_func(rank, world_size): import time # CRITICAL: Set NCCL environment variables BEFORE importing torch - # This ensures NCCL reads these settings during initialization + # NOTE: IB is disabled due to subnet manager PKEY mismatch (test-2 only has limited membership 0x7fff) + # Cluster admin needs to configure opensm partitions to assign full membership PKEYs to both ports os.environ['NCCL_DEBUG'] = 'INFO' os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT) os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1' # Updated naming - os.environ['NCCL_IB_DISABLE'] = '0' # Enable IB - os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' # Fallback to eth0 - os.environ['NCCL_IB_GID_INDEX'] = '3' # Use RoCE v2 (GID index 3) - os.environ['NCCL_IB_TC'] = '106' # Traffic class for RoCE - os.environ['NCCL_IB_TIMEOUT'] = '22' # IB timeout (higher = more tolerant) - os.environ['NCCL_IB_RETRY_CNT'] = '7' # IB retry count - os.environ['NCCL_NET_GDR_LEVEL'] = '5' # Enable GPUDirect RDMA + os.environ['NCCL_IB_DISABLE'] = '1' # Disable IB - use ethernet (PKEY misconfiguration) + os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' # Use ethernet (working) os.environ['NCCL_P2P_DISABLE'] = '0' # Enable P2P os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1' # Better compatibility @@ -137,9 +133,8 @@ def run_training_node_func(rank, world_size): print(f" MASTER_PORT: {MASTER_PORT}") print(f" RANK: {rank}") print(f" WORLD_SIZE: {world_size}") - print(f" NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (IB enabled)") - print(f" NCCL_IB_GID_INDEX: {os.environ.get('NCCL_IB_GID_INDEX')}") - print(f" NCCL_IB_TIMEOUT: {os.environ.get('NCCL_IB_TIMEOUT')}") + print(f" NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (ethernet mode)") + print(f" NCCL_SOCKET_IFNAME: {os.environ.get('NCCL_SOCKET_IFNAME')}") # STEP 3.5: Pre-flight checks print(f"[{rank}] Running pre-flight checks...")