fix abort in webui DDP mode
Former-commit-id: b90ac72d753b13a3eed9cb8b898fac2f2fe5153f
This commit is contained in:
@@ -71,28 +71,23 @@ def main():
|
||||
export_model()
|
||||
elif command == Command.TRAIN:
|
||||
if get_device_count() > 1:
|
||||
nnodes = os.environ.get("NNODES", "1")
|
||||
node_rank = os.environ.get("RANK", "0")
|
||||
nproc_per_node = os.environ.get("NPROC_PER_NODE", str(get_device_count()))
|
||||
master_addr = os.environ.get("MASTER_ADDR", "127.0.0.1")
|
||||
master_port = os.environ.get("MASTER_PORT", str(random.randint(20001, 29999)))
|
||||
logger.info("Initializing distributed tasks at: {}:{}".format(master_addr, master_port))
|
||||
subprocess.run(
|
||||
[
|
||||
"torchrun",
|
||||
"--nnodes",
|
||||
nnodes,
|
||||
"--node_rank",
|
||||
node_rank,
|
||||
"--nproc_per_node",
|
||||
nproc_per_node,
|
||||
"--master_addr",
|
||||
master_addr,
|
||||
"--master_port",
|
||||
master_port,
|
||||
launcher.__file__,
|
||||
*sys.argv[1:],
|
||||
]
|
||||
(
|
||||
"torchrun --nnodes {nnodes} --node_rank {node_rank} --nproc_per_node {nproc_per_node} "
|
||||
"--master_addr {master_addr} --master_port {master_port} {file_name} {args}"
|
||||
).format(
|
||||
nnodes=os.environ.get("NNODES", "1"),
|
||||
node_rank=os.environ.get("RANK", "0"),
|
||||
nproc_per_node=os.environ.get("NPROC_PER_NODE", str(get_device_count())),
|
||||
master_addr=master_addr,
|
||||
master_port=master_port,
|
||||
file_name=launcher.__file__,
|
||||
args=" ".join(sys.argv[1:]),
|
||||
),
|
||||
shell=True,
|
||||
)
|
||||
else:
|
||||
run_exp()
|
||||
|
||||
Reference in New Issue
Block a user