update dags

dags/test-train-pytorch.py

@@ -15,7 +15,7 @@ default_args = {
 
 # ── rank 0: стартует, пушит свой hostname в XCom ─────────────────────────────
 def run_rank_0(**context):
-    import os, resource, socket, torch
+    import os, socket, torch
     import torch.distributed as dist
     import torch.nn as nn
     import torch.optim as optim
@@ -27,9 +27,6 @@ def run_rank_0(**context):
     # Пушим MASTER_ADDR чтобы rank1 знал куда коннектиться
     context["ti"].xcom_push(key="master_addr", value=fqdn)
 
-    # Снимаем лимит на locked memory (нужно для ibv_reg_mr)
-    resource.setrlimit(resource.RLIMIT_MEMLOCK, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
-
     os.environ.update({
         "MASTER_ADDR": fqdn,
         "MASTER_PORT": MASTER_PORT,
@@ -38,6 +35,8 @@ def run_rank_0(**context):
         "NCCL_SOCKET_IFNAME": "eth0",   # bootstrap/rendezvous через eth
         "NCCL_IB_DISABLE": "0",         # трафик через IB
         "NCCL_NET": "IB",
+        "NCCL_IB_USE_INLINE": "1",      # без ibv_reg_mr pinned memory
+        "NCCL_BUFFSIZE": "1048576",
         "NCCL_P2P_DISABLE": "1",        # разные ноды — p2p не нужен
         "NCCL_SHM_DISABLE": "1",        # shm только для локальных рангов
         "NCCL_DEBUG": "INFO",
@@ -66,7 +65,7 @@ def run_rank_0(**context):
 
 # ── rank 1: ждёт XCom от rank0, потом коннектится ────────────────────────────
 def run_rank_1(**context):
-    import os, resource, socket, time, torch
+    import os, socket, time, torch
     import torch.distributed as dist
     import torch.nn as nn
     import torch.optim as optim
@@ -88,8 +87,6 @@ def run_rank_1(**context):
 
     print(f"[rank1] master_addr={master_addr}")
 
-    resource.setrlimit(resource.RLIMIT_MEMLOCK, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
-
     os.environ.update({
         "MASTER_ADDR": master_addr,
         "MASTER_PORT": MASTER_PORT,
@@ -98,6 +95,8 @@ def run_rank_1(**context):
         "NCCL_SOCKET_IFNAME": "eth0",
         "NCCL_IB_DISABLE": "0",
         "NCCL_NET": "IB",
+        "NCCL_IB_USE_INLINE": "1",      # без ibv_reg_mr pinned memory
+        "NCCL_BUFFSIZE": "1048576",
         "NCCL_P2P_DISABLE": "1",
         "NCCL_SHM_DISABLE": "1",
         "NCCL_DEBUG": "INFO",