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 
- Set NCCL_IB_DISABLE=0 to enable IB
- Add IB-specific config (GID_INDEX, TC, TIMEOUT, RETRY_CNT, NET_GDR_LEVEL)
- Update deprecated NCCL_BLOCKING_WAIT to TORCH_NCCL_BLOCKING_WAIT
- Set NCCL environment variables before torch imports for proper initialization

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
George Stykalin
2026-04-15 18:23:35 +03:00
parent 6968b18418
commit c659cd2f4b
1 changed files with 11 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -43,15 +43,16 @@ def run_training_node_func(rank, world_size):
# This ensures NCCL reads these settings during initialization # This ensures NCCL reads these settings during initialization
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['TORCH_NCCL_BLOCKING_WAIT'] = '1' # Updated naming
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' # Force ethernet (reliable) os.environ['NCCL_IB_DISABLE'] = '0' # Enable IB
os.environ['NCCL_IB_DISABLE'] = '1' # Disable IB - infrastructure not configured 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_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_LL_THRESHOLD'] = '0' # Disable LL for compatibility
os.environ['NCCL_ALGO'] = 'Ring' # Use ring algorithm
os.environ['NCCL_PROTO'] = 'Simple' # Use simple protocol
os.environ['NCCL_SOCKET_FAMILY'] = 'AF_INET' # Force IPv4 sockets
# Now import torch after NCCL config is set # Now import torch after NCCL config is set
import torch import torch
@@ -136,8 +137,9 @@ 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')}") print(f" NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (IB enabled)")
print(f" NCCL_SOCKET_IFNAME: {os.environ.get('NCCL_SOCKET_IFNAME')}") print(f" NCCL_IB_GID_INDEX: {os.environ.get('NCCL_IB_GID_INDEX')}")
print(f" NCCL_IB_TIMEOUT: {os.environ.get('NCCL_IB_TIMEOUT')}")
# 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...")