diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py index 8a0c5f2..3eef6ab 100644 --- a/dags/test-train-pytorch.py +++ b/dags/test-train-pytorch.py @@ -85,17 +85,14 @@ def run_training_node_func(rank, world_size): os.environ['WORLD_SIZE'] = str(world_size) os.environ['RANK'] = str(rank) - # NCCL Configuration for InfiniBand + # NCCL Configuration - Start with basic ethernet config (IB is causing issues) os.environ['NCCL_DEBUG'] = 'INFO' os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT) os.environ['NCCL_BLOCKING_WAIT'] = '1' - os.environ['NCCL_IB_DISABLE'] = '0' # Enable IB - # Auto-detect IB interface - don't hardcode since names vary by node - os.environ['NCCL_SOCKET_IFNAME'] = '^ibp[0-9]+s[0-9]+$' # Regex pattern for IB interfaces - os.environ['NCCL_IB_HCA'] = '^mlx5_[0-9]+$' # Regex pattern for HCA devices + os.environ['NCCL_IB_DISABLE'] = '1' # Disable IB for now - use ethernet + os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' # Use regular ethernet 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 for IB print(f"[{rank}] Environment configured:") print(f" MASTER_ADDR: {MASTER_ADDR}") @@ -130,7 +127,15 @@ def run_training_node_func(rank, world_size): # STEP 4: Initialize process group print(f"[{rank}] Initializing process group (backend=nccl)...") - print(f"[{rank}] This may take 30-60 seconds if using IB...") + print(f"[{rank}] This may take 30-60 seconds...") + print(f"[{rank}] NCCL_DEBUG={os.environ.get('NCCL_DEBUG', 'not set')}") + print(f"[{rank}] NCCL_IB_DISABLE={os.environ.get('NCCL_IB_DISABLE', 'not set')}") + print(f"[{rank}] NCCL_SOCKET_IFNAME={os.environ.get('NCCL_SOCKET_IFNAME', 'not set')}") + + # Flush output to ensure logs are visible + import sys + sys.stdout.flush() + try: dist.init_process_group( backend="nccl", @@ -144,7 +149,10 @@ def run_training_node_func(rank, world_size): print(f" My rank: {dist.get_rank()}") except Exception as e: print(f"[{rank}] ✗ Failed to initialize process group!") + print(f" Error type: {type(e).__name__}") print(f" Error: {str(e)}") + import traceback + traceback.print_exc() raise # STEP 5: Setup GPU