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

Enable InfiniBand for distributed PyTorch training

Browse Source 
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
George Stykalin
2026-04-16 16:06:42 +03:00
parent 6f88661b8b
commit 84ca598e3a
1 changed files with 3 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -40,13 +40,10 @@ def run_training_node_func(rank, world_size):
import time import time
# CRITICAL: Set NCCL environment variables BEFORE importing torch # CRITICAL: Set NCCL environment variables BEFORE importing torch
# 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_DEBUG'] = 'INFO'
os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT) os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1' # Updated naming os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1'
os.environ['NCCL_IB_DISABLE'] = '1' # Disable IB - use ethernet (PKEY misconfiguration) os.environ['NCCL_IB_DISABLE'] = '0' # Enable InfiniBand
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' # Use ethernet (working)
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
@@ -133,8 +130,7 @@ def run_training_node_func(rank, world_size):
print(f" MASTER_PORT: {MASTER_PORT}") print(f" MASTER_PORT: {MASTER_PORT}")
print(f" RANK: {rank}") print(f" RANK: {rank}")
print(f" WORLD_SIZE: {world_size}") print(f" WORLD_SIZE: {world_size}")
print(f" NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (ethernet mode)") print(f" NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (IB enabled)")
print(f" NCCL_SOCKET_IFNAME: {os.environ.get('NCCL_SOCKET_IFNAME')}")
# STEP 3.5: Pre-flight checks # STEP 3.5: Pre-flight checks
print(f"[{rank}] Running pre-flight checks...") print(f"[{rank}] Running pre-flight checks...")