| |
|
|
| import argparse |
| import os |
| import subprocess |
| import sys |
|
|
| parser = argparse.ArgumentParser(description="Nebula multi-node torchrun entry") |
| parser.add_argument("--training_script", required=True) |
| parser.add_argument("--nnodes", type=int, required=True) |
| parser.add_argument("--nproc_per_node", type=int, required=True) |
| launcher_args, training_args = parser.parse_known_args() |
|
|
| env = os.environ |
| print("===== distributed-related env =====", flush=True) |
| for key in sorted(env): |
| if any(name in key for name in ("MASTER", "RANK", "WORLD", "NODE", "WORKER", "NPROC")): |
| print(f"{key}={env[key]}", flush=True) |
| print("===================================", flush=True) |
|
|
| master_addr = env.get("MASTER_ADDR") |
| master_port = env.get("MASTER_PORT", "23456") |
| node_rank = env.get("RANK") or env.get("NODE_RANK") |
|
|
| if launcher_args.nnodes > 1 and not master_addr: |
| parser.error("MASTER_ADDR is required for multi-node training") |
| if node_rank is None: |
| parser.error("RANK or NODE_RANK is required for distributed training") |
|
|
| command = [ |
| "torchrun", |
| f"--nnodes={launcher_args.nnodes}", |
| f"--nproc_per_node={launcher_args.nproc_per_node}", |
| f"--node_rank={node_rank}", |
| f"--master_addr={master_addr or '127.0.0.1'}", |
| f"--master_port={master_port}", |
| launcher_args.training_script, |
| *training_args, |
| ] |
| print("command", " ".join(command), flush=True) |
| sys.exit(subprocess.run(command, check=False).returncode) |
|
|
