diff --git a/docs/src/_static/distributed/m3-ultra-mesh-broken.png b/docs/src/_static/distributed/m3-ultra-mesh-broken.png new file mode 100644 index 00000000..108ff58d Binary files /dev/null and b/docs/src/_static/distributed/m3-ultra-mesh-broken.png differ diff --git a/docs/src/_static/distributed/m3-ultra-mesh.png b/docs/src/_static/distributed/m3-ultra-mesh.png new file mode 100644 index 00000000..e049c024 Binary files /dev/null and b/docs/src/_static/distributed/m3-ultra-mesh.png differ diff --git a/docs/src/usage/distributed.rst b/docs/src/usage/distributed.rst index 0b83709e..54d6617f 100644 --- a/docs/src/usage/distributed.rst +++ b/docs/src/usage/distributed.rst @@ -7,22 +7,29 @@ Distributed Communication MLX supports distributed communication operations that allow the computational cost of training or inference to be shared across many physical machines. At the -moment we support three different communication backends: +moment we support several different communication backends introduced below. + +.. list-table:: + :widths: 20 80 + :header-rows: 1 + + * - Backend + - Description + * - :ref:`MPI ` + - A full featured and mature distributed communications library. + * - :ref:`RING ` + - Ring all reduce and all gather over TCP sockets. Always available and + usually faster than MPI. + * - :ref:`JACCL ` + - Low latency communication with RDMA over thunderbolt. Necessary for + things like tensor parallelism. + * - :ref:`NCCL ` + - The backend of choice for CUDA environments. -* `MPI `_ a - full-featured and mature distributed communications library -* A **ring** backend of our own that uses native TCP sockets. It should be - faster for thunderbolt connections, but it also works over Ethernet. -* `nccl `_, for use in CUDA environments. The list of all currently supported operations and their documentation can be seen in the :ref:`API docs`. -.. note:: - Some operations may not be supported or not as fast as they should be. - We are adding more and tuning the ones we have as we are figuring out the - best way to do distributed computing on Macs using MLX. - Getting Started --------------- @@ -85,7 +92,7 @@ Selecting Backend ^^^^^^^^^^^^^^^^^ You can select the backend you want to use when calling :func:`init` by passing -one of ``{'any', 'ring', 'mpi', 'nccl'}``. When passing ``any``, MLX will try all +one of ``{'any', 'ring', 'jaccl', 'mpi', 'nccl'}``. When passing ``any``, MLX will try all available backends. If they all fail then a singleton group is created. .. note:: @@ -110,6 +117,8 @@ The following examples aim to clarify the backend initialization logic in MLX: world_ring = mx.distributed.init(backend="ring") world_any = mx.distributed.init() # same as MPI because it was initialized first! +.. _training_example: + Training Example ---------------- @@ -192,16 +201,273 @@ almost identical to the example above: loss = step(model, x, y) mx.eval(loss, model.parameters()) +.. _ring_section: + +Getting Started with Ring +------------------------- + +The ring backend does not depend on any third party library so it is always +available. It uses TCP sockets so the nodes need to be reachable via a network. +As the name suggests the nodes are connected in a ring which means that rank 1 +can only communicate with rank 0 and rank 2, rank 2 only with rank 1 and rank 3 +and so on and so forth. As a result :func:`send` and :func:`recv` with +arbitrary sender and receiver are not supported in the ring backend. + +Defining a Ring +^^^^^^^^^^^^^^^ + +The easiest way to define and use a ring is via a JSON hostfile and the +``mlx.launch`` :doc:`helper script `. For each node one +defines a hostname to ssh into to run commands on this node and one or more IPs +that this node will listen to for connections. + +For example the hostfile below defines a 4 node ring. ``hostname1`` will be +rank 0, ``hostname2`` rank 1 etc. + +.. code:: json + + [ + {"ssh": "hostname1", "ips": ["123.123.123.1"]}, + {"ssh": "hostname2", "ips": ["123.123.123.2"]}, + {"ssh": "hostname3", "ips": ["123.123.123.3"]}, + {"ssh": "hostname4", "ips": ["123.123.123.4"]} + ] + +Running ``mlx.launch --hostfile ring-4.json my_script.py`` will ssh into each +node, run the script which will listen for connections in each of the provided +IPs. Specifically, ``hostname1`` will connect to ``123.123.123.2`` and accept a +connection from ``123.123.123.4`` and so on and so forth. + +Thunderbolt Ring +^^^^^^^^^^^^^^^^ + +Although the ring backend can have benefits over MPI even for Ethernet, its +main purpose is to use Thunderbolt rings for higher bandwidth communication. +Setting up such thunderbolt rings can be done manually, but is a relatively +tedious process. To simplify this, we provide the utility ``mlx.distributed_config``. + +To use ``mlx.distributed_config`` your computers need to be accessible by ssh via +Ethernet or Wi-Fi. Subsequently, connect them via thunderbolt cables and then call the +utility as follows: + +.. code:: shell + + mlx.distributed_config --verbose --hosts host1,host2,host3,host4 --backend ring + +By default the script will attempt to discover the thunderbolt ring and provide +you with the commands to configure each node as well as the ``hostfile.json`` +to use with ``mlx.launch``. If password-less ``sudo`` is available on the nodes +then ``--auto-setup`` can be used to configure them automatically. + +If you want to go through the process manually, the steps are as follows: + +* Disable the thunderbolt bridge interface +* For the cable connecting rank ``i`` to rank ``i + 1`` find the interfaces + corresponding to that cable in nodes ``i`` and ``i + 1``. +* Set up a unique subnetwork connecting the two nodes for the corresponding + interfaces. For instance if the cable corresponds to ``en2`` on node ``i`` + and ``en2`` also on node ``i + 1`` then we may assign IPs ``192.168.0.1`` and + ``192.168.0.2`` respectively to the two nodes. For more details you can see + the commands prepared by the utility script. + +.. _jaccl_section: + +Getting Started with JACCL +-------------------------- + +Starting from macOS 26.2, RDMA over thunderbolt is available and +enables low-latency communication between Macs with thunderbolt 5. MLX provides +the JACCL backend that uses this functionality to achieve communication latency +an order of magnitude lower than the ring backend. + +.. note:: + + The name JACCL (pronounced Jackal) stands for *Jack and Angelos' Collective + Communication Library* and it is an obvious pun to Nvidia's NCCL but also + tribute to *Jack Beasley* who led the development of RDMA over Thunderbolt + at Apple. + +Enabling RDMA +^^^^^^^^^^^^^ + +Until the feature matures, enabling RDMA over thunderbolt is slightly more +involved and **cannot** be done remotely even with sudo. In fact, it has to be +done in macOS recovery: + +1. `Start your computer in recovery `_. +2. Open the Terminal by going to Utilities -> Terminal. +3. Run ``rdma_ctl enable``. +4. Reboot. + +To verify that you have successfully enabled Thunderbolt RDMA you can run +``ibv_devices`` which should produce something like the following for an M3 Ultra. + +.. code-block:: bash + + ~ % ibv_devices + device node GUID + ------ ---------------- + rdma_en2 8096a9d9edbaac05 + rdma_en3 8196a9d9edbaac05 + rdma_en5 8396a9d9edbaac05 + rdma_en4 8296a9d9edbaac05 + rdma_en6 8496a9d9edbaac05 + rdma_en7 8596a9d9edbaac05 + +Defining a Mesh +^^^^^^^^^^^^^^^ + +The JACCL backend supports only fully connected topologies. Namely, there needs +to be a thunderbolt cable connecting all pairs of Macs directly. For example, in +the following topology visualizations, the left one is valid because there is a +connection from any node to any other node, while for the one on the right M3 +Ultra 1 is not connected to M3 Ultra 2. + +.. raw:: html + +
+
+ M3 Ultra thunderbolt mesh +

Fully connected mesh of four M3 Ultra.

+
+
+ M3 Ultra broken thunderbolt mesh +

Not a valid mesh (M3 Ultra 1 is not connected to M3 Ultra 2).

+
+
+ +Similar to the ring backend, the easiest way to use JACCL with MLX is to write +a JSON hostfile that will be used by ``mlx.launch``. The hostfile needs to contain + +- Hostnames to use for launching scripts via ssh +- An IP for rank 0 that is reachable by all nodes +- A list of rdma devices that connect each node to each other node + +The following JSON defines the valid 4-node mesh from the image above. + +.. code-block:: json + + [ + { + "ssh": "m3-ultra-1", + "ips": ["123.123.123.1"], + "rdma": [null, "rdma_en5", "rdma_en4", "rdma_en3"] + }, + { + "ssh": "m3-ultra-2", + "ips": [], + "rdma": ["rdma_en5", null, "rdma_en3", "rdma_en4"] + }, + { + "ssh": "m3-ultra-3", + "ips": [], + "rdma": ["rdma_en4", "rdma_en3", null, "rdma_en5"] + }, + { + "ssh": "m3-ultra-4", + "ips": [], + "rdma": ["rdma_en3", "rdma_en4", "rdma_en5", null] + } + ] + +Even though TCP/IP is not used when communicating with Thunderbolt RDMA, +disabling the thunderbolt bridge is still required as well as setting up +isolated local networks for each thunderbolt connection. + +All of the above can be done instead via ``mlx.distributed_config``. This helper +script will + +- ssh into each node +- extract the thunderbolt connectivity +- check for a valid mesh +- provide the commands to configure each node (or run them if sudo is available) +- generate the hostfile to be used with ``mlx.launch`` + +Putting It All Together +^^^^^^^^^^^^^^^^^^^^^^^^ + +For example launching a distributed MLX script that uses JACCL is fairly simple +if the nodes are reachable via ssh and have password-less sudo. + +First, connect all the thunderbolt cables. Then we can verify the connections +by using the ``mlx.distributed_config`` script to visualize them. + +.. code-block:: + + mlx.distributed_config --verbose \ + --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \ + --over thunderbolt --dot | dot -Tpng | open -f -a Preview + +After making sure that everything looks right we can auto-configure the nodes +and save the hostfile to ``m3-ultra-jaccl.json`` by running: + +.. code-block:: + + mlx.distributed_config --verbose \ + --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \ + --over thunderbolt --backend jaccl \ + --auto-setup --output m3-ultra-jaccl.json + +And now we are ready to run a distributed MLX script such as distributed inference +of a gigantic model using MLX LM. + +.. code-block:: + + mlx.launch --verbose --backend jaccl --hostfile m3-ultra-jaccl.json \ + --env MLX_METAL_FAST_SYNCH=1 -- \ # <--- important + /path/to/remote/python -m mlx_lm chat --model mlx-community/DeepSeek-R1-0528-4bit + +.. note:: + + Defining the environment variable ``MLX_METAL_FAST_SYNCH=1`` enables a + different, faster way of synchronizing between the GPU and the CPU. It is + not specific to the JACCL backend and can be used in all cases where the CPU + and GPU need to collaborate for some computation and is pretty critical for + low-latency communication since the communication is done by the CPU. + +.. _nccl_section: + +Getting Started with NCCL +------------------------- + +MLX on CUDA environments ships with the ability to talk to `NCCL +`_ which is a high-performance collective +communication library that supports both multi-gpu and multi-node setups. + +For CUDA environments, NCCL is the default backend for ``mlx.launch`` and all +it takes to run a distributed job is + +.. code-block:: + + mlx.launch -n 8 test.py + + # perfect for interactive scripts + mlx.launch -n 8 python -m mlx_lm chat --model my-model + +You can also use ``mlx.launch`` to ssh to a remote node and launch a script +with the same ease + +.. code-block:: + + mlx.launch --hosts my-cuda-node -n 8 test.py + +In many cases you may not want to use ``mlx.launch`` with the NCCL backend +because the cluster scheduler will be the one launching the processes. You can +:ref:`see which environment variables need to be defined ` in +order for the MLX NCCL backend to be initialized correctly. + +.. _mpi_section: Getting Started with MPI ------------------------ -MLX already comes with the ability to "talk" to MPI if it is installed on the -machine. Launching distributed MLX programs that use MPI can be done with -``mpirun`` as expected. However, in the following examples we will be using -``mlx.launch --backend mpi`` which takes care of some nuisances such as setting -absolute paths for the ``mpirun`` executable and the ``libmpi.dyld`` shared -library. +MLX already comes with the ability to "talk" to `MPI +`_ if it is installed +on the machine. Launching distributed MLX programs that use MPI can be done +with ``mpirun`` as expected. However, in the following examples we will be +using ``mlx.launch --backend mpi`` which takes care of some nuisances such as +setting absolute paths for the ``mpirun`` executable and the ``libmpi.dyld`` +shared library. The simplest possible usage is the following which, assuming the minimal example in the beginning of this page, should result in: @@ -269,78 +535,116 @@ Force MPI to use the most performant network interface by setting ``--mca btl_tcp_if_include `` where ```` should be the interface you want to use. -Getting Started with Ring -------------------------- +.. _no_mlx_launch: -The ring backend does not depend on any third party library so it is always -available. It uses TCP sockets so the nodes need to be reachable via a network. -As the name suggests the nodes are connected in a ring which means that rank 1 -can only communicate with rank 0 and rank 2, rank 2 only with rank 1 and rank 3 -and so on and so forth. As a result :func:`send` and :func:`recv` with -arbitrary sender and receiver is not supported in the ring backend. +Distributed Without ``mlx.launch`` +---------------------------------- -Defining a Ring -^^^^^^^^^^^^^^^ +None of the implementations of the distributed backends require launching with +``mlx.launch``. The script simply connects to each host. Starts a process per +rank and sets up the necessary environment variables before delegating to your +MLX script. See the :doc:`dedicated documentation page ` +for more details. -The easiest way to define and use a ring is via a JSON hostfile and the -``mlx.launch`` :doc:`helper script `. For each node one -defines a hostname to ssh into to run commands on this node and one or more IPs -that this node will listen to for connections. +For many use-cases this will be the easiest way to perform distributed +computations in MLX. However, there may be reasons that you cannot or should +not use ``mlx.launch``. A common such case is the use of a scheduler that +starts all the processes for you on machines undetermined at the time of +scheduling the job. -For example the hostfile below defines a 4 node ring. ``hostname1`` will be -rank 0, ``hostname2`` rank 1 etc. +Below we list the environment variables required to use each backend. -.. code:: json +Ring +^^^^^^ - [ - {"ssh": "hostname1", "ips": ["123.123.123.1"]}, - {"ssh": "hostname2", "ips": ["123.123.123.2"]}, - {"ssh": "hostname3", "ips": ["123.123.123.3"]}, - {"ssh": "hostname4", "ips": ["123.123.123.4"]} - ] +**MLX_RANK** should contain a single 0-based integer that defines the rank of +the process. -Running ``mlx.launch --hostfile ring-4.json my_script.py`` will ssh into each -node, run the script which will listen for connections in each of the provided -IPs. Specifically, ``hostname1`` will connect to ``123.123.123.2`` and accept a -connection from ``123.123.123.4`` and so on and so forth. +**MLX_HOSTFILE** should contain the path to a json file that contains IPs and +ports for each rank to listen to, something like the following: -Thunderbolt Ring -^^^^^^^^^^^^^^^^ +.. code-block:: json -Although the ring backend can have benefits over MPI even for Ethernet, its -main purpose is to use Thunderbolt rings for higher bandwidth communication. -Setting up such thunderbolt rings can be done manually, but is a relatively -tedious process. To simplify this, we provide the utility ``mlx.distributed_config``. + [ + ["123.123.1.1:5000", "123.123.1.2:5000"], + ["123.123.2.1:5000", "123.123.2.2:5000"], + ["123.123.3.1:5000", "123.123.3.2:5000"], + ["123.123.4.1:5000", "123.123.4.2:5000"] + ] -To use ``mlx.distributed_config`` your computers need to be accessible by ssh via -Ethernet or Wi-Fi. Subsequently, connect them via thunderbolt cables and then call the -utility as follows: +**MLX_RING_VERBOSE** is optional and if set to 1 it enables some more logging +from the distributed backend. -.. code:: shell +JACCL +^^^^^ - mlx.distributed_config --verbose --hosts host1,host2,host3,host4 +**MLX_RANK** should contain a single 0-based integer that defines the rank of +the process. -By default the script will attempt to discover the thunderbolt ring and provide -you with the commands to configure each node as well as the ``hostfile.json`` -to use with ``mlx.launch``. If password-less ``sudo`` is available on the nodes -then ``--auto-setup`` can be used to configure them automatically. +**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. -To validate your connection without configuring anything -``mlx.distributed_config`` can also plot the ring using DOT format. +**MLX_IBV_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: -.. code:: shell +.. code-block:: json - mlx.distributed_config --verbose --hosts host1,host2,host3,host4 --dot >ring.dot - dot -Tpng ring.dot >ring.png - open ring.png + [ + [null, "rdma_en5", "rdma_en4", "rdma_en3"], + ["rdma_en5", null, "rdma_en3", "rdma_en4"], + ["rdma_en4", "rdma_en3", null, "rdma_en5"], + ["rdma_en3", "rdma_en4", "rdma_en5", null] + ] -If you want to go through the process manually, the steps are as follows: -* Disable the thunderbolt bridge interface -* For the cable connecting rank ``i`` to rank ``i + 1`` find the interfaces - corresponding to that cable in nodes ``i`` and ``i + 1``. -* Set up a unique subnetwork connecting the two nodes for the corresponding - interfaces. For instance if the cable corresponds to ``en2`` on node ``i`` - and ``en2`` also on node ``i + 1`` then we may assign IPs ``192.168.0.1`` and - ``192.168.0.2`` respectively to the two nodes. For more details you can see - the commands prepared by the utility script. +NCCL +^^^^^ + +**MLX_RANK** should contain a single 0-based integer that defines the rank of +the process. + +**MLX_WORLD_SIZE** should contain the total number of processes that will be +launched. + +**NCCL_HOST_IP** and **NCCL_PORT** should contain the IP and port that all +hosts can connect to to establish the NCCL communication. + +**CUDA_VISIBLE_DEVICES** should contain the local index of the gpu that +corresponds to this process. + +Of course any `other environment variable +`_ that is +used by NCCL can be set. + +.. _tips_and_tricks: + +Tips and Tricks +---------------- + +This is a small collection of tips to help you utilize better the distributed +communication capabilities of MLX. + +- *Test locally first.* + + You can use the pattern ``mlx.launch -n2 -- my_script.py`` to run a small + scale test on a single node first. + +- *Batch your communication.* + + As described in the :ref:`training example `, performing a + lot of small communications can hurt performance. Copy the approach of + :func:`mlx.nn.average_gradients` to gather many small communications in a + single large one. + +- *Visualize the connectivity.* + + Use ``mlx.distributed_config --hosts h1,h2,h3 --over thunderbolt --dot`` to + visualize the connnections and make sure that the cables are connected + correctly. See the :ref:`JACCL section ` for examples. + +- *Use the debugger.* + + ``mlx.launch`` is meant for interactive use. It broadcasts stdin to all + processes and gathers stdout from all processes. This makes using ``pdb`` a + breeze. diff --git a/docs/src/usage/launching_distributed.rst b/docs/src/usage/launching_distributed.rst index 2e956a48..e1e3128b 100644 --- a/docs/src/usage/launching_distributed.rst +++ b/docs/src/usage/launching_distributed.rst @@ -7,13 +7,106 @@ Launching Distributed Programs .. currentmodule:: mlx.core.distributed -Installing the MLX python package provides a helper script ``mlx.launch`` that -can be used to run python scripts distributed on several nodes. It allows -launching using either the MPI backend or the ring backend. See the -:doc:`distributed docs ` for the different backends. +The MLX python package provides two utilities to help you configure +your Macs for distributed computation and also launch distributed programs on +multiple nodes or with many processes in a single node. These utilities are aptly named -Usage ------ +- ``mlx.launch`` +- ``mlx.distributed_config`` + +See the :doc:`distributed docs ` for an introduction and +getting-started guides to the various backends. + +``mlx.distributed_config`` +--------------------------- + +Unless you are launching distributed jobs locally for development or multi-gpu +CUDA environments, then you have several Macs that you need to configure for +distributed communication with MLX. + +``mlx.distributed_config`` aims to automate the process of configuring the +network interfaces (especially for communication over thunderbolt) and also +creating the hostfile to be used with ``mlx.launch``. + +We will analyse 3 cases of using ``mlx.distributed_config`` + +1. RDMA over thunderbolt using JACCL +2. TCP/IP over thunderbolt using the ring backend +3. TCP/IP over ethernet using the ring backend + +JACCL +^^^^^^^ + +After following :ref:`the steps to enable RDMA ` you can run the +following command to configure the nodes and create the hostfile. + +.. code-block:: + + mlx.distributed_config --verbose --backend jaccl \ + --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 --over thunderbolt \ + --auto-setup --output m3-ultra-jaccl.json + +Let's walk through the steps that the script takes to configure the nodes. + +1. ssh to all nodes to verify that they are reachable +2. Extract the thunderbolt connectivity. Namely run commands on each node to + calculate which node is connected to which other node. +3. Verify that we have a valid fully connected mesh +4. Check that RDMA is enabled +5. Extract the ethernet IP from interface en0 +6. Disable the thunderbolt bridge and set up peer to peer networks for each + thunderbolt cable +7. Write the hostfile + +Knowing the above steps allows you to manually configure the nodes but also +debug any configuration issue. For instance changing the Ethernet IP to a +different interface directly in the config is possible (as long as it is +reachable from all nodes). + +The ``--auto-setup`` argument requires password-less sudo on each node. If it +isn't available then the configuration script will print commands to be run on +each node. + +Ring over thunderbolt +^^^^^^^^^^^^^^^^^^^^^ + +Setting up a ring backend over thunderbolt only requires changing the +``--backend`` from ``jaccl`` to ``ring``. + +The steps are very similar with the main difference being that instead of +verifying that the nodes are fully connected, the script attempts to identify a +ring topology (or multiple rings). + +Ring over Ethernet +^^^^^^^^^^^^^^^^^^ + +Configuring the ring backend over ethernet doesn't require setting up network +interface and as such it simply extracts the ``en0`` IP from each node and +writes the hostfile. + +Debugging cable connections +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``mlx.distributed_config`` can help you debug the connectivity of your nodes +over thunderbolt by exporting a graph of the connections. + +Running + +.. code-block:: + + mlx.distributed_config --verbose \ + --hosts host1,host2,host3,host4 \ + --over thunderbolt --dot + +will export a `GraphViz `_ representation of the +connections between the nodes which makes it very easy to figure out which +cable is not connected correctly. + +See :ref:`the JACCL section ` for an example. + + +``mlx.launch`` +-------------- The minimal usage example of ``mlx.launch`` is simply @@ -33,6 +126,10 @@ the rest if one of them fails unexpectedly or if ``mlx.launch`` is terminated. It also takes care of forwarding the output of each remote process to stdout and stderr respectively. +Importantly, it also broadcasts stdin to each process which enables interactive +programs to work in distributed mode as well as debugging using the interactive +debugger. + Providing Hosts ^^^^^^^^^^^^^^^^ @@ -63,10 +160,62 @@ host and on the same path. A good checklist to debug errors is the following: ``mlx.launch --print-python`` to see what that path is. * the script you want to run is available on all hosts at the same path +If you are launching from a node with a completely different setup than the +nodes that the program will run on, you can specify ``--no-verify-script`` so +that ``mlx.launch`` does not attempt to verify that the executable and script +exist locally before launching the distributed job. + +.. _ring_specifics: + +Ring Specifics +^^^^^^^^^^^^^^ + +The :ref:`ring ` backend, which is also the default +backend, can be explicitly selected with the argument ``--backend ring``. The +ring backend has some specific requirements and arguments that are different to +other backends: + +* The argument ``--hosts`` only accepts IPs and not hostnames. If we need to + ssh to a hostname that does not correspond to the IP we want to bind to we + have to provide a hostfile. +* ``--starting-port`` defines the port to bind to on the remote hosts. + Specifically rank 0 for the first IP will use this port and each subsequent + IP or rank will add 1 to this port. +* ``--connections-per-ip`` allows us to increase the number of connections + between neighboring nodes. This corresponds to ``--mca btl_tcp_links 2`` for + ``mpirun``. + +.. _jaccl_specifics: + +JACCL Specifics +^^^^^^^^^^^^^^^^ + +The :ref:`JACCL ` backend can be selected with the argument +``--backend jaccl``. A hostfile is necessary to launch with this backend +because it needs to contain the RDMA devices connecting each node to each other +node. + +NCCL Specifics +^^^^^^^^^^^^^^ + +The :ref:`NCCL ` backend is the default backend for CUDA +environments. When launching from a Mac to a Linux machine with CUDA then the +backend should be selected using ``--backend nccl``. + +The ``--repeat-hosts, -n`` argument should be used to launch multi-node and +multi-gpu jobs. For instance + +.. code-block:: + + mlx.launch --backend nccl --hosts linux-1,linux-2 -n 8 --no-verify-script -- ./my-job.sh + +will attempt to launch 16 processes, 8 on each node that will all run +``my-job.sh``. + .. _mpi_specifics: MPI Specifics -------------- +^^^^^^^^^^^^^ One can use MPI by passing ``--backend mpi`` to ``mlx.launch``. In that case, ``mlx.launch`` is a thin wrapper over ``mpirun``. Moreover, @@ -83,23 +232,3 @@ to choose a specific interface for the byte-transfer-layer of MPI we can call .. code:: shell mlx.launch --backend mpi --mpi-arg '--mca btl_tcp_if_include en0' --hostfile hosts.json my_script.py - - -.. _ring_specifics: - -Ring Specifics --------------- - -The ring backend, which is also the default backend, can be explicitly selected -with the argument ``--backend ring``. The ring backend has some specific -requirements and arguments that are different to MPI: - -* The argument ``--hosts`` only accepts IPs and not hostnames. If we need to - ssh to a hostname that does not correspond to the IP we want to bind to we - have to provide a hostfile. -* ``--starting-port`` defines the port to bind to on the remote hosts. - Specifically rank 0 for the first IP will use this port and each subsequent - IP or rank will add 1 to this port. -* ``--connections-per-ip`` allows us to increase the number of connections - between neighboring nodes. This corresponds to ``--mca btl_tcp_links 2`` for - ``mpirun``. diff --git a/mlx/distributed/CMakeLists.txt b/mlx/distributed/CMakeLists.txt index b7762f6a..807fd0fb 100644 --- a/mlx/distributed/CMakeLists.txt +++ b/mlx/distributed/CMakeLists.txt @@ -4,6 +4,11 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/ops.cpp ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp) +if(MLX_BUILD_CPU AND NOT WIN32) + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp) +endif() + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/mpi) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/ring) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/nccl) +add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/jaccl) diff --git a/mlx/distributed/distributed.cpp b/mlx/distributed/distributed.cpp index 2f5ea802..3cde6a26 100644 --- a/mlx/distributed/distributed.cpp +++ b/mlx/distributed/distributed.cpp @@ -5,6 +5,7 @@ #include "mlx/backend/cuda/cuda.h" #include "mlx/distributed/distributed.h" #include "mlx/distributed/distributed_impl.h" +#include "mlx/distributed/jaccl/jaccl.h" #include "mlx/distributed/mpi/mpi.h" #include "mlx/distributed/nccl/nccl.h" #include "mlx/distributed/ring/ring.h" @@ -102,7 +103,27 @@ class EmptyGroup : public GroupImpl { } // namespace detail bool is_available() { - return mpi::is_available() || ring::is_available() || nccl::is_available(); + return mpi::is_available() || ring::is_available() || nccl::is_available() || + jaccl::is_available(); +} + +bool is_available(const std::string& bk) { + if (bk == "any") { + return is_available(); + } + if (bk == "mpi") { + return mpi::is_available(); + } + if (bk == "ring") { + return ring::is_available(); + } + if (bk == "nccl") { + return nccl::is_available(); + } + if (bk == "jaccl") { + return jaccl::is_available(); + } + return false; } int Group::rank() const { @@ -135,6 +156,8 @@ Group init(bool strict /* = false */, const std::string& bk /* = "any" */) { group = ring::init(strict); } else if (bk == "nccl") { group = nccl::init(strict); + } else if (bk == "jaccl") { + group = jaccl::init(strict); } else if (bk == "any") { if (mlx::core::cu::is_available()) { group = nccl::init(false); @@ -148,13 +171,17 @@ Group init(bool strict /* = false */, const std::string& bk /* = "any" */) { group = mpi::init(false); bk_ = "mpi"; } + if (group == nullptr) { + group = jaccl::init(false); + bk_ = "jaccl"; + } if (group == nullptr && strict) { throw std::runtime_error("[distributed] Couldn't initialize any backend"); } } else { std::ostringstream msg; - msg << "[distributed] The only valid values for backend are 'any', 'mpi' " - << "and 'ring' but '" << bk << "' was provided."; + msg << "[distributed] The only valid values for backend are 'any', 'mpi', 'nccl', " + << "'jaccl' and 'ring' but '" << bk << "' was provided."; throw std::invalid_argument(msg.str()); } diff --git a/mlx/distributed/distributed.h b/mlx/distributed/distributed.h index fa5c42a1..a6971dd9 100644 --- a/mlx/distributed/distributed.h +++ b/mlx/distributed/distributed.h @@ -16,6 +16,7 @@ class GroupImpl; /* Check if a communication backend is available */ bool is_available(); +bool is_available(const std::string& bk); /** * A distributed::Group represents a group of independent mlx processes that diff --git a/mlx/distributed/jaccl/CMakeLists.txt b/mlx/distributed/jaccl/CMakeLists.txt new file mode 100644 index 00000000..dbd60c8e --- /dev/null +++ b/mlx/distributed/jaccl/CMakeLists.txt @@ -0,0 +1,7 @@ +if(MLX_BUILD_CPU + AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin" + AND MACOS_SDK_VERSION GREATER_EQUAL 26.2) + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp) +else() + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_jaccl.cpp) +endif() diff --git a/mlx/distributed/jaccl/jaccl.cpp b/mlx/distributed/jaccl/jaccl.cpp new file mode 100644 index 00000000..92718fdc --- /dev/null +++ b/mlx/distributed/jaccl/jaccl.cpp @@ -0,0 +1,1197 @@ +// Copyright © 2025 Apple Inc. + +#include +#include +#include +#include +#include + +#include + +#include "mlx/backend/cpu/encoder.h" +#include "mlx/distributed/distributed_impl.h" +#include "mlx/distributed/reduction_ops.h" +#include "mlx/distributed/utils.h" +#include "mlx/dtype_utils.h" + +#define LOAD_SYMBOL(symbol, variable) \ + { \ + variable = (decltype(variable))dlsym(librdma_handle_, #symbol); \ + char* error = dlerror(); \ + if (error != nullptr) { \ + std::cerr << IBV_TAG << " " << error << std::endl; \ + librdma_handle_ = nullptr; \ + return; \ + } \ + } + +constexpr const char* IBV_TAG = "[jaccl]"; +constexpr int NUM_BUFFERS = 2; +constexpr int BUFFER_SIZE = 4096; +constexpr int MAX_SEND_WR = 32; +constexpr int MAX_RECV_WR = 32; +constexpr int SEND_WR = 1; +constexpr int RECV_WR = 2; +constexpr int MAX_PEERS = 8; + +using GroupImpl = mlx::core::distributed::detail::GroupImpl; +using json = nlohmann::json; +namespace detail = mlx::core::distributed::detail; +namespace allocator = mlx::core::allocator; + +struct IBVWrapper { + IBVWrapper() { + librdma_handle_ = dlopen("librdma.dylib", RTLD_NOW | RTLD_GLOBAL); + if (librdma_handle_ == nullptr) { + return; + } + + LOAD_SYMBOL(ibv_get_device_list, get_device_list); + LOAD_SYMBOL(ibv_get_device_name, get_device_name); + LOAD_SYMBOL(ibv_open_device, open_device); + LOAD_SYMBOL(ibv_free_device_list, free_device_list); + LOAD_SYMBOL(ibv_close_device, close_device); + + LOAD_SYMBOL(ibv_alloc_pd, alloc_pd); + LOAD_SYMBOL(ibv_create_qp, create_qp); + LOAD_SYMBOL(ibv_create_cq, create_cq); + LOAD_SYMBOL(ibv_destroy_cq, destroy_cq); + LOAD_SYMBOL(ibv_destroy_qp, destroy_qp); + LOAD_SYMBOL(ibv_dealloc_pd, dealloc_pd); + + LOAD_SYMBOL(ibv_query_port, query_port); + LOAD_SYMBOL(ibv_query_gid, query_gid); + LOAD_SYMBOL(ibv_modify_qp, modify_qp); + LOAD_SYMBOL(ibv_reg_mr, reg_mr); + LOAD_SYMBOL(ibv_dereg_mr, dereg_mr); + + // Not really symbols but leaving them here in case they become symbols in + // the future. + // + // LOAD_SYMBOL(ibv_post_send, post_send); + // LOAD_SYMBOL(ibv_post_recv, post_recv); + // LOAD_SYMBOL(ibv_poll_cq, poll_cq); + } + + bool is_available() { + return librdma_handle_ != nullptr; + } + + void* librdma_handle_; + + // API + ibv_device** (*get_device_list)(int*); + const char* (*get_device_name)(ibv_device*); + ibv_context* (*open_device)(ibv_device*); + void (*free_device_list)(ibv_device**); + int (*close_device)(ibv_context*); + + ibv_pd* (*alloc_pd)(ibv_context*); + ibv_qp* (*create_qp)(ibv_pd*, ibv_qp_init_attr*); + ibv_cq* (*create_cq)(ibv_context*, int, void*, ibv_comp_channel*, int); + int (*destroy_cq)(ibv_cq*); + int (*destroy_qp)(ibv_qp*); + int (*dealloc_pd)(ibv_pd*); + + int (*query_port)(ibv_context*, uint8_t, ibv_port_attr*); + int (*query_gid)(ibv_context*, uint8_t, int, ibv_gid*); + int (*modify_qp)(ibv_qp*, ibv_qp_attr*, int); + ibv_mr* (*reg_mr)(ibv_pd*, void*, size_t, int); + int (*dereg_mr)(ibv_mr*); +}; + +IBVWrapper& ibv() { + static IBVWrapper wrapper; + return wrapper; +} + +template +struct is_container : std::false_type {}; + +template +struct is_container< + T, + std::void_t> + : std::true_type {}; + +std::ostream& operator<<(std::ostream& os, const ibv_gid& gid) { + os << std::hex << std::setfill('0'); + for (int i = 0; i < 16; i += 2) { + uint16_t part = (gid.raw[i] << 8) | gid.raw[i + 1]; + os << std::setw(4) << part; + if (i < 14) + os << ":"; + } + os << std::dec; + return os; +} + +void* page_aligned_alloc(size_t num_bytes) { + static size_t page_size = sysconf(_SC_PAGESIZE); + void* buf; + if (posix_memalign(&buf, page_size, num_bytes)) { + return nullptr; + } + return buf; +} + +/** + * Contains the information that defines a destination to a remote device. + * Basically we can compute our own destination and share it with remote hosts + * over the side channel. + */ +struct Destination { + int local_id; + int queue_pair_number; + int packet_sequence_number; + ibv_gid global_identifier; +}; + +std::ostream& operator<<(std::ostream& os, const Destination& dst) { + os << dst.local_id << " " << dst.queue_pair_number << " " + << dst.packet_sequence_number << " " << dst.global_identifier; + return os; +} + +/** + * A buffer that can be registered to a number of protection domains. + */ +class SharedBuffer { + public: + SharedBuffer(size_t num_bytes) + : data_(page_aligned_alloc(num_bytes)), num_bytes_(num_bytes) {} + ~SharedBuffer() { + for (auto& [pd, mr] : memory_regions_) { + ibv().dereg_mr(mr); + } + if (data_ != nullptr) { + std::free(data_); + } + } + + SharedBuffer(const SharedBuffer&) = delete; + SharedBuffer& operator=(const SharedBuffer&) = delete; + SharedBuffer(SharedBuffer&& b) : data_(nullptr), num_bytes_(0) { + std::swap(data_, b.data_); + std::swap(num_bytes_, b.num_bytes_); + std::swap(memory_regions_, b.memory_regions_); + } + + void register_to_protection_domain(ibv_pd* protection_domain) { + auto [it, inserted] = memory_regions_.insert({protection_domain, nullptr}); + if (!inserted) { + throw std::runtime_error( + "[jaccl] Buffer can be registered once per protection domain"); + } + + it->second = ibv().reg_mr( + protection_domain, + data_, + num_bytes_, + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | + IBV_ACCESS_REMOTE_WRITE); + if (!it->second) { + throw std::runtime_error("[jaccl] Register memory region failed"); + } + } + + size_t size() const { + return num_bytes_; + } + + uint32_t local_key(ibv_pd* protection_domain) const { + return memory_regions_.at(protection_domain)->lkey; + } + + ibv_sge to_scatter_gather_entry(ibv_pd* protection_domain) const { + ibv_sge entry; + entry.addr = reinterpret_cast(data_); + entry.length = size(); + entry.lkey = local_key(protection_domain); + return entry; + } + + template + T* data() { + return static_cast(data_); + } + + template + T* begin() { + return static_cast(data_); + } + + template + T* end() { + return static_cast(data_) + size() / sizeof(T); + } + + private: + void* data_; + size_t num_bytes_; + std::unordered_map memory_regions_; +}; + +/** + * Manipulates an RDMA connection. Enables (among other things) + * + * - Creating a queue pair + * - Sending and receiving + * - Checking completion + */ +struct Connection { + ibv_context* ctx; + ibv_pd* protection_domain; + ibv_cq* completion_queue; + ibv_qp* queue_pair; + Destination src; // holds the local information + + Connection(ibv_context* ctx_) + : ctx(ctx_), + protection_domain(nullptr), + completion_queue(nullptr), + queue_pair(nullptr) { + src.local_id = -1; + } + + Connection(Connection&& c) : Connection(nullptr) { + std::swap(ctx, c.ctx); + std::swap(protection_domain, c.protection_domain); + std::swap(completion_queue, c.completion_queue); + std::swap(queue_pair, c.queue_pair); + std::swap(src, c.src); + } + + Connection(const Connection&) = delete; + Connection& operator=(Connection&) = delete; + + ~Connection() { + if (queue_pair != nullptr) { + ibv().destroy_qp(queue_pair); + } + if (completion_queue != nullptr) { + ibv().destroy_cq(completion_queue); + } + if (protection_domain != nullptr) { + ibv().dealloc_pd(protection_domain); + } + if (ctx != nullptr) { + ibv().close_device(ctx); + } + } + + void allocate_protection_domain() { + protection_domain = ibv().alloc_pd(ctx); + if (protection_domain == nullptr) { + throw std::runtime_error("[jaccl] Couldn't allocate protection domain"); + } + } + + void create_completion_queue(int num_entries) { + completion_queue = ibv().create_cq(ctx, num_entries, nullptr, nullptr, 0); + if (completion_queue == nullptr) { + throw std::runtime_error("[jaccl] Couldn't create completion queue"); + } + } + + void create_queue_pair() { + ibv_qp_init_attr init_attr; + init_attr.qp_context = ctx; + init_attr.qp_context = ctx; + init_attr.send_cq = completion_queue; + init_attr.recv_cq = completion_queue; + init_attr.srq = nullptr; + init_attr.cap.max_send_wr = MAX_SEND_WR; + init_attr.cap.max_recv_wr = MAX_RECV_WR; + init_attr.cap.max_send_sge = 1; + init_attr.cap.max_recv_sge = 1; + init_attr.cap.max_inline_data = 0; + init_attr.qp_type = IBV_QPT_UC; + init_attr.sq_sig_all = 0; + + queue_pair = ibv().create_qp(protection_domain, &init_attr); + + if (queue_pair == nullptr) { + throw std::runtime_error("[jaccl] Couldn't create queue pair"); + } + } + + const Destination& info() { + if (queue_pair == nullptr || src.local_id >= 0) { + return src; + } + + ibv_port_attr port_attr; + ibv().query_port(ctx, 1, &port_attr); + ibv_gid gid; + ibv().query_gid(ctx, 1, 1, &gid); + + src.local_id = port_attr.lid; + src.queue_pair_number = queue_pair->qp_num; + src.packet_sequence_number = 7; // TODO: Change to sth random + src.global_identifier = gid; + + return src; + } + + void queue_pair_init() { + ibv_qp_attr attr = {}; + attr.qp_state = IBV_QPS_INIT; + attr.port_num = 1; + attr.pkey_index = 0; + attr.qp_access_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | + IBV_ACCESS_REMOTE_WRITE; + + int mask = + IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS; + + if (int status = ibv().modify_qp(queue_pair, &attr, mask); status != 0) { + std::ostringstream msg; + msg << "[jaccl] Changing queue pair to INIT failed with errno " << status; + throw std::invalid_argument(msg.str()); + } + } + + void queue_pair_rtr(const Destination& dst) { + ibv_qp_attr attr = {}; + memset(&attr, 0, sizeof(attr)); + attr.qp_state = IBV_QPS_RTR; + attr.path_mtu = IBV_MTU_1024; + attr.rq_psn = dst.packet_sequence_number; + attr.dest_qp_num = dst.queue_pair_number; + attr.ah_attr.dlid = dst.local_id; + attr.ah_attr.sl = 0; + attr.ah_attr.src_path_bits = 0; + attr.ah_attr.port_num = 1; + attr.ah_attr.is_global = 0; + + if (dst.global_identifier.global.interface_id) { + attr.ah_attr.is_global = 1; + attr.ah_attr.grh.hop_limit = 1; + attr.ah_attr.grh.dgid = dst.global_identifier; + attr.ah_attr.grh.sgid_index = 1; + } + + int mask = IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | + IBV_QP_RQ_PSN; + + if (int status = ibv().modify_qp(queue_pair, &attr, mask); status != 0) { + std::ostringstream msg; + msg << "[jaccl] Changing queue pair to RTR failed with errno " << status; + throw std::invalid_argument(msg.str()); + } + } + + void queue_pair_rts() { + ibv_qp_attr attr = {}; + attr.qp_state = IBV_QPS_RTS; + attr.sq_psn = src.packet_sequence_number; + + int mask = IBV_QP_STATE | IBV_QP_SQ_PSN; + + if (int status = ibv().modify_qp(queue_pair, &attr, mask); status != 0) { + std::ostringstream msg; + msg << "[jaccl] Changing queue pair to RTS failed with errno " << status; + throw std::invalid_argument(msg.str()); + } + } + + void post_send(const SharedBuffer& buff, uint64_t work_request_id) { + ibv_send_wr work_request, *bad_work_request; + + auto entry = buff.to_scatter_gather_entry(protection_domain); + work_request.wr_id = work_request_id; + work_request.sg_list = &entry; + work_request.num_sge = 1; + work_request.opcode = IBV_WR_SEND; + work_request.send_flags = IBV_SEND_SIGNALED; + work_request.next = nullptr; + + if (int status = + ibv_post_send(queue_pair, &work_request, &bad_work_request); + status != 0) { + std::ostringstream msg; + msg << "[jaccl] Send failed with error code " << status; + throw std::invalid_argument(msg.str()); + } + } + + void post_recv(const SharedBuffer& buff, uint64_t work_request_id) { + ibv_recv_wr work_request, *bad_work_request; + + auto entry = buff.to_scatter_gather_entry(protection_domain); + work_request.wr_id = work_request_id; + work_request.sg_list = &entry; + work_request.num_sge = 1; + work_request.next = nullptr; + + if (int status = + ibv_post_recv(queue_pair, &work_request, &bad_work_request); + status != 0) { + std::ostringstream msg; + msg << "[jaccl] Recv failed with error code " << status; + throw std::invalid_argument(msg.str()); + } + } +}; + +/** + * Implement a TCP side channel to exchange information about the RDMA + * connections. + * + * Implements a simple all gather where every node sends to rank 0 and rank 0 + * broadcasts to every node. + */ +class SideChannel { + public: + SideChannel(int rank, int size, const char* addr) : rank_(rank), size_(size) { + auto address = detail::parse_address(addr); + + if (rank_ == 0) { + detail::TCPSocket server(IBV_TAG); + server.listen(IBV_TAG, address); + + for (int i = 0; i < size - 1; i++) { + sockets_.push_back(server.accept(IBV_TAG)); + } + + std::vector ranks(size - 1); + for (int i = 0; i < size - 1; i++) { + sockets_[i].recv( + IBV_TAG, reinterpret_cast(&ranks[i]), sizeof(int)); + ranks[i]--; + } + for (int i = 0; i < size - 1; i++) { + while (i != ranks[i]) { + std::swap(sockets_[i], sockets_[ranks[i]]); + std::swap(ranks[i], ranks[ranks[i]]); + } + } + } else { + sockets_.push_back(detail::TCPSocket::connect( + IBV_TAG, address, 4, 1000, [](int attempt, int wait) { + std::cerr << IBV_TAG << " Connection attempt " << attempt + << " waiting " << wait << " ms" << std::endl; + })); + sockets_[0].send(IBV_TAG, reinterpret_cast(&rank_), sizeof(int)); + } + } + + SideChannel(const SideChannel&) = delete; + SideChannel& operator=(const SideChannel&) = delete; + + SideChannel(SideChannel&& sc) + : rank_(sc.rank_), size_(sc.size_), sockets_(std::move(sc.sockets_)) { + sc.rank_ = -1; + sc.size_ = -1; + } + + template + std::vector all_gather(const T& v) { + std::vector result(size_); + + // T is a container of stuff like std::vector or std::string + if constexpr (is_container::value) { + using U = typename T::value_type; + + // Share the lengths first and set the communication size to be the + // maximum length of the containers. + auto lengths = all_gather(v.size()); + auto max_len = *std::max_element(lengths.begin(), lengths.end()); + for (auto& s : result) { + s.resize(max_len); + } + + // All gather of length max_len + if (rank_ == 0) { + std::copy(v.begin(), v.end(), result[rank_].begin()); + for (int i = 1; i < size_; i++) { + sockets_[i - 1].recv(IBV_TAG, result[i].data(), sizeof(U) * max_len); + } + for (int i = 1; i < size_; i++) { + for (int j = 0; j < size_; j++) { + sockets_[i - 1].send( + IBV_TAG, result[j].data(), sizeof(U) * max_len); + } + } + } else { + std::copy(v.begin(), v.end(), result[rank_].begin()); + sockets_[0].send(IBV_TAG, result[rank_].data(), sizeof(U) * max_len); + for (int i = 0; i < size_; i++) { + sockets_[0].recv(IBV_TAG, result[i].data(), sizeof(U) * max_len); + } + } + + // Resize the outputs back to the original length + for (int i = 0; i < size_; i++) { + result[i].resize(lengths[i]); + } + } + + // T is a scalar + else { + if (rank_ == 0) { + result[rank_] = v; + for (int i = 1; i < size_; i++) { + sockets_[i - 1].recv(IBV_TAG, &result[i], sizeof(T)); + } + for (int i = 1; i < size_; i++) { + sockets_[i - 1].send(IBV_TAG, result.data(), size_ * sizeof(T)); + } + } else { + sockets_[0].send(IBV_TAG, &v, sizeof(T)); + sockets_[0].recv(IBV_TAG, result.data(), size_ * sizeof(T)); + } + } + + return result; + } + + private: + int rank_; + int size_; + std::vector sockets_; +}; + +/** + * Manages a set of connections. Among other things it uses a side channel to + * exchange the necessary information and then configure the connections to be + * ready for RDMA operations. + */ +class ConnectionManager { + public: + ConnectionManager( + int rank, + const std::vector& device_names, + const char* coordinator_addr) + : rank_(rank), + size_(device_names.size()), + side_channel_(rank_, size_, coordinator_addr) { + create_contexts(device_names); + if (connections_[rank_].ctx != nullptr) { + throw std::runtime_error("[jaccl] Malformed device file"); + } + } + + int rank() const { + return rank_; + } + + int size() const { + return size_; + } + + /** + * Performs the connection initialization. Namely, after this call all + * Connection objects should have a queue pair in RTS state. + */ + void initialize(int num_buffers, size_t num_bytes) { + // Create the queue pairs + for (auto& conn : connections_) { + if (conn.ctx == nullptr) { + continue; + } + conn.allocate_protection_domain(); + conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); + conn.create_queue_pair(); + } + + allocate_buffers(num_buffers, num_bytes); + + // First init all connections + for (int peer = 0; peer < size_; peer++) { + if (peer == rank_) { + continue; + } + connections_[peer].queue_pair_init(); + } + + // Gather the information to be exchanged, this also serves as a barrier so + // that all peers have initialized their connections before attempting to + // transition to RTS. + std::vector info; + for (auto& conn : connections_) { + info.emplace_back(conn.info()); + } + auto all_infos = side_channel_.all_gather(info); + + // Transition queue pairs to RTS + for (int peer = 0; peer < size_; peer++) { + if (peer == rank_) { + continue; + } + auto peer_info = all_infos[peer][rank_]; + connections_[peer].queue_pair_rtr(peer_info); + connections_[peer].queue_pair_rts(); + } + } + + void allocate_buffers(int num_buffers, size_t num_bytes) { + // Deregister any buffers and free the memory + buffers_.clear(); + + // Allocate the memory + for (int i = 0; i < num_buffers; i++) { + for (int j = 0; j < size_; j++) { + buffers_.emplace_back(num_bytes); + } + } + + for (int i = 0; i < num_buffers; i++) { + for (int j = 0; j < size_; j++) { + // This is our send buffer so register it with all pds so we can send + // it to all connected devices. + if (j == rank_) { + for (auto& conn : connections_) { + if (conn.ctx != nullptr) { + buffers_[i * size_ + j].register_to_protection_domain( + conn.protection_domain); + } + } + } + + // This is the recv buffer from rank j so register it to rank j's + // protection domain. + else { + buffers_[i * size_ + j].register_to_protection_domain( + connections_[j].protection_domain); + } + } + } + } + + void send_to(int rank, int buff) { + connections_[rank].post_send( + buffers_[buff * size_ + rank_], SEND_WR << 16 | buff << 8 | rank); + } + + void recv_from(int rank, int buff) { + connections_[rank].post_recv( + buffers_[buff * size_ + rank], RECV_WR << 16 | buff << 8 | rank); + } + + /** + * Poll all connections and save the work completions and return the + * corresponding length. + */ + int poll(int num_completions, ibv_wc* work_completions) { + int completions = 0; + for (int r = 0; r < size_; r++) { + if (r == rank_) { + continue; + } + if (completions >= num_completions) { + return completions; + } + + int c = ibv_poll_cq( + connections_[r].completion_queue, + num_completions - completions, + work_completions + completions); + + completions += c; + } + return completions; + } + + /** + * + */ + int poll(int rank, int num_completions, ibv_wc* work_completions) { + return ibv_poll_cq( + connections_[rank].completion_queue, num_completions, work_completions); + } + + SharedBuffer& send_buffer(int buff) { + return buffers_[buff * size_ + rank_]; + } + + SharedBuffer& buffer(int rank, int buff) { + return buffers_[buff * size_ + rank]; + } + + void barrier() { + side_channel_.all_gather(0); + } + + private: + void create_contexts(const std::vector& device_names) { + int num_devices = 0; + ibv_device** devices = ibv().get_device_list(&num_devices); + for (auto& name : device_names) { + // Empty so add a nullptr context + if (name.empty()) { + connections_.emplace_back(nullptr); + continue; + } + + // Search for the name and try to open the device + for (int i = 0; i < num_devices; i++) { + if (name == ibv().get_device_name(devices[i])) { + auto ctx = ibv().open_device(devices[i]); + if (ctx == nullptr) { + std::ostringstream msg; + msg << "[jaccl] Could not open device " << name; + throw std::runtime_error(msg.str()); + } + connections_.emplace_back(ctx); + break; + } + } + } + ibv().free_device_list(devices); + } + + int rank_; + int size_; + SideChannel side_channel_; + std::vector connections_; + std::vector buffers_; +}; + +std::vector load_device_names(int rank, const char* dev_file) { + std::vector device_names; + std::ifstream f(dev_file); + + json devices = json::parse(f); + devices = devices[rank]; + for (auto it = devices.begin(); it != devices.end(); it++) { + std::string n; + if (!it->is_null()) { + n = *it; + } + device_names.emplace_back(std::move(n)); + } + + return device_names; +} + +namespace mlx::core::distributed::jaccl { + +class IBVGroup : public GroupImpl { + public: + IBVGroup(ConnectionManager cm) + : cm_(std::move(cm)), rank_(cm.rank()), size_(cm.size()) {} + + Stream communication_stream(StreamOrDevice s) override { + return to_stream(s, Device::cpu); + } + + int rank() override { + return cm_.rank(); + } + + int size() override { + return cm_.size(); + } + + void all_sum(const array& input, array& output, Stream stream) override { + dispatch_all_types(output.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + all_reduce(input, output, stream, detail::SumOp{}); + }); + } + + void all_max(const array& input, array& output, Stream stream) override { + dispatch_all_types(output.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + all_reduce(input, output, stream, detail::MaxOp{}); + }); + } + + void all_min(const array& input, array& output, Stream stream) override { + dispatch_all_types(output.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + all_reduce(input, output, stream, detail::MinOp{}); + }); + } + + void all_gather(const array& input, array& output, Stream stream) override { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + size_t n_bytes = input.nbytes(); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.set_output_array(output); + encoder.dispatch([in_ptr, out_ptr, n_bytes, this]() { + // Copy our data to the appropriate place + std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); + + // Fully connected all gather + char* data = out_ptr; + char* our_data = out_ptr + rank_ * n_bytes; + constexpr int64_t N = BUFFER_SIZE; + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2; + int64_t total = static_cast(n_bytes); + int num_peers = size_ - 1; + + // Counters to maintain the state of transfers + int in_flight = 0; + int read_offset = 0; + int completed_send_count[PIPELINE] = {0}; + int write_offset[MAX_PEERS] = {0}; + + // Prefill the pipeline + int buff = 0; + while (read_offset < total && buff < PIPELINE) { + post_recv_all(buff); + std::copy( + our_data + read_offset, + our_data + std::min(read_offset + N, total), + cm_.send_buffer(buff).begin()); + post_send_all(buff); + + buff++; + in_flight += 2 * num_peers; + read_offset += N; + } + + // Main loop + // + // Keep going until we have no longer data in flight. + while (in_flight > 0) { + ibv_wc wc[WC_NUM]; + int n = cm_.poll(WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + // Send completed. If all sends completed then send the next chunk. + if (work_type == SEND_WR && read_offset < total) { + completed_send_count[buff]++; + if (completed_send_count[buff] == num_peers) { + std::copy( + our_data + read_offset, + our_data + std::min(read_offset + N, total), + cm_.send_buffer(buff).begin()); + post_send_all(buff); + + completed_send_count[buff] = 0; + in_flight += num_peers; + read_offset += N; + } + } + + // Recv completed. If we have more chunks then post another recv. + else if (work_type == RECV_WR) { + std::copy( + cm_.buffer(rank, buff).begin(), + cm_.buffer(rank, buff).begin() + + std::min(N, total - write_offset[rank]), + data + rank * n_bytes + write_offset[rank]); + write_offset[rank] += N; + if (write_offset[rank] + N * (PIPELINE - 1) < total) { + cm_.recv_from(rank, buff); + in_flight++; + } + } + } + } + }); + } + + void send(const array& input, int dst, Stream stream) override { + auto data = input.data(); + int64_t n_bytes = input.nbytes(); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.dispatch([data, n_bytes, dst, this]() { + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE; + constexpr int N = BUFFER_SIZE; + + int in_flight = 0; + int64_t read_offset = 0; + + // Prefill the pipeline + int buff = 0; + while (read_offset < n_bytes && buff < PIPELINE) { + std::copy( + data + read_offset, + data + std::min(read_offset + N, n_bytes), + cm_.send_buffer(buff).begin()); + cm_.send_to(dst, buff); + + buff++; + read_offset += N; + in_flight++; + } + + // Main loop + while (in_flight > 0) { + // Poll the hardware for completions. + // + // If a send was completed and we have more data to send then go ahead + // and send them. + ibv_wc wc[WC_NUM]; + int n = cm_.poll(WC_NUM, wc); + for (int i = 0; i < n; i++) { + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + if (read_offset < n_bytes) { + std::copy( + data + read_offset, + data + std::min(read_offset + N, n_bytes), + cm_.send_buffer(buff).begin()); + cm_.send_to(dst, buff); + + read_offset += N; + in_flight++; + } + } + } + }); + } + + void recv(array& out, int src, Stream stream) override { + auto data = out.data(); + int64_t n_bytes = out.nbytes(); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_output_array(out); + encoder.dispatch([data, n_bytes, src, this]() { + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE; + constexpr int N = BUFFER_SIZE; + + int in_flight = 0; + int64_t write_offset = 0; + + // Prefill the pipeline + int buff = 0; + while (write_offset < n_bytes && buff < PIPELINE) { + cm_.recv_from(src, buff); + + in_flight++; + buff++; + } + + // Main loop + while (in_flight > 0) { + // Poll the hardware for completions. + // + // If a recv was completed copy it to the output and if we have more + // data to fetch post another recv. + ibv_wc wc[WC_NUM]; + int n = cm_.poll(WC_NUM, wc); + for (int i = 0; i < n; i++) { + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + std::copy( + cm_.buffer(src, buff).begin(), + cm_.buffer(src, buff).begin() + + std::min(n_bytes - write_offset, static_cast(N)), + data + write_offset); + write_offset += N; + + if (write_offset + (PIPELINE - 1) * N < n_bytes) { + cm_.recv_from(src, buff); + + in_flight++; + } + } + } + }); + } + + void sum_scatter(const array& input, array& output, Stream stream) override { + throw std::runtime_error("[jaccl] sum_scatter not supported."); + } + + std::shared_ptr split(int color, int key = -1) override { + throw std::runtime_error("[jaccl] Group split not supported."); + } + + private: + void post_recv_all(int buffer) { + for (int i = 0; i < size_; i++) { + if (i == rank_) { + continue; + } + cm_.recv_from(i, buffer); + } + } + + void post_send_all(int buffer) { + for (int i = 0; i < size_; i++) { + if (i == rank_) { + continue; + } + cm_.send_to(i, buffer); + } + } + + template + void all_reduce( + const array& input, + array& output, + Stream stream, + ReduceOp reduce_op) { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.set_output_array(output); + encoder.dispatch([in_ptr, out_ptr, size = input.size(), this, reduce_op]() { + // If not inplace all reduce then copy the input to the output first + if (in_ptr != out_ptr) { + std::memcpy(out_ptr, in_ptr, size * sizeof(T)); + } + + // Fully connected all reduce + T* data = out_ptr; + constexpr int64_t N = BUFFER_SIZE / sizeof(T); + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2; + int64_t total = static_cast(size); + int num_peers = size_ - 1; + + // Counters to maintain the state of transfers + int in_flight = 0; + int read_offset = 0; + int completed_send_count[PIPELINE] = {0}; + int completed_recv_begin[MAX_PEERS] = {0}; + int completed_recv_end[MAX_PEERS] = {0}; + + // Prefill the pipeline + int buff = 0; + while (read_offset < total && buff < PIPELINE) { + post_recv_all(buff); + std::copy( + data + read_offset, + data + std::min(read_offset + N, total), + cm_.send_buffer(buff).begin()); + post_send_all(buff); + + buff++; + in_flight += 2 * num_peers; + read_offset += N; + } + + // Main loop + // + // Keep going until we have no longer data in flight. + while (in_flight > 0) { + // Poll the hardware for completions. + // + // If a send was completed mark how many completions we have received + // for that buffer. If we have sent the buffer to all peers we can + // reuse the buffer so copy the next chunk of data and send it to all. + // + // If a receive is completed then advance the pointer of completed + // receives. + ibv_wc wc[WC_NUM]; + int n = cm_.poll(WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + if (work_type == SEND_WR && read_offset < total) { + completed_send_count[buff]++; + if (completed_send_count[buff] == num_peers) { + std::copy( + data + read_offset, + data + std::min(read_offset + N, total), + cm_.send_buffer(buff).begin()); + post_send_all(buff); + + completed_send_count[buff] = 0; + in_flight += num_peers; + read_offset += N; + } + } + + else if (work_type == RECV_WR) { + completed_recv_end[rank]++; + } + } + + // Process the completed recv + // + // For each rank we have a range of completed recv defined by a begin + // and end inclusive and exlusive in standard C++ fashion. + // + // When there is an unprocessed receive we first check if we have + // finished sending the write location. If so then we reduce in-place + // and then check if there is more to be received and post a recv. + for (int r = 0; r < size_; r++) { + int s = completed_recv_begin[r]; + int e = completed_recv_end[r]; + int w = s * N; + while (w < read_offset && e - s > 0) { + int buff = s % PIPELINE; + reduce_op( + cm_.buffer(r, buff).begin(), + data + w, + std::min(N, total - w)); + w += N; + s++; + if (w + (PIPELINE - 1) * N < total) { + cm_.recv_from(r, buff); + in_flight++; + } + } + completed_recv_begin[r] = s; + } + } + }); + } + + ConnectionManager cm_; + int rank_; + int size_; +}; + +bool is_available() { + return ibv().is_available(); +} + +std::shared_ptr init(bool strict /* = false */) { + const char* dev_file = std::getenv("MLX_IBV_DEVICES"); + const char* coordinator = std::getenv("MLX_JACCL_COORDINATOR"); + const char* rank_str = std::getenv("MLX_RANK"); + + if (!is_available() || !dev_file || !coordinator || !rank_str) { + 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) " + << "but provided MLX_RANK=\"" << ((rank_str) ? rank_str : "") + << "\", MLX_IBV_DEVICES=\"" << ((dev_file) ? dev_file : "") + << "\" and MLX_JACCL_COORDINATOR=\"" + << ((coordinator) ? coordinator : ""); + throw std::runtime_error(msg.str()); + } + return nullptr; + } + + auto rank = std::atoi(rank_str); + auto device_names = load_device_names(rank, dev_file); + + auto cm = ConnectionManager(rank, device_names, coordinator); + if (cm.size() > MAX_PEERS) { + std::ostringstream msg; + msg << "[jaccl] The maximum number of supported peers is " << MAX_PEERS + << " but " << cm.size() << " was provided"; + throw std::runtime_error(msg.str()); + } + + cm.initialize(NUM_BUFFERS, BUFFER_SIZE); + cm.barrier(); + + return std::make_shared(std::move(cm)); +} + +} // namespace mlx::core::distributed::jaccl diff --git a/mlx/distributed/jaccl/jaccl.h b/mlx/distributed/jaccl/jaccl.h new file mode 100644 index 00000000..d07f9ccc --- /dev/null +++ b/mlx/distributed/jaccl/jaccl.h @@ -0,0 +1,12 @@ +// Copyright © 2025 Apple Inc. + +#include "mlx/distributed/distributed.h" + +namespace mlx::core::distributed::jaccl { + +using GroupImpl = mlx::core::distributed::detail::GroupImpl; + +bool is_available(); +std::shared_ptr init(bool strict = false); + +} // namespace mlx::core::distributed::jaccl diff --git a/mlx/distributed/jaccl/no_jaccl.cpp b/mlx/distributed/jaccl/no_jaccl.cpp new file mode 100644 index 00000000..12fd6ab2 --- /dev/null +++ b/mlx/distributed/jaccl/no_jaccl.cpp @@ -0,0 +1,20 @@ +// Copyright © 2025 Apple Inc. + +#include "mlx/distributed/jaccl/jaccl.h" + +namespace mlx::core::distributed::jaccl { + +using GroupImpl = mlx::core::distributed::detail::GroupImpl; + +bool is_available() { + return false; +} + +std::shared_ptr init(bool strict /* = false */) { + if (strict) { + throw std::runtime_error("Cannot initialize jaccl distributed backend."); + } + return nullptr; +} + +} // namespace mlx::core::distributed::jaccl diff --git a/mlx/distributed/reduction_ops.h b/mlx/distributed/reduction_ops.h new file mode 100644 index 00000000..02777be3 --- /dev/null +++ b/mlx/distributed/reduction_ops.h @@ -0,0 +1,38 @@ +// Copyright © 2025 Apple Inc. + +namespace mlx::core::distributed::detail { + +template +struct SumOp { + void operator()(const T* input, T* output, size_t N) const { + while (N-- > 0) { + *output += *input; + input++; + output++; + } + } +}; + +template +struct MaxOp { + void operator()(const T* input, T* output, size_t N) const { + while (N-- > 0) { + *output = std::max(*output, *input); + input++; + output++; + } + } +}; + +template +struct MinOp { + void operator()(const T* input, T* output, size_t N) const { + while (N-- > 0) { + *output = std::min(*output, *input); + input++; + output++; + } + } +}; + +} // namespace mlx::core::distributed::detail diff --git a/mlx/distributed/ring/ring.cpp b/mlx/distributed/ring/ring.cpp index 23537c4d..9679d9ff 100644 --- a/mlx/distributed/ring/ring.cpp +++ b/mlx/distributed/ring/ring.cpp @@ -1,9 +1,6 @@ // Copyright © 2024 Apple Inc. -#include #include -#include -#include #include #include #include @@ -22,6 +19,8 @@ #include "mlx/backend/cpu/encoder.h" #include "mlx/distributed/distributed.h" #include "mlx/distributed/distributed_impl.h" +#include "mlx/distributed/reduction_ops.h" +#include "mlx/distributed/utils.h" #include "mlx/threadpool.h" #ifndef SOL_TCP @@ -94,6 +93,7 @@ constexpr const size_t ALL_SUM_SIZE = 8 * 1024 * 1024; constexpr const size_t ALL_SUM_BUFFERS = 2; constexpr const int CONN_ATTEMPTS = 5; constexpr const int CONN_WAIT = 1000; +constexpr const char* RING_TAG = "[ring]"; using GroupImpl = mlx::core::distributed::detail::GroupImpl; using json = nlohmann::json; @@ -296,55 +296,6 @@ class CommunicationThreads { std::unordered_map threads_; }; -struct address_t { - sockaddr_storage addr; - socklen_t len; - - const sockaddr* get() const { - return (struct sockaddr*)&addr; - } -}; - -/** - * Parse a sockaddr from an ip and port provided as strings. - */ -address_t parse_address(const std::string& ip, const std::string& port) { - struct addrinfo hints, *res; - memset(&hints, 0, sizeof(hints)); - hints.ai_family = AF_UNSPEC; - hints.ai_socktype = SOCK_STREAM; - - int status = getaddrinfo(ip.c_str(), port.c_str(), &hints, &res); - if (status != 0) { - std::ostringstream msg; - msg << "Can't parse address " << ip << ":" << port; - throw std::runtime_error(msg.str()); - } - - address_t result; - memcpy(&result.addr, res->ai_addr, res->ai_addrlen); - result.len = res->ai_addrlen; - freeaddrinfo(res); - - return result; -} - -/** - * Parse a sockaddr provided as an : string. - */ -address_t parse_address(const std::string& ip_port) { - auto colon = ip_port.find(":"); - if (colon == std::string::npos) { - std::ostringstream msg; - msg << "Can't parse address " << ip_port; - throw std::runtime_error(msg.str()); - } - std::string ip(ip_port.begin(), ip_port.begin() + colon); - std::string port(ip_port.begin() + colon + 1, ip_port.end()); - - return parse_address(ip, port); -} - /** * Load all addresses from the json hostfile. The hostfile is a list of * addresses in order of rank. For each rank there can be many addresses so @@ -357,15 +308,15 @@ address_t parse_address(const std::string& ip_port) { * ["ip3:5000", "ip3:5001"], * ] */ -std::vector> load_nodes(const char* hostfile) { - std::vector> nodes; +std::vector> load_nodes(const char* hostfile) { + std::vector> nodes; std::ifstream f(hostfile); json hosts = json::parse(f); for (auto& h : hosts) { - std::vector host; + std::vector host; for (auto& ips : h) { - host.push_back(parse_address(ips.get())); + host.push_back(std::move(detail::parse_address(ips.get()))); } nodes.push_back(std::move(host)); } @@ -377,73 +328,15 @@ std::vector> load_nodes(const char* hostfile) { * Create a socket and accept one connection for each of the provided * addresses. */ -std::vector accept_connections(const std::vector& addresses) { +std::vector accept_connections( + const std::vector& addresses) { std::vector sockets; int success; for (auto& address : addresses) { - // Create the socket to wait for connections from the peers - int sock = socket(AF_INET, SOCK_STREAM, 0); - if (sock < 0) { - std::ostringstream msg; - msg << "[ring] Couldn't create socket (error: " << errno << ")"; - throw std::runtime_error(msg.str()); - } - - // Make sure we can launch immediately after shutdown by setting the - // reuseaddr option so that we don't get address already in use errors - int enable = 1; - success = setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int)); - if (success < 0) { - shutdown(sock, 2); - close(sock); - std::ostringstream msg; - msg << "[ring] Couldn't enable reuseaddr (error: " << errno << ")"; - throw std::runtime_error(msg.str()); - } - success = setsockopt(sock, SOL_SOCKET, SO_REUSEPORT, &enable, sizeof(int)); - if (success < 0) { - shutdown(sock, 2); - close(sock); - std::ostringstream msg; - msg << "[ring] Couldn't enable reuseport (error: " << errno << ")"; - throw std::runtime_error(msg.str()); - } - - // Bind the socket to the address and port - success = bind(sock, address.get(), address.len); - if (success < 0) { - shutdown(sock, 2); - close(sock); - std::ostringstream msg; - msg << "[ring] Couldn't bind socket (error: " << errno << ")"; - throw std::runtime_error(msg.str()); - } - - // Wait for connections - success = listen(sock, 0); - if (success < 0) { - shutdown(sock, 2); - close(sock); - std::ostringstream msg; - msg << "[ring] Couldn't listen (error: " << errno << ")"; - throw std::runtime_error(msg.str()); - } - - int peer_socket = accept(sock, nullptr, nullptr); - if (peer_socket < 0) { - shutdown(sock, 2); - close(sock); - std::ostringstream msg; - msg << "[ring] Accept failed (error: " << errno << ")"; - throw std::runtime_error(msg.str()); - } - - // Close the listening socket - shutdown(sock, 2); - close(sock); - - sockets.push_back(peer_socket); + detail::TCPSocket socket(RING_TAG); + socket.listen(RING_TAG, address); + sockets.push_back(socket.accept(RING_TAG).detach()); } return sockets; @@ -454,93 +347,42 @@ std::vector accept_connections(const std::vector& addresses) { * provided addresses. */ std::vector make_connections( - const std::vector& addresses, + const std::vector& addresses, bool verbose) { std::vector sockets; int success; for (auto& address : addresses) { - int sock; - - // Attempt to connect to the peer CONN_ATTEMPTS times with exponential - // backoff. TODO: Do we need that? - for (int attempt = 0; attempt < CONN_ATTEMPTS; attempt++) { - // Create the socket - sock = socket(AF_INET, SOCK_STREAM, 0); - if (sock < 0) { - std::ostringstream msg; - msg << "[ring] Couldn't create socket (error: " << errno << ")"; - throw std::runtime_error(msg.str()); - } - - if (attempt > 0) { - int wait = (1 << (attempt - 1)) * CONN_WAIT; - log_info( - verbose, - "Attempt", - attempt, - "wait", - wait, - "ms (error:", - errno, - ")"); - std::this_thread::sleep_for(std::chrono::milliseconds(wait)); - } - - success = connect(sock, address.get(), address.len); - if (success == 0) { - break; - } - } - if (success < 0) { - std::ostringstream msg; - msg << "[ring] Couldn't connect (error: " << errno << ")"; - throw std::runtime_error(msg.str()); - } - - sockets.push_back(sock); + sockets.push_back(detail::TCPSocket::connect( + RING_TAG, + address, + CONN_ATTEMPTS, + CONN_WAIT, + [verbose](int attempt, int wait) { + log_info( + verbose, + "Attempt", + attempt, + "waiting", + wait, + "ms (error:", + errno, + ")"); + }) + .detach()); } return sockets; } -template -struct SumOp { - void operator()(const T* input, T* output, size_t N) { - while (N-- > 0) { - *output += *input; - input++; - output++; - } - } -}; - -template -struct MaxOp { - void operator()(const T* input, T* output, size_t N) { - while (N-- > 0) { - *output = std::max(*output, *input); - input++; - output++; - } - } -}; - -template -struct MinOp { - void operator()(const T* input, T* output, size_t N) { - while (N-- > 0) { - *output = std::min(*output, *input); - input++; - output++; - } - } -}; } // namespace class RingGroup : public GroupImpl { public: - RingGroup(int rank, std::vector> nodes, bool verbose) + RingGroup( + int rank, + std::vector> nodes, + bool verbose) : rank_(rank), verbose_(verbose), pool_(0) { if (rank_ > 0 && rank_ >= nodes.size()) { throw std::runtime_error( @@ -633,17 +475,17 @@ class RingGroup : public GroupImpl { void all_sum(const array& input, array& output, Stream stream) override { SWITCH_TYPE( - output, all_reduce>(input, output, stream, SumOp())); + output, all_reduce(input, output, stream, detail::SumOp())); } void all_max(const array& input, array& output, Stream stream) override { SWITCH_TYPE( - output, all_reduce>(input, output, stream, MaxOp())); + output, all_reduce(input, output, stream, detail::MaxOp())); } void all_min(const array& input, array& output, Stream stream) override { SWITCH_TYPE( - output, all_reduce>(input, output, stream, MinOp())); + output, all_reduce(input, output, stream, detail::MinOp())); } std::shared_ptr split(int color, int key = -1) override { diff --git a/mlx/distributed/utils.cpp b/mlx/distributed/utils.cpp new file mode 100644 index 00000000..ef52ec09 --- /dev/null +++ b/mlx/distributed/utils.cpp @@ -0,0 +1,204 @@ +// Copyright © 2025 Apple Inc. + +#include +#include +#include +#include +#include + +#include "mlx/distributed/utils.h" + +namespace mlx::core::distributed::detail { + +/** + * Parse a sockaddr from an ip and port provided as strings. + */ +address_t parse_address(const std::string& ip, const std::string& port) { + struct addrinfo hints, *res; + std::memset(&hints, 0, sizeof(hints)); + hints.ai_family = AF_UNSPEC; + hints.ai_socktype = SOCK_STREAM; + + int status = getaddrinfo(ip.c_str(), port.c_str(), &hints, &res); + if (status != 0) { + std::ostringstream msg; + msg << "Can't parse address " << ip << ":" << port; + throw std::runtime_error(msg.str()); + } + + address_t result; + memcpy(&result.addr, res->ai_addr, res->ai_addrlen); + result.len = res->ai_addrlen; + freeaddrinfo(res); + + return result; +} + +/** + * Parse a sockaddr provided as an : string. + */ +address_t parse_address(const std::string& ip_port) { + auto colon = ip_port.find(":"); + if (colon == std::string::npos) { + std::ostringstream msg; + msg << "Can't parse address " << ip_port; + throw std::runtime_error(msg.str()); + } + std::string ip(ip_port.begin(), ip_port.begin() + colon); + std::string port(ip_port.begin() + colon + 1, ip_port.end()); + + return parse_address(ip, port); +} + +TCPSocket::TCPSocket(const char* tag) { + sock_ = socket(AF_INET, SOCK_STREAM, 0); + if (sock_ < 0) { + std::ostringstream msg; + msg << tag << " Couldn't create socket (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } +} + +TCPSocket::TCPSocket(TCPSocket&& s) { + sock_ = s.sock_; + s.sock_ = -1; +} + +TCPSocket& TCPSocket::operator=(TCPSocket&& s) { + if (this != &s) { + sock_ = s.sock_; + s.sock_ = -1; + } + return *this; +} + +TCPSocket::TCPSocket(int s) : sock_(s) {} + +TCPSocket::~TCPSocket() { + if (sock_ > 0) { + shutdown(sock_, 2); + close(sock_); + } +} + +int TCPSocket::detach() { + int s = sock_; + sock_ = -1; + return s; +} + +void TCPSocket::listen(const char* tag, const address_t& addr) { + int success; + + // Make sure we can launch immediately after shutdown by setting the + // reuseaddr option so that we don't get address already in use errors + int enable = 1; + success = setsockopt(sock_, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int)); + if (success < 0) { + std::ostringstream msg; + msg << tag << " Couldn't enable reuseaddr (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + success = setsockopt(sock_, SOL_SOCKET, SO_REUSEPORT, &enable, sizeof(int)); + if (success < 0) { + std::ostringstream msg; + msg << tag << " Couldn't enable reuseport (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + + // Bind the socket to the address and port + success = bind(sock_, addr.get(), addr.len); + if (success < 0) { + std::ostringstream msg; + msg << tag << " Couldn't bind socket (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + + // Prepare waiting for connections + success = ::listen(sock_, 0); + if (success < 0) { + std::ostringstream msg; + msg << tag << " Couldn't listen (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } +} + +TCPSocket TCPSocket::accept(const char* tag) { + int peer = ::accept(sock_, nullptr, nullptr); + if (peer < 0) { + std::ostringstream msg; + msg << tag << " Accept failed (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + + return TCPSocket(peer); +} + +void TCPSocket::send(const char* tag, const void* data, size_t len) { + while (len > 0) { + auto n = ::send(sock_, data, len, 0); + if (n <= 0) { + std::ostringstream msg; + msg << tag << " Send failed with errno=" << errno; + throw std::runtime_error(msg.str()); + } + len -= n; + data = static_cast(data) + n; + } +} + +void TCPSocket::recv(const char* tag, void* data, size_t len) { + while (len > 0) { + auto n = ::recv(sock_, data, len, 0); + if (n <= 0) { + std::ostringstream msg; + msg << tag << " Recv failed with errno=" << errno; + throw std::runtime_error(msg.str()); + } + len -= n; + data = static_cast(data) + n; + } +} + +TCPSocket TCPSocket::connect( + const char* tag, + const address_t& addr, + int num_retries, + int wait, + std::function cb) { + int sock, success; + + // Attempt to connect `num_retries` times with exponential backoff. + for (int attempt = 0; attempt < num_retries; attempt++) { + // Create the socket + sock = socket(AF_INET, SOCK_STREAM, 0); + if (sock < 0) { + std::ostringstream msg; + msg << tag << " Couldn't create socket to connect (error: " << errno + << ")"; + throw std::runtime_error(msg.str()); + } + + success = ::connect(sock, addr.get(), addr.len); + if (success == 0) { + break; + } + + cb(attempt, wait); + if (wait > 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(wait)); + } + + wait <<= 1; + } + + if (success < 0) { + std::ostringstream msg; + msg << tag << " Couldn't connect (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + + return TCPSocket(sock); +} + +} // namespace mlx::core::distributed::detail diff --git a/mlx/distributed/utils.h b/mlx/distributed/utils.h new file mode 100644 index 00000000..213dd59a --- /dev/null +++ b/mlx/distributed/utils.h @@ -0,0 +1,67 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#include +#include +#include + +namespace mlx::core::distributed::detail { + +struct address_t { + sockaddr_storage addr; + socklen_t len; + + const sockaddr* get() const { + return (struct sockaddr*)&addr; + } +}; + +/** + * Parse a sockaddr from an ip and port provided as strings. + */ +address_t parse_address(const std::string& ip, const std::string& port); + +/** + * Parse a sockaddr provided as an : string. + */ +address_t parse_address(const std::string& ip_port); + +/** + * Small wrapper over a TCP socket to simplify initiating connections. + */ +class TCPSocket { + public: + TCPSocket(const char* tag); + TCPSocket(const TCPSocket&) = delete; + TCPSocket& operator=(const TCPSocket&) = delete; + TCPSocket(TCPSocket&& s); + TCPSocket& operator=(TCPSocket&&); + ~TCPSocket(); + + void listen(const char* tag, const address_t& addr); + TCPSocket accept(const char* tag); + + void send(const char* tag, const void* data, size_t len); + void recv(const char* tag, void* data, size_t len); + + int detach(); + + operator int() const { + return sock_; + } + + static TCPSocket connect( + const char* tag, + const address_t& addr, + int num_retries = 1, + int wait = 0, + std::function cb = nullptr); + + private: + TCPSocket(int sock); + + int sock_; +}; + +} // namespace mlx::core::distributed::detail diff --git a/python/mlx/_distributed_utils/common.py b/python/mlx/_distributed_utils/common.py new file mode 100644 index 00000000..16bf3f2b --- /dev/null +++ b/python/mlx/_distributed_utils/common.py @@ -0,0 +1,95 @@ +# Copyright © 2025 Apple Inc. + +import argparse +import ipaddress +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + + +@dataclass +class Host: + rank: int + ssh_hostname: str + ips: list[str] + rdma: list[Optional[str]] + + +class OptionalBoolAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + if option_string.startswith("--no-"): + setattr(namespace, self.dest, False) + else: + setattr(namespace, self.dest, True) + + +def positive_number(x): + x = int(x) + if x <= 0: + raise ValueError("Number should be positive") + return x + + +def log(verbose, *args, **kwargs): + if not verbose: + return + kwargs["file"] = sys.stderr + print("\033[32m[INFO]", *args, "\033[0m", **kwargs) + + +def log_warning(*args, **kwargs): + kwargs["file"] = sys.stderr + print("\033[33m[WARN]", *args, "\033[0m", **kwargs) + + +def log_error(*args, **kwargs): + kwargs["file"] = sys.stderr + print("\033[31m[ERROR]", *args, "\033[0m", **kwargs) + + +def parse_hostlist(parser, hostlist, repeats): + hosts = [] + for i, h in enumerate(hostlist.split(",")): + if h == "": + raise ValueError("Hostname cannot be empty") + try: + ipaddress.ip_address(h) + ips = [h] + except ValueError: + ips = [] + for i in range(repeats): + hosts.append(Host(i, h, ips, [])) + return hosts + + +def parse_hostfile(parser, hostfile): + """Parse the json hostfile that contains both the hostnames to ssh into and + the ips to communicate over when using the ring backend. + + Example: + + [ + {"ssh": "hostname1", "ips": ["123.123.123.1"], "rdma": [null, "rdma_en2", "rdma_en3"]}, + {"ssh": "hostname2", "ips": ["123.123.123.2"], "rdma": ["rdma_en2", null, "rdma_en3"]}, + ... + {"ssh": "hostnameN", "ips": ["123.123.123.N"], "rdma": ["rdma_en2", "rdma_en3", null]}, + ] + + Args: + hostfile (str): The path to the json file containing the host + information + """ + hostfile = Path(hostfile) + if not hostfile.exists(): + parser.error(f"Hostfile {str(hostfile)} doesn't exist") + + try: + hosts = [] + with open(hostfile) as f: + for i, h in enumerate(json.load(f)): + hosts.append(Host(i, h["ssh"], h.get("ips", []), h.get("rdma", []))) + return hosts + except Exception as e: + parser.error(f"Failed to parse hostfile {str(hostfile)} ({str(e)})") diff --git a/python/mlx/_distributed_utils/config.py b/python/mlx/_distributed_utils/config.py new file mode 100644 index 00000000..ac4e69bf --- /dev/null +++ b/python/mlx/_distributed_utils/config.py @@ -0,0 +1,570 @@ +# Copyright © 2025 Apple Inc. + +import argparse +import json +import shlex +import sys +import threading +from collections import defaultdict +from dataclasses import dataclass +from subprocess import DEVNULL, run +from typing import Optional + +import mlx.core as mx + +from .common import ( + Host, + OptionalBoolAction, + log, + log_error, + parse_hostfile, + parse_hostlist, +) + + +@dataclass +class SSHInfo: + can_ssh: bool + has_sudo: bool + + def __bool__(self): + return self.can_ssh + + +@dataclass +class ThunderboltPort: + iface: str + uuid: str + connected_to: Optional[str] + + +@dataclass +class ThunderboltHost: + name: str + ports: list[ThunderboltPort] + + +def add_ethernet_ips(hosts, verbose=False): + # Get the ips for each host + for h in hosts: + log(verbose, "Getting the ip from", h.ssh_hostname) + h.ips.append( + run( + ["ssh", h.ssh_hostname, "ipconfig", "getifaddr", "en0"], + capture_output=True, + text=True, + ).stdout.strip() + ) + + +def check_rdma(hosts, verbose=False): + # Check whether the hosts are capable of RDMA over thunderbolt + warn = False + for h in hosts: + log(verbose, "Checking that", h.ssh_hostname, "supports RDMA") + rdma_devs = ( + run(["ssh", h.ssh_hostname, "ibv_devices"], capture_output=True, text=True) + .stdout.strip() + .split() + ) + rdma_devs = [d for d in rdma_devs if d.startswith("rdma_")] + if not rdma_devs: + log_warning(h.ssh_hostname, "does not seem to have RDMA enabled") + warn = True + + if warn: + log_warning() + log_warning( + "Some of the hosts don't have RDMA enabled or they don't support RDMA." + ) + log_warning() + log_warning( + "See https://ml-explore.github.io/mlx/build/html/usage/distributed.html" + ) + log_warning("for instructions on how to enable RDMA.") + + +def can_auto_setup(hosts, sshinfo, auto_setup=False): + has_sudo = all(info.has_sudo for info in sshinfo) + if not has_sudo and auto_setup: + log_warning( + "Automatic setup requested but the following hosts do not have passwordless sudo" + ) + for h, i in zip(hosts, sshinfo): + if not i.has_sudo: + log_warning(" - ", h.ssh_hostname) + return has_sudo + + +class IPConfigurator: + def __init__(self, hosts, tb_hosts, uuid_reverse_index): + assigned = set() + ips = defaultdict(list) + ip0 = 0 + ip1 = 0 + for src_node, h in enumerate(tb_hosts): + for src_port, p in enumerate(h.ports): + if not p.connected_to: + continue + if p.connected_to not in uuid_reverse_index: + continue + if (src_node, src_port) in assigned: + continue + + dst_node, dst_port = uuid_reverse_index[p.connected_to] + + ip_src = f"192.168.{ip0}.{ip1 + 1}" + ip_dst = f"192.168.{ip0}.{ip1 + 2}" + iface_src = p.iface + iface_dst = tb_hosts[dst_node].ports[dst_port].iface + + ips[src_node, dst_node].append((iface_src, ip_src)) + ips[dst_node, src_node].append((iface_dst, ip_dst)) + + assigned.add((src_node, src_port)) + assigned.add((dst_node, dst_port)) + + ip1 += 4 + if ip1 > 255: + ip0 += 1 + ip1 = 0 + if ip0 > 255: + raise ValueError("Ran out of available local IPs") + + self.ips = ips + self.hosts = hosts + self.tb_hosts = tb_hosts + + def setup(self, verbose=False, auto_setup=False): + netmask = "255.255.255.252" + for i, (h, th) in enumerate(zip(self.hosts, self.tb_hosts)): + command = "" + command += "sudo ifconfig bridge0 down\n" + for j in range(len(self.hosts)): + if i == j or (i, j) not in self.ips: + continue + for (iface, ip), (_, peer) in zip(self.ips[i, j], self.ips[j, i]): + command += f"sudo ifconfig {iface} inet {ip} netmask {netmask}\n" + command += f"sudo route change {peer} -interface {iface}\n" + if auto_setup: + print(f"Running auto setup for {h.ssh_hostname}") + command = command.strip().replace("\n", " ; ") + command = ["ssh", h.ssh_hostname, command] + log(verbose, shlex.join(command)) + run(command) + else: + msg = f"Setup for {h.ssh_hostname}" + print(msg) + print("=" * len(msg)) + print(command) + input("Enter to continue") + print() + + +def parse_hardware_ports(ports_string): + ports = {} + port_name = None + for l in ports_string.decode("utf-8").split("\n"): + if l.startswith("Hardware Port:"): + port_name = l.strip()[15:] + elif l.startswith("Device:"): + ports[port_name] = l.strip()[8:] + port_name = None + return ports + + +def extract_connectivity(hosts, verbose): + # Extract the current connectivity from the remote hosts + thunderbolt_connections = [] + for h in hosts: + log(verbose, "Getting connectivity from", h.ssh_hostname) + thunderbolt_connections.append( + json.loads( + run( + [ + "ssh", + h.ssh_hostname, + "system_profiler", + "SPThunderboltDataType", + "-json", + ], + capture_output=True, + ).stdout + ) + ) + interface_maps = [] + for h in hosts: + log(verbose, "Getting interface names from", h.ssh_hostname) + interface_maps.append( + parse_hardware_ports( + run( + [ + "ssh", + h.ssh_hostname, + "networksetup", + "-listallhardwareports", + ], + capture_output=True, + ).stdout + ) + ) + + # Parse the connectivity into some simple dataclasses + tb_hosts = [] + for c, iface_map in zip(thunderbolt_connections, interface_maps): + name = "" + ports = [] + for t in c["SPThunderboltDataType"]: + uuid = t.get("domain_uuid_key") + if uuid is None: + continue + name = t["device_name_key"] + tag = t["receptacle_1_tag"]["receptacle_id_key"] + items = t.get("_items", []) + connected_items = [item for item in items if "domain_uuid_key" in item] + connected_to = ( + connected_items[0]["domain_uuid_key"] if connected_items else None + ) + iface = iface_map[f"Thunderbolt {tag}"] + ports.append(ThunderboltPort(iface, uuid, connected_to)) + tb_hosts.append(ThunderboltHost(name, sorted(ports, key=lambda x: x.iface))) + + # Create a reverse index to be able to map uuids to (host, port) quickly + uuid_reverse_index = {} + for i, h in enumerate(tb_hosts): + for j, p in enumerate(h.ports): + uuid_reverse_index[p.uuid] = (i, j) + + return tb_hosts, uuid_reverse_index + + +def make_connectivity_matrix(tb_hosts, uuid_reverse_index): + connectivity = [] + for i, h in enumerate(tb_hosts): + c = [0] * len(tb_hosts) + for p in h.ports: + if p.connected_to in uuid_reverse_index: + j, _ = uuid_reverse_index[p.connected_to] + c[j] += 1 + connectivity.append(c) + return connectivity + + +def tb_connectivity_to_dot(hosts, tb_hosts, uuid_reverse_index): + # Make ids per node + names = [] + for i in range(len(tb_hosts)): + n = "" + j = i + while True: + n += chr(97 + j % 26) + j //= 26 + if j == 0: + break + names.append(n) + + print("graph G {") + print(" node [shape=rectangle];") + for i, h in enumerate(hosts): + print(f' {names[i]} [label="{h.ssh_hostname}"];') + for i, h in enumerate(tb_hosts): + for p in h.ports: + if not p.connected_to: + continue + dst = uuid_reverse_index[p.connected_to] + if dst[0] < i: + continue + print(f" {names[i]} -- {names[dst[0]]}", end="") + print(f' [label="{p.iface}/{tb_hosts[dst[0]].ports[dst[1]].iface}"]') + print("}") + + +def extract_rings(connectivity): + rings = [] + existing_rings = set() + num_nodes = len(connectivity) + + def dfs(start_node, node, path, visited): + path.append(node) + visited.add(node) + for j in range(num_nodes): + if connectivity[node][j] <= 0: + continue + if j == start_node: + yield path[:] + if j not in visited: + yield from dfs(start_node, j, path, visited) + path.pop() + visited.remove(node) + + for start in range(num_nodes): + for r in dfs(start, start, [], set()): + cnt = min(connectivity[r[i]][r[(i + 1) % len(r)]] for i in range(len(r))) + rkey = tuple(sorted(r)) + if rkey not in existing_rings: + rings.append((r, cnt)) + existing_rings.add(rkey) + + return sorted(rings, key=lambda x: -len(x[0])) + + +def check_valid_mesh(hosts, connectivity, strict=True): + num_nodes = len(connectivity) + for i in range(num_nodes): + for j in range(num_nodes): + if i == j: + continue + if connectivity[i][j] <= 0: + if strict: + log_error( + f"Incomplete mesh, {hosts[i].ssh_hostname} is not connected to {hosts[j].ssh_hostname}" + ) + log_error() + log_error("Try passing --dot to visualize the connectivity") + sys.exit(1) + else: + return False + return True + + +def check_ssh_connections(hosts): + results = [None] * len(hosts) + + def _check(hostname, i): + info = SSHInfo(False, False) + results[i] = info + + # Check for ssh + result = run( + [ + "ssh", + "-o", + "BatchMode=yes", + "-o", + "ConnectTimeout=5", + hostname, + "echo", + "success", + ], + stdout=DEVNULL, + stderr=DEVNULL, + ) + info.can_ssh = result.returncode == 0 + if not info.can_ssh: + return + + # Check for sudo + result = run( + [ + "ssh", + "-o", + "BatchMode=yes", + "-o", + "ConnectTimeout=5", + hostname, + "sudo", + "ls", + ], + stdout=DEVNULL, + stderr=DEVNULL, + ) + info.has_sudo = result.returncode == 0 + + threads = [ + threading.Thread(target=_check, args=(h.ssh_hostname, i)) + for i, h in enumerate(hosts) + ] + for t in threads: + t.start() + for t in threads: + t.join() + + if not all(results): + log_error("Could not ssh to the following hosts:") + for i, h in enumerate(hosts): + if not results[i]: + log_error(" - ", h.ssh_hostname) + log_error() + log_error("Maybe they are not set-up for password-less ssh?") + sys.exit(1) + + return results + + +def prepare_ethernet_hostfile(args, hosts): + log(args.verbose, f"Preparing an ethernet hostfile") + add_ethernet_ips(hosts, args.verbose) + + hostfile = [] + for h in hosts: + hostfile.append(dict(ssh=h.ssh_hostname, ips=h.ips)) + + if args.output_hostfile: + with open(args.output_hostfile, "w") as f: + json.dump(hostfile, f, indent=4) + else: + print("Hostfile") + print("========") + print(json.dumps(hostfile, indent=4)) + + +def configure_ring(args, hosts, ips, ring, sshinfo): + log(args.verbose, "Prepare a ring hostfile") + ring, count = ring + hostfile = [] + for i, node in enumerate(ring): + h = hosts[node] + peer = ring[i - 1] + hostfile.append( + { + "ssh": h.ssh_hostname, + "ips": [ips.ips[node, peer][c][1] for c in range(count)], + "rdma": [], + } + ) + + has_sudo = can_auto_setup(hosts, sshinfo, args.auto_setup) + ips.setup(verbose=args.verbose, auto_setup=args.auto_setup and has_sudo) + + if args.output_hostfile: + with open(args.output_hostfile, "w") as f: + json.dump(hostfile, f, indent=4) + else: + print("Hostfile") + print("========") + print(json.dumps(hostfile, indent=4)) + + +def configure_jaccl(args, hosts, ips, sshinfo): + log(args.verbose, "Prepare a jaccl hostfile") + check_rdma(hosts, args.verbose) + add_ethernet_ips(hosts, args.verbose) + + hostfile = [] + for i, h in enumerate(hosts): + rdma = [] + for j in range(len(hosts)): + if i == j: + rdma.append(None) + else: + rdma.append(f"rdma_{ips.ips[i, j][0][0]}") + hostfile.append({"ssh": h.ssh_hostname, "ips": h.ips, "rdma": rdma}) + + has_sudo = can_auto_setup(hosts, sshinfo, args.auto_setup) + ips.setup(verbose=args.verbose, auto_setup=args.auto_setup and has_sudo) + + if args.output_hostfile: + with open(args.output_hostfile, "w") as f: + json.dump(hostfile, f, indent=4) + else: + print("Hostfile") + print("========") + print(json.dumps(hostfile, indent=4)) + + +def prepare_tb_hostfile(args, hosts, sshinfo): + log(args.verbose, f"Preparing for communication over thunderbolt") + tb_hosts, uuid_reverse_index = extract_connectivity(hosts, args.verbose) + + if args.dot: + tb_connectivity_to_dot(hosts, tb_hosts, uuid_reverse_index) + return + + ips = IPConfigurator(hosts, tb_hosts, uuid_reverse_index) + connectivity = make_connectivity_matrix(tb_hosts, uuid_reverse_index) + + if args.backend is None: + rings = extract_rings(connectivity) + has_mesh = check_valid_mesh(hosts, connectivity, False) + has_ring = len(rings) > 0 and len(rings[0][0]) == len(hosts) + + if not has_ring and not has_mesh: + log_error("Neither thunderbolt mesh nor ring found.") + log_error("Perhaps run with --dot to generate a plot of the connectivity.") + sys.exit(1) + + elif has_ring: + configure_ring(args, hosts, ips, rings[0], sshinfo) + + else: + configure_jaccl(args, hosts, ips, sshinfo) + + elif args.backend == "ring": + rings = extract_rings(connectivity) + has_ring = len(rings) > 0 and len(rings[0][0]) == len(hosts) + if not has_ring: + log_error("Could not find a full ring.") + log_error() + log_error("Try passing --dot to visualize the connectivity") + if len(rings) > 0: + log_error("Rings found:") + for r in rings: + log_error(f" - {','.join(hosts[i].ssh_hostname for i in r)}") + sys.exit(1) + configure_ring(args, hosts, ips, rings[0], sshinfo) + + elif args.backend == "jaccl": + check_valid_mesh(hosts, connectivity) + configure_jaccl(args, hosts, ips, sshinfo) + + +def main(): + parser = argparse.ArgumentParser( + description="Configure remote machines for use with MLX distributed" + ) + parser.add_argument( + "--verbose", action="store_true", help="Print debug messages in stdout" + ) + parser.add_argument( + "--hosts", default="127.0.0.1", help="A comma separated list of hosts" + ) + parser.add_argument("--hostfile", help="The file containing the hosts") + parser.add_argument( + "--over", + choices=["thunderbolt", "ethernet"], + default="thunderbolt", + help="What type of connectivity to configure", + required=True, + ) + parser.add_argument( + "--output-hostfile", help="If provided, save the hostfile to this path" + ) + parser.add_argument( + "--auto-setup", + "--no-auto-setup", + action=OptionalBoolAction, + nargs=0, + dest="auto_setup", + default=None, + ) + parser.add_argument( + "--dot", action="store_true", help="Output the topology in DOT format and exit" + ) + parser.add_argument( + "--backend", + choices=["ring", "jaccl"], + default=None, + help="Which distributed backend to configure", + ) + args = parser.parse_args() + + if args.hostfile is not None: + hosts = parse_hostfile(parser, args.hostfile) + else: + hosts = parse_hostlist(parser, args.hosts, 1) + + # Check that we can ssh + log( + args.verbose, + f"Checking for ssh access for {', '.join(h.ssh_hostname for h in hosts)}", + ) + sshinfo = check_ssh_connections(hosts) + + # Prepare a hostfile for communication over ethernet using the ips of the + # provided hostnames + if args.over == "ethernet": + prepare_ethernet_hostfile(args, hosts) + + # Configure the macs for communication over thunderbolt, both via RDMA and IP + else: + prepare_tb_hostfile(args, hosts, sshinfo) diff --git a/python/mlx/_distributed_utils/launch.py b/python/mlx/_distributed_utils/launch.py new file mode 100644 index 00000000..56c7323e --- /dev/null +++ b/python/mlx/_distributed_utils/launch.py @@ -0,0 +1,557 @@ +# Copyright © 2025 Apple Inc. + +import argparse +import base64 +import json +import os +import shlex +import shutil +import sys +import tempfile +import threading +from collections import Counter +from itertools import chain +from pathlib import Path +from queue import Empty as QueueEmpty +from queue import Queue +from select import select +from subprocess import PIPE, Popen, run + +import mlx.core as mx + +from .common import log, log_warning, parse_hostfile, parse_hostlist, positive_number + + +class CommandProcess: + @property + def process(self): + """Return the Popen object that refers to the current command.""" + raise NotImplementedError() + + @property + def exit_status(self): + """Return a tuple (returncode, killed) for the command. It should be + (None, None) while the command is running normally.""" + raise NotImplementedError() + + def preprocess_output(self, data: str, is_stdout=False): + """Preprocess the output of the command so that extra data can be + capture or the format changed on the fly.""" + raise NotImplementedError() + + def terminate(self): + """Terminate or return the exit code.""" + raise NotImplementedError() + + +class RemoteProcess(CommandProcess): + def __init__(self, rank, host, python, cwd, files, env, command): + is_local = host == "127.0.0.1" + cmd = RemoteProcess.make_launch_script(rank, cwd, files, env, command) + if not is_local: + cmd = f"ssh -tt -o LogLevel=QUIET {host} {shlex.quote(cmd)}" + + self._host = host + self._pidfile = None + self._is_local = is_local + self._process = Popen( + cmd, + shell=True, + executable="/bin/bash", + stdin=PIPE, + stdout=PIPE, + stderr=PIPE, + ) + + self._killed = False + + @property + def process(self): + return self._process + + @property + def exit_status(self): + return self._process.poll(), self._killed + + def preprocess_output(self, data, is_stdout=False): + if self._pidfile is None: + pidfile, *rest = data.split("\n", maxsplit=1) + self._pidfile = pidfile + return rest[0] if rest else "" + + return data + + def terminate(self): + if self._killed: + return + + self._process.terminate() + self._process.wait() + + # Kill the remote program if possible + cmd = RemoteProcess.make_kill_script(self._pidfile) + if not self._is_local: + cmd = f"ssh {self._host} {shlex.quote(cmd)}" + c = run( + cmd, + check=True, + shell=True, + executable="/bin/bash", + capture_output=True, + text=True, + ) + + self._killed = c.stdout.strip() == "1" + + @staticmethod + def make_launch_script(rank, cwd, files, env, command): + script = "" + + # Disable echo + script = "stty -echo; " + + # Write the PID to a file so we can kill the process if needed + script += "pidfile=$(mktemp); " + script += "echo $$ > $pidfile; " + script += 'printf "%s\\n" $pidfile; ' + + # Change the working directory if one was requested. Otherwise attempt to + # change to the current one but don't fail if it wasn't possible. + d = cwd or os.getcwd() + script += f"if [[ -d {repr(d)} ]]; then " + script += f" cd {repr(d)}; " + if cwd is not None: + script += "else " + script += f" echo 'Failed to change directory to' {repr(d)} >2; " + script += "fi; " + + # Add the environment variables that were requested + for e in env: + key, *value = e.split("=", maxsplit=1) + value = shlex.quote(value[0]) if len(value) > 0 else "" + if not all(c.isalnum() or c == "_" for c in key): + log_warning( + f"'{e}' is an invalid environment variable so it is ignored" + ) + continue + script += f"export {key}={value}; " + + # Make the temporary files + for env_name, content in files.items(): + script += "fname=$(mktemp); " + script += f"echo {shlex.quote(content)} >$fname; " + script += f"export {env_name}=$fname; " + + # Finally add the rank + script += f"export MLX_RANK={rank}; " + + # Replace the process with the script + script += f"cmd=({' '.join(map(shlex.quote, command))}); " + script += 'exec "${cmd[@]}"' + + return script + + @staticmethod + def make_kill_script(pidfile): + script = "" + script += f"pid=$(cat {pidfile}); " + script += "if ps -p $pid >/dev/null; then " + script += " kill $pid; " + script += " echo 1; " + script += "else " + script += " echo 0; " + script += "fi; " + script += f"rm {pidfile}" + + return script + + +def _launch_with_io(command_class, arguments, verbose): + stop = False + exit_codes = [(None, None)] * len(arguments) + + def _thread_fn(rank, *args, **kwargs): + stdin_queue = kwargs.pop("stdin_queue") + stdout_queue = kwargs.pop("stdout_queue") + stderr_queue = kwargs.pop("stderr_queue") + + command = command_class(rank, *args, **kwargs) + p = command.process + os.set_blocking(p.stdout.fileno(), False) + os.set_blocking(p.stderr.fileno(), False) + os.set_blocking(p.stdin.fileno(), False) + + to_read = [p.stdout.fileno(), p.stderr.fileno()] + to_write = [p.stdin.fileno()] + + stdin_buffer = b"" + while p.poll() is None: + try: + stdin_buffer += stdin_queue.get_nowait() + except QueueEmpty: + pass + rlist, wlist, _ = select(to_read, to_write, [], 1.0) + for fd in rlist: + is_stdout = fd == p.stdout.fileno() + msg = os.read(fd, 8192).decode(errors="ignore") + msg = command.preprocess_output(msg, is_stdout) + if is_stdout: + stdout_queue.put(msg.encode()) + else: + stderr_queue.put(msg.encode()) + for fd in wlist: + if len(stdin_buffer) > 0: + n = os.write(fd, stdin_buffer) + stdin_buffer = stdin_buffer[n:] + if stop: + command.terminate() + break + exit_codes[rank] = command.exit_status + + if exit_codes[rank][1]: + log_warning(f"Node with rank {rank} was killed") + elif exit_codes[rank][0] != 0: + log_warning(f"Node with rank {rank} exited with code {exit_codes[rank][0]}") + else: + log(verbose, f"Node with rank {rank} completed") + + stdin_queues = [] + stdout_queues = [] + stderr_queues = [] + threads = [] + for i, (args, kwargs) in enumerate(arguments): + stdin_queues.append(Queue()) + stdout_queues.append(Queue()) + stderr_queues.append(Queue()) + t = threading.Thread( + target=_thread_fn, + args=args, + kwargs=kwargs + | { + "stdin_queue": stdin_queues[-1], + "stdout_queue": stdout_queues[-1], + "stderr_queue": stderr_queues[-1], + }, + ) + t.start() + threads.append(t) + + os.set_blocking(sys.stdin.fileno(), False) + os.set_blocking(sys.stdout.fileno(), True) + os.set_blocking(sys.stderr.fileno(), True) + while not stop or any(not q.empty() for q in chain(stdout_queues, stderr_queues)): + # Broadcast user input to the jobs + rlist, _, _ = select([sys.stdin.fileno()], [], [], 0.1) + for fd in rlist: + stdin_buffer = os.read(fd, 8192) + for q in stdin_queues: + q.put(stdin_buffer) + + # Gather job output + for q in stdout_queues: + try: + while not q.empty(): + sys.stdout.buffer.write(q.get_nowait()) + except QueueEmpty: + pass + for q in stderr_queues: + try: + while not q.empty(): + sys.stderr.buffer.write(q.get_nowait()) + except QueueEmpty: + pass + sys.stdout.buffer.flush() + sys.stderr.buffer.flush() + + # Check if all are running and terminate otherwise + if any(t.is_alive() for t in threads): + for i, t in enumerate(threads): + if not t.is_alive(): + if exit_codes[i][0] != 0: + stop = True + break + else: + break + + # Wait for the jobs to finish + for t in threads: + t.join() + + # Process any remaining outputs + for q in stdout_queues: + while not q.empty(): + sys.stdout.buffer.write(q.get()) + for q in stderr_queues: + while not q.empty(): + sys.stderr.buffer.write(q.get()) + sys.stdout.buffer.flush() + sys.stderr.buffer.flush() + + +def launch_ring(parser, hosts, args, command): + if any(len(h.ips) == 0 for h in hosts): + parser.error( + "The ring backend requires IPs to be provided instead of hostnames" + ) + + port = args.starting_port + ring_hosts = [] + for h in hosts: + node = [] + for ip in h.ips: + for i in range(args.connections_per_ip): + node.append(f"{ip}:{port}") + port += 1 + ring_hosts.append(node) + hostfile = json.dumps(ring_hosts) if len(ring_hosts) > 1 else "" + + files = {"MLX_HOSTFILE": hostfile} + env = args.env + if args.verbose: + env.append("MLX_RING_VERBOSE=1") + cwd = args.cwd + + log(args.verbose, "Running", shlex.join(command)) + + _launch_with_io( + RemoteProcess, + [ + ((rank, h.ssh_hostname, args.python, cwd, files, env, command), {}) + for rank, h in enumerate(hosts) + ], + args.verbose, + ) + + +def launch_nccl(parser, hosts, args, command): + if not hosts[0].ips: + raise ValueError("Rank 0 should have an IP reachable from all other ranks") + + master_host = hosts[0].ips[0] + master_port = args.nccl_port + world_size = len(hosts) + + env = args.env + cwd = args.cwd + if args.verbose: + env.append("NCCL_DEBUG=INFO") + env.append(f"NCCL_HOST_IP={master_host}") + env.append(f"NCCL_PORT={master_port}") + env.append(f"MLX_WORLD_SIZE={world_size}") + + log(args.verbose, "Running", shlex.join(command)) + + _launch_with_io( + RemoteProcess, + [ + ( + ( + rank, + h.ssh_hostname, + args.python, + cwd, + {}, + env + [f"CUDA_VISIBLE_DEVICES={rank % args.repeat_hosts}"], + command, + ), + {}, + ) + for rank, h in enumerate(hosts) + ], + args.verbose, + ) + + +def launch_jaccl(parser, hosts, args, command): + if not hosts[0].ips: + raise ValueError("Rank 0 should have an IP reachable from all other ranks") + + have_rdmas = all(len(h.rdma) == len(hosts) for h in hosts) + have_nulls = all(h.rdma[i] is None for i, h in enumerate(hosts)) + if not have_rdmas or not have_nulls: + raise ValueError("Malformed hostfile for jaccl backend") + + coordinator = hosts[0].ips[0] + 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])} + + log(args.verbose, "Running", shlex.join(command)) + + _launch_with_io( + RemoteProcess, + [ + ((rank, h.ssh_hostname, args.python, cwd, files, env, command), {}) + for rank, h in enumerate(hosts) + ], + args.verbose, + ) + + +def get_mpi_libname(): + try: + ompi_info = run(["which", "ompi_info"], check=True, capture_output=True) + ompi_info = ompi_info.stdout.strip().decode() + + if platform.system() == "Darwin": + otool_output = run( + ["otool", "-L", ompi_info], check=True, capture_output=True + ) + else: + otool_output = run(["ldd", ompi_info], check=True, capture_output=True) + otool_output = otool_output.stdout.decode() + + # StopIteration if not found + libmpi_line = next( + filter(lambda line: "libmpi" in line, otool_output.splitlines()) + ) + return libmpi_line.strip().split()[0].removeprefix("@rpath/") + except: + return None + + +def launch_mpi(parser, hosts, args, command): + mpirun = run(["which", "mpirun"], check=True, capture_output=True) + mpirun = mpirun.stdout.strip().decode() + + # Compatibility with homebrew and pip installs + mpi_libname = get_mpi_libname() + if mpi_libname is not None: + dyld = Path(mpirun).parent.parent / "lib" + args.env = [ + f"DYLD_LIBRARY_PATH={str(dyld)}", + f"MLX_MPI_LIBNAME={mpi_libname}", + ] + args.env + + log(args.verbose, f"Using '{mpirun}'") + with tempfile.NamedTemporaryFile(mode="w") as f: + hosts = Counter((h.ssh_hostname for h in hosts)) + for h, n in hosts.items(): + print(f"{h} slots={n}", file=f) + f.flush() + + cmd = [ + mpirun, + "--output", + ":raw", # do not line buffer output + "--hostfile", + f.name, + *(["-cwd", args.cwd] if args.cwd else []), + *sum((["-x", e] for e in args.env), []), + *sum([shlex.split(arg) for arg in args.mpi_arg], []), + "--", + *command, + ] + log(args.verbose, "Running", " ".join(cmd)) + try: + run(cmd) + except KeyboardInterrupt: + pass + + +def main(): + parser = argparse.ArgumentParser(description="Launch an MLX distributed program") + parser.add_argument( + "--print-python", + action="store_true", + help="Print the path to the current python executable and exit", + ) + parser.add_argument( + "--verbose", action="store_true", help="Print debug messages in stdout" + ) + parser.add_argument( + "--hosts", default="127.0.0.1", help="A comma separated list of hosts" + ) + parser.add_argument( + "--repeat-hosts", + "-n", + type=positive_number, + default=1, + help="Repeat each host a given number of times", + ) + parser.add_argument("--hostfile", help="The file containing the hosts") + parser.add_argument( + "--backend", + choices=["ring", "mpi", "nccl", "jaccl"], + default="nccl" if mx.cuda.is_available() else "ring", + help="Which distributed backend to launch", + ) + parser.add_argument( + "--env", + action="append", + default=[], + help="Set environment variables for the jobs", + ) + parser.add_argument( + "--mpi-arg", + action="append", + default=[], + help="Arguments to pass directly to mpirun", + ) + parser.add_argument( + "--connections-per-ip", + default=1, + type=int, + help="How many connections per ip to use for the ring backend", + ) + parser.add_argument( + "--starting-port", + "-p", + type=int, + default=32323, + help="For the ring backend listen on this port increasing by 1 per rank and IP", + ) + parser.add_argument( + "--cwd", help="Set the working directory on each node to the provided one" + ) + parser.add_argument( + "--nccl-port", + type=int, + default=12345, + help="The port to use for the NCCL communication (only for nccl backend)", + ) + parser.add_argument( + "--no-verify-script", + action="store_false", + dest="verify_script", + help="Do not verify that the script exists", + ) + parser.add_argument( + "--python", default=sys.executable, help="Use this python on the remote hosts" + ) + + args, rest = parser.parse_known_args() + + if args.print_python: + print(args.python) + return + + if len(rest) == 0: + parser.error("No script is provided") + if rest[0] == "--": + rest.pop(0) + + # Try to extract a list of hosts and corresponding ips + if args.hostfile is not None: + hosts = parse_hostfile(parser, args.hostfile) + else: + hosts = parse_hostlist(parser, args.hosts, args.repeat_hosts) + + # Check if the script is a file and convert it to a full path + if (script := Path(rest[0])).exists() and script.is_file(): + rest[0:1] = [args.python, str(script.resolve())] + elif (command := shutil.which(rest[0])) is not None: + rest[0] = command + elif args.verify_script: + raise ValueError(f"Invalid script or command {rest[0]}") + + # Launch + if args.backend == "ring": + launch_ring(parser, hosts, args, rest) + if args.backend == "mpi": + launch_mpi(parser, hosts, args, rest) + if args.backend == "nccl": + launch_nccl(parser, hosts, args, rest) + if args.backend == "jaccl": + launch_jaccl(parser, hosts, args, rest) diff --git a/python/mlx/distributed_run.py b/python/mlx/distributed_run.py deleted file mode 100644 index e4b50a5c..00000000 --- a/python/mlx/distributed_run.py +++ /dev/null @@ -1,909 +0,0 @@ -# Copyright © 2025 Apple Inc. - -import argparse -import base64 -import ipaddress -import json -import os -import platform -import shlex -import shutil -import sys -import tempfile -import threading -import time -from collections import Counter -from dataclasses import dataclass -from pathlib import Path -from queue import Empty as QueueEmpty -from queue import Queue -from select import select -from subprocess import PIPE, Popen, run -from typing import Optional - -import mlx.core as mx - - -@dataclass -class Host: - rank: int - ssh_hostname: str - ips: list[str] - - -@dataclass -class ThunderboltPort: - iface: str - uuid: str - connected_to: Optional[str] - - -@dataclass -class ThunderboltHost: - name: str - ports: list[ThunderboltPort] - - -def parse_hardware_ports(ports_string): - ports = {} - port_name = None - for l in ports_string.decode("utf-8").split("\n"): - if l.startswith("Hardware Port:"): - port_name = l.strip()[15:] - elif l.startswith("Device:"): - ports[port_name] = l.strip()[8:] - port_name = None - return ports - - -def get_num_nvidia_gpus(): - result = run(["nvidia-smi", "-L"], capture_output=True, text=True, check=True) - return len(result.stdout.strip().split("\n")) - - -def extract_rings(hosts, index): - def usable_port(i, j, used_ports): - return (i, j) not in used_ports and hosts[i].ports[j].connected_to is not None - - def dfs(start_node, node, path, visited, used_ports): - path.append(node) - visited.add(node) - for j, p in enumerate(hosts[node].ports): - if not usable_port(node, j, used_ports): - continue - next_node, _ = index[p.connected_to] - if next_node == start_node: - yield path[:] - if next_node not in visited: - yield from dfs(start_node, next_node, path, visited, used_ports) - path.pop() - visited.remove(node) - - # Concretize maps the found cycle to real thunderbolt ports. It also adds - # those ports to the used set so next cycles can't use them again. - def concretize(cycle, used_ports): - concrete_path = [] - for n1, n2 in zip(cycle, cycle[1:] + cycle[:1]): - for j, p in enumerate(hosts[n1].ports): - if not usable_port(n1, j, used_ports): - continue - n2_hat, nj = index[p.connected_to] - if n2 == n2_hat: - concrete_path.append(((n1, j), (n2, nj))) - used_ports.add((n1, j)) - used_ports.add((n2, nj)) - break - if concrete_path[-1][0][0] != n1: - raise RuntimeError("Couldn't concretize the cycle") - return concrete_path - - # Normalize tries to ensure that the cycles have the same direction so we can - # use them together. We achieve this by selecting the direction such that - # the smallest rank hosts connect to larger rank hosts. - def normalize(path): - small_to_large = sum(1 for p in path if p[0][0] < p[1][0]) - if small_to_large > len(path) - small_to_large: - return path - else: - return [(p[1], p[0]) for p in path] - - rings = [] - used_ports = set() - for start_node in range(len(hosts)): - while True: - ring = [] - for r in dfs(start_node, start_node, [], set(), used_ports): - if len(r) > len(ring): - ring = r - # Break early since we won't find a bigger ring no matter what - if len(ring) == len(hosts): - break - if not ring: - break - try: - rings.append(normalize(concretize(ring, used_ports))) - except RuntimeError: - if len(rings) > 0: - return rings - raise - - return rings - - -def positive_number(x): - x = int(x) - if x <= 0: - raise ValueError("Number should be positive") - return x - - -def log(verbose, *args, **kwargs): - if not verbose: - return - print("\033[32m[INFO]", *args, "\033[0m", **kwargs) - - -def log_warning(*args, **kwargs): - kwargs["file"] = sys.stderr - print("\033[33m[WARN]", *args, "\033[0m", **kwargs) - - -def log_error(*args, **kwargs): - kwargs["file"] = sys.stderr - print("\033[31m[ERROR]", *args, "\033[0m", **kwargs) - - -def parse_hostfile(parser, hostfile): - """Parse the json hostfile that contains both the hostnames to ssh into and - the ips to communicate over when using the ring backend. - - Example: - - [ - {"ssh": "hostname1", "ips": ["123.123.123.1"]}, - {"ssh": "hostname2", "ips": ["123.123.123.2"]}, - ... - {"ssh": "hostnameN", "ips": ["123.123.123.N"]}, - ] - - Args: - hostfile (str): The path to the json file containing the host - information - """ - hostfile = Path(hostfile) - if not hostfile.exists(): - parser.error(f"Hostfile {str(hostfile)} doesn't exist") - - try: - hosts = [] - with open(hostfile) as f: - for i, h in enumerate(json.load(f)): - hosts.append(Host(i, h["ssh"], h.get("ips", []))) - return hosts - except Exception as e: - parser.error(f"Failed to parse hostfile {str(hostfile)} ({str(e)})") - - -def parse_hostlist(parser, hostlist, repeats): - hosts = [] - for i, h in enumerate(hostlist.split(",")): - if h == "": - raise ValueError("Hostname cannot be empty") - try: - ipaddress.ip_address(h) - ips = [h] - except ValueError: - ips = [] - for i in range(repeats): - hosts.append(Host(i, h, ips)) - return hosts - - -def make_monitor_script(rank, hostfile, cwd, env, command, verbose): - # Imports that are used throughout - script = "" - script += "import os\n" - script += "import sys\n" - script += "import tempfile\n" - script += "from pathlib import Path\n" - - # Write the PID to a file so we can kill the process if needed - script += "_, pidfile = tempfile.mkstemp() \n" - script += "open(pidfile, 'w').write(str(os.getpid()))\n" - script += "print(pidfile, flush=True)\n" - - # Change the working directory if one was requested. Otherwise attempt to - # change to the current one but don't fail if it wasn't possible. - d = cwd or os.getcwd() - script += f"if Path({repr(d)}).exists():\n" - script += f" os.chdir({repr(d)})\n" - if cwd is not None: - script += "else:\n" - script += ( - f" print('Failed to change directory to', {repr(d)}, file=sys.stderr)\n" - ) - script += f" sys.exit(1)\n" - - # Add the environment variables that were given to us - script += "env = dict(os.environ)\n" - for e in env: - key, *value = e.split("=", maxsplit=1) - value = shlex.quote(value[0]) if len(value) > 0 else "" - if not all(c.isalnum() or c == "_" for c in key): - log_warning(f"'{e}' is an invalid environment variable so it is ignored") - continue - script += f"env[{repr(key)}] = {repr(value)}\n" - - # Add the environment variables to enable the ring distributed backend - if hostfile != "": - script += "_, hostfile = tempfile.mkstemp()\n" - script += "with open(hostfile, 'w') as f:\n" - script += f" f.write({repr(hostfile)})\n" - if verbose: - script += "env['MLX_RING_VERBOSE'] = '1'\n" - script += "env['MLX_HOSTFILE'] = hostfile\n" - script += f"env['MLX_RANK'] = '{rank}'\n" - script += "\n" - - # Replace the process with the script - script += f"command = [{','.join(map(repr, command))}]\n" - script += "os.execve(command[0], command, env)\n" - - return script - - -def launch_ring(parser, hosts, args, command): - stop = False - exit_codes = [None] * len(hosts) - - def node_thread(rank, host, hostfile, input_queue): - is_local = host == "127.0.0.1" - script = make_monitor_script( - rank, hostfile, args.cwd, args.env, command, args.verbose - ) - script_b64 = base64.b64encode(script.encode()).decode() - cmd = f'{sys.executable} -c "import base64; exec(base64.b64decode(\\"{script_b64}\\"));"' - if not is_local: - cmd = f"ssh {host} '{cmd}'" - p = Popen( - cmd, - shell=True, - stdin=PIPE, - stdout=PIPE, - stderr=PIPE, - ) - os.set_blocking(p.stdout.fileno(), False) - os.set_blocking(p.stderr.fileno(), False) - os.set_blocking(p.stdin.fileno(), False) - - # Repeat the stdout and stderr to the local machine - to_read = [p.stdout.fileno(), p.stderr.fileno()] - to_write = [p.stdin.fileno(), sys.stdout.fileno(), sys.stderr.fileno()] - pidfile = "" - stdin_buffer = b"" - stdout_buffer = b"" - stderr_buffer = b"" - while p.poll() is None: - try: - stdin_buffer += input_queue.get_nowait() - except QueueEmpty: - pass - rlist, wlist, _ = select(to_read, to_write, [], 1.0) - for fd in rlist: - msg = os.read(fd, 8192).decode(errors="ignore") - - # Fetch the PID file first if we haven't already - if pidfile == "": - pidfile, *msg = msg.split("\n", maxsplit=1) - msg = msg[0] if msg else "" - - is_stdout = fd == p.stdout.fileno() - if is_stdout: - stdout_buffer += msg.encode() - else: - stderr_buffer += msg.encode() - for fd in wlist: - if fd == p.stdin.fileno() and len(stdin_buffer) > 0: - n = os.write(fd, stdin_buffer) - stdin_buffer = stdin_buffer[n:] - elif fd == sys.stdout.fileno() and len(stdout_buffer) > 0: - n = os.write(fd, stdout_buffer) - stdout_buffer = stdout_buffer[n:] - elif fd == sys.stderr.fileno() and len(stderr_buffer) > 0: - n = os.write(fd, stderr_buffer) - stderr_buffer = stderr_buffer[n:] - if stop: - p.terminate() - break - p.wait() - exit_codes[rank] = p.returncode - - # Kill the remote program if possible - cmd = "" - cmd += f"pid=$(cat {pidfile}); " - cmd += "if ps -p $pid >/dev/null; then " - cmd += " kill $pid; " - cmd += " echo 1; " - cmd += "else " - cmd += " echo 0; " - cmd += "fi; " - cmd += f"rm {pidfile}" - if not is_local: - cmd = f"ssh {host} '{cmd}'" - c = run(cmd, check=True, shell=True, capture_output=True, text=True) - if c.stdout.strip() == "1": - log_warning(f"Node with rank {rank} was killed") - elif p.returncode != 0: - log_warning(f"Node with rank {rank} exited with code {p.returncode}") - else: - log(args.verbose, f"Node with rank {rank} completed") - - if all(len(h.ips) == 0 for h in hosts): - parser.error( - "The ring backend requires IPs to be provided instead of hostnames" - ) - - port = args.starting_port - ring_hosts = [] - for h in hosts: - node = [] - for ip in h.ips: - for i in range(args.connections_per_ip): - node.append(f"{ip}:{port}") - port += 1 - ring_hosts.append(node) - hostfile = json.dumps(ring_hosts) if len(ring_hosts) > 1 else "" - - log(args.verbose, "Running", shlex.join(command)) - - input_queues = [] - threads = [] - for i, h in enumerate(hosts): - if i + 1 == len(hosts): - time.sleep(1.0) - input_queues.append(Queue()) - t = threading.Thread( - target=node_thread, args=(i, h.ssh_hostname, hostfile, input_queues[-1]) - ) - t.start() - threads.append(t) - - os.set_blocking(sys.stdin.fileno(), False) - while not stop: - rlist, _, _ = select([sys.stdin.fileno()], [], [], 1.0) - for fd in rlist: - stdin_buffer = os.read(fd, 8192) - for q in input_queues: - q.put(stdin_buffer) - if any(t.is_alive() for t in threads): - for i, t in enumerate(threads): - if not t.is_alive(): - if exit_codes[i] != 0: - stop = True - break - else: - break - for t in threads: - t.join() - - -def get_mpi_libname(): - try: - ompi_info = run(["which", "ompi_info"], check=True, capture_output=True) - ompi_info = ompi_info.stdout.strip().decode() - - if platform.system() == "Darwin": - otool_output = run( - ["otool", "-L", ompi_info], check=True, capture_output=True - ) - else: - otool_output = run(["ldd", ompi_info], check=True, capture_output=True) - otool_output = otool_output.stdout.decode() - - # StopIteration if not found - libmpi_line = next( - filter(lambda line: "libmpi" in line, otool_output.splitlines()) - ) - return libmpi_line.strip().split()[0].removeprefix("@rpath/") - except: - return None - - -def launch_mpi(parser, hosts, args, command): - mpirun = run(["which", "mpirun"], check=True, capture_output=True) - mpirun = mpirun.stdout.strip().decode() - - # Compatibility with homebrew and pip installs - mpi_libname = get_mpi_libname() - if mpi_libname is not None: - dyld = Path(mpirun).parent.parent / "lib" - args.env = [ - f"DYLD_LIBRARY_PATH={str(dyld)}", - f"MLX_MPI_LIBNAME={mpi_libname}", - ] + args.env - - log(args.verbose, f"Using '{mpirun}'") - with tempfile.NamedTemporaryFile(mode="w") as f: - hosts = Counter((h.ssh_hostname for h in hosts)) - for h, n in hosts.items(): - print(f"{h} slots={n}", file=f) - f.flush() - - cmd = [ - mpirun, - "--output", - ":raw", # do not line buffer output - "--hostfile", - f.name, - *(["-cwd", args.cwd] if args.cwd else []), - *sum((["-x", e] for e in args.env), []), - *sum([shlex.split(arg) for arg in args.mpi_arg], []), - "--", - *command, - ] - log(args.verbose, "Running", " ".join(cmd)) - try: - run(cmd) - except KeyboardInterrupt: - pass - - -def launch_nccl(parser, hosts, args, command): - master_host = hosts[0].ips[0] - - if master_host != "127.0.0.1": - raise ValueError("The NCCL backend only supports localhost for now.") - master_port = args.nccl_port - world_size = len(hosts) - - base_env = os.environ.copy() - base_env.update( - { - "NCCL_DEBUG": base_env.get( - "NCCL_DEBUG", "INFO" if args.verbose else "DEBUG" - ), - "NCCL_SOCKET_IFNAME": "lo", # Use loopback for local communication - "NCCL_HOST_IP": master_host, - "NCCL_PORT": str(master_port), - "MLX_WORLD_SIZE": str(world_size), - } - ) - procs = [] - num_gpus = get_num_nvidia_gpus() - if num_gpus == 0: - raise RuntimeError("Cannot run NCCL backend with no GPUs.") - if args.repeat_hosts > num_gpus: - raise RuntimeError("NCCL requires a separate GPU per process.") - - try: - for rank in range(world_size): - env = base_env.copy() - mlx_rank = str(rank % args.repeat_hosts) - env["MLX_RANK"] = mlx_rank - env["CUDA_VISIBLE_DEVICES"] = mlx_rank - p = Popen(command, env=env) - procs.append(p) - - for p in procs: - ret = p.wait() - if ret != 0: - raise RuntimeError(f"Rank process exited with {ret}") - - except (RuntimeError, KeyboardInterrupt) as err: - for p in procs: - if p.poll() is None: - try: - p.kill() - except Exception: - pass - raise - - -def check_ssh_connections(hosts): - results = [False] * len(hosts) - - def _check(hostname, i): - result = run( - [ - "ssh", - "-o", - "BatchMode=yes", - "-o", - "ConnectTimeout=5", - hostname, - "echo", - "success", - ], - stdout=PIPE, - stderr=PIPE, - ) - results[i] = result.returncode == 0 - - threads = [ - threading.Thread(target=_check, args=(h.ssh_hostname, i)) - for i, h in enumerate(hosts) - ] - for t in threads: - t.start() - for t in threads: - t.join() - - if not all(results): - log_error("Could not ssh to the following hosts:") - for i, h in enumerate(hosts): - if not results[i]: - log_error(" - ", h.ssh_hostname) - log_error() - log_error("Maybe they are not set-up for password-less ssh?") - sys.exit(1) - - -def prepare_tb_ring(args, hosts): - log( - args.verbose, - f"Preparing a thunderbolt ring for {', '.join(h.ssh_hostname for h in hosts)}", - ) - - # Check that we can ssh - check_ssh_connections(hosts) - if args.auto_setup and args.verbose: - log_warning( - "--auto-setup is requested which requires password-less sudo", - "on the remote hosts", - ) - - # Extract the current connectivity from the remote hosts - thunderbolt_connections = [] - for h in hosts: - log(args.verbose, "Getting connectivity from", h.ssh_hostname) - thunderbolt_connections.append( - json.loads( - run( - [ - "ssh", - h.ssh_hostname, - "system_profiler", - "SPThunderboltDataType", - "-json", - ], - capture_output=True, - ).stdout - ) - ) - interface_maps = [] - for h in hosts: - log(args.verbose, "Getting interface names from", h.ssh_hostname) - interface_maps.append( - parse_hardware_ports( - run( - [ - "ssh", - h.ssh_hostname, - "networksetup", - "-listallhardwareports", - ], - capture_output=True, - ).stdout - ) - ) - - # Parse the connectivity into some simple dataclasses - tb_hosts = [] - for c, iface_map in zip(thunderbolt_connections, interface_maps): - name = "" - ports = [] - for t in c["SPThunderboltDataType"]: - uuid = t.get("domain_uuid_key") - if uuid is None: - continue - name = t["device_name_key"] - tag = t["receptacle_1_tag"]["receptacle_id_key"] - items = t.get("_items", []) - connected_items = [item for item in items if "domain_uuid_key" in item] - connected_to = ( - connected_items[0]["domain_uuid_key"] if connected_items else None - ) - iface = iface_map[f"Thunderbolt {tag}"] - ports.append(ThunderboltPort(iface, uuid, connected_to)) - tb_hosts.append(ThunderboltHost(name, sorted(ports, key=lambda x: x.iface))) - - # Create a reverse index to be able to map uuids to (host, port) quickly - uuid_reverse_index = {} - for i, h in enumerate(tb_hosts): - for j, p in enumerate(h.ports): - uuid_reverse_index[p.uuid] = (i, j) - - # Find the rings by simply walking and marking visited (host, port) tuples - # and keeping the largest rings greedily. - log(args.verbose, "Extracting rings from the parsed connectivity") - rings = extract_rings(tb_hosts, uuid_reverse_index) - - # Just output a DOT graphical representation of the found rings - if args.dot: - names = [] - for i in range(len(tb_hosts)): - n = "" - j = i - while True: - n += chr(97 + j % 26) - j //= 26 - if j == 0: - break - names.append(n) - - print("graph G {") - print(" node [shape=rectangle];") - for i, h in enumerate(hosts): - print(f' {names[i]} [label="{h.ssh_hostname}"];') - for r in rings: - for (i, _), (j, _) in r: - print(f" {names[i]} -- {names[j]};") - print("}") - return - - # Assign IPs to each interface such that the interfaces can communicate - ips = {} - pairs = {} - expecting = set() - ip0 = 0 - ip1 = 0 - netmask = "255.255.255.252" - for r in rings: - for a, b in r: - ips[a] = f"192.168.{ip0}.{ip1 + 1}" - ips[b] = f"192.168.{ip0}.{ip1 + 2}" - pairs[a] = b - pairs[b] = a - expecting.add(b) - ip1 += 4 - if ip1 > 255: - ip0 += 1 - ip1 = 0 - if ip0 > 255: - raise ValueError("Ran out of available local IPs for the ring") - - # Extract the host order from the first ring - hostmap = dict((r[0][0], r[1][0]) for r in rings[0]) - first_host = min(hostmap.keys()) - order = [first_host] - while hostmap[order[-1]] != first_host: - order.append(hostmap[order[-1]]) - - # Create the hostfile - hostfile = [] - for i in order: - h = hosts[i] - host = { - "ssh": h.ssh_hostname, - "ips": [ - ips[i, j] - for j, p in enumerate(tb_hosts[i].ports) - if (i, j) in expecting - ], - } - hostfile.append(host) - - if not args.hostfile_only: - for i, h in enumerate(hosts): - command = "" - command += "sudo ifconfig bridge0 down\n" - for j, p in enumerate(tb_hosts[i].ports): - if (i, j) not in ips: - continue - iface = p.iface - ip = ips[i, j] - peer = ips[pairs[i, j]] - command += f"sudo ifconfig {iface} inet {ip} netmask {netmask}\n" - command += f"sudo route change {peer} -interface {iface}\n" - if args.auto_setup: - print(f"Running auto setup for {h.ssh_hostname}") - command = command.strip().replace("\n", " && ") - command = ["ssh", h.ssh_hostname, command] - log(args.verbose, shlex.join(command)) - run(command) - else: - msg = f"Setup for {h.ssh_hostname}" - print(msg) - print("=" * len(msg)) - print(command) - input("Enter to continue") - print() - - if args.output_hostfile: - with open(args.output_hostfile, "w") as f: - json.dump(hostfile, f, indent=4) - else: - print("Hostfile") - print("========") - print(json.dumps(hostfile, indent=4)) - - -def prepare_hostfile(args, hosts): - log( - args.verbose, - f"Preparing an ethernet hostfile for {', '.join(h.ssh_hostname for h in hosts)}", - ) - - # Check that we can ssh - check_ssh_connections(hosts) - - # Get the ips for each host - for h in hosts: - log(args.verbose, "Getting the ip from", h.ssh_hostname) - h.ips.append( - run( - ["ssh", h.ssh_hostname, "ipconfig", "getifaddr", "en0"], - capture_output=True, - text=True, - ).stdout.strip() - ) - - hostfile = [] - for h in hosts: - hostfile.append(dict(ssh=h.ssh_hostname, ips=h.ips)) - - if args.output_hostfile: - with open(args.output_hostfile, "w") as f: - json.dump(hostfile, f, indent=4) - else: - print("Hostfile") - print("========") - print(json.dumps(hostfile, indent=4)) - - -def distributed_config(): - parser = argparse.ArgumentParser( - description="Configure remote machines for use with MLX distributed" - ) - parser.add_argument( - "--verbose", action="store_true", help="Print debug messages in stdout" - ) - parser.add_argument( - "--backend", - choices=["ring", "mpi", "nccl"], - default="nccl" if mx.cuda.is_available() else "ring", - help="Which distributed backend to configure", - ) - parser.add_argument( - "--over", - choices=["thunderbolt", "ethernet"], - default="thunderbolt", - help="What type of connectivity to configure", - ) - parser.add_argument( - "--hosts", default="127.0.0.1", help="A comma separated list of hosts" - ) - parser.add_argument("--hostfile", help="The file containing the hosts") - parser.add_argument( - "--dot", action="store_true", help="Output the topology in DOT format and exit" - ) - parser.add_argument( - "--hostfile-only", action="store_true", help="If set only compute the hostfile" - ) - parser.add_argument( - "--output-hostfile", help="If provided, save the hostfile to this path" - ) - parser.add_argument( - "--auto-setup", - action="store_true", - help="If set we will attempt to automatically configure the machines via ssh", - ) - args = parser.parse_args() - - if args.backend == "mpi" and args.over == "thunderbolt": - raise ValueError( - ( - "The configuration of MPI over thunderbolt is " - "not supported yet by mlx.distributed_config" - ) - ) - - if args.hostfile is not None: - hosts = parse_hostfile(parser, args.hostfile) - else: - hosts = parse_hostlist(parser, args.hosts, 1) - - if args.over == "thunderbolt": - prepare_tb_ring(args, hosts) - else: - prepare_hostfile(args, hosts) - - -def main(): - parser = argparse.ArgumentParser(description="Launch an MLX distributed program") - parser.add_argument( - "--print-python", - action="store_true", - help="Print the path to the current python executable and exit", - ) - parser.add_argument( - "--verbose", action="store_true", help="Print debug messages in stdout" - ) - parser.add_argument( - "--hosts", default="127.0.0.1", help="A comma separated list of hosts" - ) - parser.add_argument( - "--repeat-hosts", - "-n", - type=positive_number, - default=1, - help="Repeat each host a given number of times", - ) - parser.add_argument("--hostfile", help="The file containing the hosts") - parser.add_argument( - "--backend", - choices=["ring", "mpi", "nccl"], - default="nccl" if mx.cuda.is_available() else "ring", - help="Which distributed backend to launch", - ) - parser.add_argument( - "--env", - action="append", - default=[], - help="Set environment variables for the jobs", - ) - parser.add_argument( - "--mpi-arg", - action="append", - default=[], - help="Arguments to pass directly to mpirun", - ) - parser.add_argument( - "--connections-per-ip", - default=1, - type=int, - help="How many connections per ip to use for the ring backend", - ) - parser.add_argument( - "--starting-port", - "-p", - type=int, - default=5000, - help="For the ring backend listen on this port increasing by 1 per rank and IP", - ) - parser.add_argument( - "--cwd", help="Set the working directory on each node to the provided one" - ) - parser.add_argument( - "--nccl-port", - type=int, - default=12345, - help="The port to use for the NCCL communication (only for nccl backend)", - ) - - args, rest = parser.parse_known_args() - - if args.print_python: - print(sys.executable) - return - - if len(rest) == 0: - parser.error("No script is provided") - if rest[0] == "--": - rest.pop(0) - - # Try to extract a list of hosts and corresponding ips - if args.hostfile is not None: - hosts = parse_hostfile(parser, args.hostfile) - else: - hosts = parse_hostlist(parser, args.hosts, args.repeat_hosts) - - # Check if the script is a file and convert it to a full path - if (script := Path(rest[0])).exists(): - rest[0:1] = [sys.executable, str(script.resolve())] - elif (command := shutil.which(rest[0])) is not None: - rest[0] = command - else: - raise ValueError(f"Invalid script or command {rest[0]}") - - # Launch - if args.backend == "ring": - launch_ring(parser, hosts, args, rest) - if args.backend == "mpi": - launch_mpi(parser, hosts, args, rest) - if args.backend == "nccl": - launch_nccl(parser, hosts, args, rest) - - -if __name__ == "__main__": - main() diff --git a/python/src/distributed.cpp b/python/src/distributed.cpp index d147c278..9f4a7cb5 100644 --- a/python/src/distributed.cpp +++ b/python/src/distributed.cpp @@ -52,9 +52,25 @@ void init_distributed(nb::module_& parent_module) { m.def( "is_available", - &mx::distributed::is_available, + [](const std::string& backend) { + return mx::distributed::is_available(backend); + }, + "backend"_a = "any", + nb::sig("def is_available(backend: str = 'any') -> bool"), R"pbdoc( Check if a communication backend is available. + + Note, this function returns whether MLX has the capability of + instantiating that distributed backend not whether it is possible to + create a communication group. For that purpose one should use + ``init(strict=True)``. + + Args: + backend (str, optional): The name of the backend to check for availability. + It takes the same values as :func:`init()`. Default: ``"any"``. + + Returns: + bool: Whether the distributed backend is available. )pbdoc"); m.def( @@ -79,10 +95,10 @@ void init_distributed(nb::module_& parent_module) { in case ``mx.distributed.is_available()`` returns False otherwise it throws a runtime error. Default: ``False`` backend (str, optional): Which distributed backend to initialize. - Possible values ``mpi``, ``ring``, ``nccl``, ``any``. If set to ``any`` all - available backends are tried and the first one that succeeds - becomes the global group which will be returned in subsequent - calls. Default: ``any`` + Possible values ``mpi``, ``ring``, ``nccl``, ``jaccl``, ``any``. If + set to ``any`` all available backends are tried and the first one + that succeeds becomes the global group which will be returned in + subsequent calls. Default: ``any`` Returns: Group: The group representing all the launched processes. diff --git a/setup.py b/setup.py index 1678e13a..0439009a 100644 --- a/setup.py +++ b/setup.py @@ -265,8 +265,8 @@ if __name__ == "__main__": } entry_points = { "console_scripts": [ - "mlx.launch = mlx.distributed_run:main", - "mlx.distributed_config = mlx.distributed_run:distributed_config", + "mlx.launch = mlx._distributed_utils.launch:main", + "mlx.distributed_config = mlx._distributed_utils.config:main", ] } install_requires = []