gstykalin/test-dags
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix PyTorch DDP initialization for K8s with IB

Browse Source 
- Fix MASTER_ADDR with full namespace suffix
- Use ethernet instead of IB (IB was causing hanging)
- Add better error handling and diagnostic logging
- Add NCCL environment variable logging

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
George Stykalin
2026-04-15 12:20:52 +03:00
parent 93d8b86105
commit ef78f2a7b8
1 changed files with 15 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
dags/test-train-pytorch.py
View File

@@ -85,17 +85,14 @@ def run_training_node_func(rank, world_size):
os.environ['WORLD_SIZE'] = str(world_size) os.environ['WORLD_SIZE'] = str(world_size)
os.environ['RANK'] = str(rank) 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_DEBUG'] = 'INFO'
os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT) os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
os.environ['NCCL_BLOCKING_WAIT'] = '1' os.environ['NCCL_BLOCKING_WAIT'] = '1'
os.environ['NCCL_IB_DISABLE'] = '0' # Enable IB os.environ['NCCL_IB_DISABLE'] = '1' # Disable IB for now - use ethernet
# Auto-detect IB interface - don't hardcode since names vary by node os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' # Use regular ethernet
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_P2P_DISABLE'] = '0' # Enable P2P os.environ['NCCL_P2P_DISABLE'] = '0' # Enable P2P
os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1' # Better compatibility 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"[{rank}] Environment configured:")
print(f" MASTER_ADDR: {MASTER_ADDR}") print(f" MASTER_ADDR: {MASTER_ADDR}")
@@ -130,7 +127,15 @@ def run_training_node_func(rank, world_size):
# STEP 4: Initialize process group # STEP 4: Initialize process group
print(f"[{rank}] Initializing process group (backend=nccl)...") 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: try:
dist.init_process_group( dist.init_process_group(
backend="nccl", backend="nccl",
@@ -144,7 +149,10 @@ def run_training_node_func(rank, world_size):
print(f" My rank: {dist.get_rank()}") print(f" My rank: {dist.get_rank()}")
except Exception as e: except Exception as e:
print(f"[{rank}] ✗ Failed to initialize process group!") print(f"[{rank}] ✗ Failed to initialize process group!")
print(f" Error type: {type(e).__name__}")
print(f" Error: {str(e)}") print(f" Error: {str(e)}")
import traceback
traceback.print_exc()
raise raise
# STEP 5: Setup GPU # STEP 5: Setup GPU