diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py
index 61b7043..788af8a 100644
--- a/dags/test-train-pytorch.py
+++ b/dags/test-train-pytorch.py
@@ -40,11 +40,11 @@ def run_training_node_func(rank, world_size):
   import time
 
   # CRITICAL: Set NCCL environment variables BEFORE importing torch
-  os.environ['NCCL_DEBUG'] = 'TRACE'
-  os.environ['NCCL_DEBUG_SUBSYS'] = 'NET,COLL'
+  os.environ['NCCL_DEBUG'] = 'INFO'
   os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
   os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1'
   os.environ['NCCL_IB_DISABLE'] = '0'  # Enable InfiniBand
+  os.environ['NCCL_IB_GPU_DIRECT_RDMA_MAX_DISTANCE'] = '6'  # GPU at bus 0, IB at bus 6
   os.environ['NCCL_P2P_DISABLE'] = '0'  # Enable P2P
   os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1'  # Better compatibility