File size: 1,420 Bytes
c2b1b26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """Windows-friendly two-process DDP launcher for PyTorch builds without libuv."""
from __future__ import annotations
import argparse
import os
import subprocess
import sys
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--nproc_per_node", type=int, default=2)
parser.add_argument("--master_addr", default="127.0.0.1")
parser.add_argument("--master_port", default="29500")
parser.add_argument("training_script")
parser.add_argument("training_args", nargs=argparse.REMAINDER)
args = parser.parse_args()
processes = []
for local_rank in range(args.nproc_per_node):
env = os.environ.copy()
env.update(
{
"MASTER_ADDR": args.master_addr,
"MASTER_PORT": args.master_port,
"WORLD_SIZE": str(args.nproc_per_node),
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"USE_LIBUV": "0",
}
)
cmd = [sys.executable, args.training_script, *args.training_args]
processes.append(subprocess.Popen(cmd, env=env))
exit_code = 0
for process in processes:
exit_code = max(exit_code, process.wait())
if exit_code:
for process in processes:
if process.poll() is None:
process.terminate()
raise SystemExit(exit_code)
if __name__ == "__main__":
main()
|