diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py index 557c64e..2e9e260 100644 --- a/dags/test-train-pytorch.py +++ b/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",