mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-07-18 01:12:42 +00:00
[bug] Fix the version check bug in colossalai run when generating the cmd. (#4713)
* Fix the version check bug in colossalai run when generating the cmd. * polish code
This commit is contained in:
parent
3e05c07bb8
commit
1e0e080837
@ -156,7 +156,8 @@ def get_launch_command(
|
|||||||
torch_version = version.parse(torch.__version__)
|
torch_version = version.parse(torch.__version__)
|
||||||
assert torch_version.major >= 1
|
assert torch_version.major >= 1
|
||||||
|
|
||||||
if torch_version.minor < 9:
|
if torch_version.major == 1 and torch_version.minor < 9:
|
||||||
|
# torch distributed launch cmd with torch < 1.9
|
||||||
cmd = [
|
cmd = [
|
||||||
sys.executable,
|
sys.executable,
|
||||||
"-m",
|
"-m",
|
||||||
@ -177,7 +178,8 @@ def get_launch_command(
|
|||||||
value = extra_launch_args.pop(key)
|
value = extra_launch_args.pop(key)
|
||||||
default_torchrun_rdzv_args[key] = value
|
default_torchrun_rdzv_args[key] = value
|
||||||
|
|
||||||
if torch_version.minor < 10:
|
if torch_version.major == 1 and torch_version.minor == 9:
|
||||||
|
# torch distributed launch cmd with torch == 1.9
|
||||||
cmd = [
|
cmd = [
|
||||||
sys.executable,
|
sys.executable,
|
||||||
"-m",
|
"-m",
|
||||||
@ -187,6 +189,7 @@ def get_launch_command(
|
|||||||
f"--node_rank={node_rank}",
|
f"--node_rank={node_rank}",
|
||||||
]
|
]
|
||||||
else:
|
else:
|
||||||
|
# torch distributed launch cmd with torch > 1.9
|
||||||
cmd = [
|
cmd = [
|
||||||
"torchrun",
|
"torchrun",
|
||||||
f"--nproc_per_node={nproc_per_node}",
|
f"--nproc_per_node={nproc_per_node}",
|
||||||
|
Loading…
Reference in New Issue
Block a user