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

Restore NCCL INFO debug level, remove per-batch noise

Browse Source 
TRACE was too verbose. INFO shows clean transport selection lines
(NET/IB, GPU Direct RDMA status) for both ranks via stderr redirect.
Per-batch logging removed; per-epoch summary kept with timing.

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
George Stykalin
2026-04-16 16:27:45 +03:00
parent fdb3e590a6
commit dc2ffc270c
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -44,7 +44,7 @@ def run_training_node_func(rank, world_size):
os.dup2(sys.stdout.fileno(), sys.stderr.fileno()) os.dup2(sys.stdout.fileno(), sys.stderr.fileno())
# CRITICAL: Set NCCL environment variables BEFORE importing torch # CRITICAL: Set NCCL environment variables BEFORE importing torch
os.environ['NCCL_DEBUG'] = 'TRACE' 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' os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1'
os.environ['NCCL_IB_DISABLE'] = '0' # Enable InfiniBand os.environ['NCCL_IB_DISABLE'] = '0' # Enable InfiniBand
@@ -285,7 +285,6 @@ def run_training_node_func(rank, world_size):
optimizer.step() optimizer.step()
epoch_loss += loss.item() epoch_loss += loss.item()
print(f"[{rank}] Epoch {epoch+1}/{num_epochs} Batch {batch_idx+1}/{num_batches} | Loss: {loss.item():.6f}")
avg_loss = epoch_loss / num_batches avg_loss = epoch_loss / num_batches
epoch_time = time.time() - epoch_start epoch_time = time.time() - epoch_start