From 062aa80b84a66d476ad90fc84df00c3a336d7239 Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <anastasiia.filippovaa@gmail.com>
Date: Thu, 7 Aug 2025 13:20:55 +0200
Subject: [PATCH] minor changer to mlx.launch

---
 python/mlx/distributed_run.py | 15 +--------------
 1 file changed, 1 insertion(+), 14 deletions(-)

diff --git a/python/mlx/distributed_run.py b/python/mlx/distributed_run.py
index 80c4b8977..fa5e094bd 100644
--- a/python/mlx/distributed_run.py
+++ b/python/mlx/distributed_run.py
@@ -418,7 +418,7 @@ def launch_mpi(parser, hosts, args, command):
 def launch_nccl(parser, hosts, args, command):
     master_host = hosts[0].ips[0]
     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.update(
@@ -814,19 +814,6 @@ def main():
         default=12345,
         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()
     if rest[0] == "--":