diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py index 9bde1c5..34e7d2c 100644 --- a/dags/test-train-pytorch.py +++ b/dags/test-train-pytorch.py @@ -3,7 +3,6 @@ from airflow.providers.standard.operators.python import PythonOperator import pendulum from datetime import timedelta -WORLD_SIZE = 2 MASTER_PORT = "29500" default_args = { @@ -13,153 +12,101 @@ default_args = { } -# ── rank 0: стартует, пушит свой hostname в XCom ───────────────────────────── +# ------------------------- +# RANK 0 +# ------------------------- def run_rank_0(**context): - import os, socket, torch - import torch.distributed as dist - import torch.nn as nn - import torch.optim as optim + import os, socket, torch, torch.distributed as dist hostname = socket.gethostname() - fqdn = f"{hostname}.airflow-worker-gpu" - print(f"[rank0] hostname={fqdn}") + master_addr = f"{hostname}.airflow-worker-gpu" - # Пушим MASTER_ADDR чтобы rank1 знал куда коннектиться - context["ti"].xcom_push(key="master_addr", value=fqdn) + # share master with rank1 + context["ti"].xcom_push(key="master_addr", value=master_addr) os.environ.update({ - "MASTER_ADDR": fqdn, + "MASTER_ADDR": master_addr, "MASTER_PORT": MASTER_PORT, "WORLD_SIZE": "2", - "NCCL_IB_GID_INDEX": "0", "RANK": "0", - "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 (minimal stable IB config) "NCCL_DEBUG": "INFO", - "NCCL_DEBUG_SUBSYS": "NET,INIT", - "TORCH_NCCL_BLOCKING_WAIT": "1", + "NCCL_SOCKET_IFNAME": "eth0", + + # IB stability (IMPORTANT) + "NCCL_IB_GID_INDEX": "0", + "NCCL_IB_USE_INLINE": "0", + "NCCL_IB_TIMEOUT": "22", + "NCCL_IB_RETRY_CNT": "7", }) - dist.init_process_group( - backend="nccl", - init_method="env://", - rank=0, - world_size=2, - timeout=timedelta(minutes=5), - ) - print("[rank0] process group OK") + dist.init_process_group("nccl", init_method="env://") torch.cuda.set_device(0) - device = torch.device("cuda:0") - print(f"[rank0] GPU={torch.cuda.get_device_name(0)}") - _train(rank=0, device=device) + x = torch.ones(1, device="cuda") + dist.all_reduce(x) + + print(f"[rank0] all_reduce result = {x.item()}") dist.destroy_process_group() - print("[rank0] DONE") -# ── rank 1: ждёт XCom от rank0, потом коннектится ──────────────────────────── +# ------------------------- +# RANK 1 +# ------------------------- def run_rank_1(**context): - import os, socket, time, torch - import torch.distributed as dist - import torch.nn as nn - import torch.optim as optim + import os, time, torch, torch.distributed as dist - # Polling XCom пока rank0 не запишет master_addr + # wait for master master_addr = None - for attempt in range(30): + for _ in range(30): master_addr = context["ti"].xcom_pull( task_ids="train_rank_0", key="master_addr", ) if master_addr: break - print(f"[rank1] waiting for master_addr (attempt {attempt+1}/30)...") - time.sleep(5) + time.sleep(2) if not master_addr: - raise RuntimeError("rank1: timeout waiting for master_addr from rank0") - - print(f"[rank1] master_addr={master_addr}") + raise RuntimeError("master_addr not found") os.environ.update({ "MASTER_ADDR": master_addr, "MASTER_PORT": MASTER_PORT, "WORLD_SIZE": "2", "RANK": "1", - "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 (same as rank0) "NCCL_DEBUG": "INFO", - "NCCL_DEBUG_SUBSYS": "NET,INIT", - "TORCH_NCCL_BLOCKING_WAIT": "1", + "NCCL_SOCKET_IFNAME": "eth0", + + # IB stability + "NCCL_IB_GID_INDEX": "0", + "NCCL_IB_USE_INLINE": "0", + "NCCL_IB_TIMEOUT": "22", + "NCCL_IB_RETRY_CNT": "7", }) - dist.init_process_group( - backend="nccl", - init_method="env://", - rank=1, - world_size=2, - timeout=timedelta(minutes=5), - ) - print("[rank1] process group OK") + dist.init_process_group("nccl", init_method="env://") torch.cuda.set_device(0) - device = torch.device("cuda:0") - print(f"[rank1] GPU={torch.cuda.get_device_name(0)}") - _train(rank=1, device=device) + x = torch.ones(1, device="cuda") + dist.all_reduce(x) + + print(f"[rank1] all_reduce result = {x.item()}") dist.destroy_process_group() - print("[rank1] DONE") -# ── общий train loop ────────────────────────────────────────────────────────── -def _train(rank: int, device): - import torch - import torch.distributed as dist - import torch.nn as nn - import torch.optim as optim - - # allreduce smoke-test - probe = torch.ones(1, device=device) * rank - dist.all_reduce(probe, op=dist.ReduceOp.SUM) - assert probe.item() == 1.0, f"allreduce mismatch: {probe.item()}" - print(f"[rank{rank}] allreduce smoke-test PASSED") - - model = nn.Sequential( - nn.Linear(10, 128), - nn.ReLU(), - nn.Linear(128, 10), - ).to(device) - - ddp_model = nn.parallel.DistributedDataParallel(model, device_ids=[0]) - optimizer = optim.SGD(ddp_model.parameters(), lr=0.001) - loss_fn = nn.MSELoss() - - for epoch in range(5): - x = torch.randn(32, 10, device=device) - y = torch.randn(32, 10, device=device) - optimizer.zero_grad() - loss = loss_fn(ddp_model(x), y) - loss.backward() - optimizer.step() - print(f"[rank{rank}] epoch={epoch} loss={loss.item():.4f}") - - -# ── DAG ─────────────────────────────────────────────────────────────────────── +# ------------------------- +# DAG +# ------------------------- with DAG( - dag_id="ddp_ib_xcom", + dag_id="ib_ddp_final_test", start_date=pendulum.today("UTC").add(days=-1), schedule=None, catchup=False, @@ -172,42 +119,12 @@ with DAG( task_id="train_rank_0", python_callable=run_rank_0, queue="gpu", - executor_config={ - "pod_override": { - "spec": { - "containers": [{ - "name": "base", - "securityContext": { - "capabilities": { - "add": ["IPC_LOCK", "SYS_RESOURCE"] - } - } - }] - } - } - }, ) rank1 = PythonOperator( task_id="train_rank_1", python_callable=run_rank_1, queue="gpu", - executor_config={ - "pod_override": { - "spec": { - "containers": [{ - "name": "base", - "securityContext": { - "capabilities": { - "add": ["IPC_LOCK", "SYS_RESOURCE"] - } - } - }] - } - } - }, ) - # rank1 стартует сразу, но ждёт master_addr через XCom polling - [rank0, rank1] - + rank0 >> rank1