zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-29 03:22:54 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

minor changer to mlx.launch

Browse Source
This commit is contained in:
Anastasiia Filippova 2025-08-07 13:20:55 +02:00
parent f540b1d612
commit 062aa80b84
1 changed files with 1 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
python/mlx/distributed_run.py
View File

@ -418,7 +418,7 @@ def launch_mpi(parser, hosts, args, command):
def launch_nccl(parser, hosts, args, command): def launch_nccl(parser, hosts, args, command):
master_host = hosts[0].ips[0] master_host = hosts[0].ips[0]
master_port = args.nccl_port master_port = args.nccl_port
world_size = args.nproc_per_node * args.nnodes world_size = args.repeat_hosts * len(hosts)
base_env = os.environ.copy() base_env = os.environ.copy()
base_env.update( base_env.update(
@ -814,19 +814,6 @@ def main():
default=12345, default=12345,
help="The port to use for the NCCL communication (only for nccl backend)", help="The port to use for the NCCL communication (only for nccl backend)",
) )
parser.add_argument(
"--nproc-per-node",
default=8,
type=int,
help="The number of processes to launch on each node (only for nccl backend)",
)
# TODO: Add support for multiple nodes
parser.add_argument(
"--nnodes",
default=1,
type=int,
help="The number of nodes to launch (only for nccl backend)",
)
args, rest = parser.parse_known_args() args, rest = parser.parse_known_args()
if rest[0] == "--": if rest[0] == "--":