diff --git a/docs/src/usage/distributed.rst b/docs/src/usage/distributed.rst index 54d6617f..e0b39228 100644 --- a/docs/src/usage/distributed.rst +++ b/docs/src/usage/distributed.rst @@ -584,7 +584,7 @@ the process. **MLX_JACCL_COORDINATOR** should contain the IP and port that rank 0 can listen to all the other ranks connect to in order to establish the RDMA connections. -**MLX_IBV_DEVICES** should contain the path to a json file that contains the +**MLX_JACCL_DEVICES** should contain the path to a json file that contains the ibverbs device names that connect each node to each other node, something like the following: diff --git a/mlx/distributed/jaccl/jaccl.cpp b/mlx/distributed/jaccl/jaccl.cpp index 92718fdc..afd0b6da 100644 --- a/mlx/distributed/jaccl/jaccl.cpp +++ b/mlx/distributed/jaccl/jaccl.cpp @@ -1159,7 +1159,7 @@ bool is_available() { } std::shared_ptr init(bool strict /* = false */) { - const char* dev_file = std::getenv("MLX_IBV_DEVICES"); + const char* dev_file = std::getenv("MLX_JACCL_DEVICES"); const char* coordinator = std::getenv("MLX_JACCL_COORDINATOR"); const char* rank_str = std::getenv("MLX_RANK"); @@ -1167,9 +1167,9 @@ std::shared_ptr init(bool strict /* = false */) { if (strict) { std::ostringstream msg; msg << "[jaccl] You need to provide via environment variables a rank (MLX_RANK), " - << "a device file (MLX_IBV_DEVICES) and a coordinator ip/port (MLX_JACCL_COORDINATOR) " + << "a device file (MLX_JACCL_DEVICES) and a coordinator ip/port (MLX_JACCL_COORDINATOR) " << "but provided MLX_RANK=\"" << ((rank_str) ? rank_str : "") - << "\", MLX_IBV_DEVICES=\"" << ((dev_file) ? dev_file : "") + << "\", MLX_JACCL_DEVICES=\"" << ((dev_file) ? dev_file : "") << "\" and MLX_JACCL_COORDINATOR=\"" << ((coordinator) ? coordinator : ""); throw std::runtime_error(msg.str()); diff --git a/python/mlx/_distributed_utils/launch.py b/python/mlx/_distributed_utils/launch.py index 7f144799..3aa1a4f0 100644 --- a/python/mlx/_distributed_utils/launch.py +++ b/python/mlx/_distributed_utils/launch.py @@ -376,7 +376,7 @@ def launch_jaccl(parser, hosts, args, command): env = args.env cwd = args.cwd env.append(f"MLX_JACCL_COORDINATOR={coordinator}:{args.starting_port}") - files = {"MLX_IBV_DEVICES": json.dumps([h.rdma for h in hosts])} + files = {"MLX_JACCL_DEVICES": json.dumps([h.rdma for h in hosts])} log(args.verbose, "Running", shlex.join(command))