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

update dags

Browse Source
This commit is contained in:
George Stykalin
2026-04-15 18:21:41 +03:00
parent 971bff62f7
commit 6968b18418
1 changed files with 21 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -36,12 +36,28 @@ def run_training_node_func(rank, world_size):
"""Execute distributed training with synchronization"""
import os
import socket
from datetime import datetime
import time
# CRITICAL: Set NCCL environment variables BEFORE importing torch
# This ensures NCCL reads these settings during initialization
os.environ['NCCL_DEBUG'] = 'INFO'
os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
os.environ['NCCL_BLOCKING_WAIT'] = '1'
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' # Force ethernet (reliable)
os.environ['NCCL_IB_DISABLE'] = '1' # Disable IB - infrastructure not configured
os.environ['NCCL_P2P_DISABLE'] = '0' # Enable P2P
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
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from datetime import datetime
import time
print(f"{'='*60}")
print(f"Node Rank {rank}/{world_size} - Starting at {datetime.now()}")
@@ -115,23 +131,13 @@ def run_training_node_func(rank, world_size):
os.environ['WORLD_SIZE'] = str(world_size)
os.environ['RANK'] = str(rank)
# NCCL Configuration - Use ethernet (IB not accessible in containers)
os.environ['NCCL_DEBUG'] = 'INFO'
os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
os.environ['NCCL_BLOCKING_WAIT'] = '1'
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' # Use ethernet (reliable)
os.environ['NCCL_IB_DISABLE'] = '1' # Disable IB - /proc/drivers/infiniband not accessible
os.environ['NCCL_P2P_DISABLE'] = '0' # Enable P2P
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
print(f"[{rank}] Environment configured:")
print(f" MASTER_ADDR: {MASTER_ADDR}")
print(f" MASTER_ADDR: {master_addr}")
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')}")
print(f" NCCL_SOCKET_IFNAME: {os.environ.get('NCCL_SOCKET_IFNAME')}")
# STEP 3.5: Pre-flight checks
print(f"[{rank}] Running pre-flight checks...")
@@ -170,9 +176,6 @@ def run_training_node_func(rank, world_size):
# STEP 5: Initialize process group
print(f"[{rank}] Initializing process group (backend=nccl)...")
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