| .. _usage_distributed: |
|
|
| Distributed Communication |
| ========================= |
|
|
| .. currentmodule:: mlx.core.distributed |
|
|
| 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 two different communication backends: |
|
|
| * `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ a |
| full-featured and mature distributed communications library |
| * A **ring** backend of our own that uses native TCP sockets and should be |
| faster for thunderbolt connections. |
|
|
| The list of all currently supported operations and their documentation can be |
| seen in the :ref:`API docs<distributed>`. |
|
|
| .. 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 |
| --------------- |
|
|
| A distributed program in MLX is as simple as: |
|
|
| .. code:: python |
|
|
| import mlx.core as mx |
|
|
| world = mx.distributed.init() |
| x = mx.distributed.all_sum(mx.ones(10)) |
| print(world.rank(), x) |
|
|
| The program above sums the array ``mx.ones(10)`` across all |
| distributed processes. However, when this script is run with ``python`` only |
| one process is launched and no distributed communication takes place. Namely, |
| all operations in ``mx.distributed`` are noops when the distributed group has a |
| size of one. This property allows us to avoid code that checks if we are in a |
| distributed setting similar to the one below: |
|
|
| .. code:: python |
|
|
| import mlx.core as mx |
|
|
| x = ... |
| world = mx.distributed.init() |
| # No need for the check we can simply do x = mx.distributed.all_sum(x) |
| if world.size() > 1: |
| x = mx.distributed.all_sum(x) |
|
|
| Running Distributed Programs |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ |
|
|
| MLX provides ``mlx.launch`` a helper script to launch distributed programs. |
| Continuing with our initial example we can run it on localhost with 4 processes using |
|
|
| .. code:: shell |
|
|
| $ mlx.launch -n 4 my_script.py |
| 3 array([4, 4, 4, ..., 4, 4, 4], dtype=float32) |
| 2 array([4, 4, 4, ..., 4, 4, 4], dtype=float32) |
| 1 array([4, 4, 4, ..., 4, 4, 4], dtype=float32) |
| 0 array([4, 4, 4, ..., 4, 4, 4], dtype=float32) |
|
|
| We can also run it on some remote hosts by providing their IPs (provided that |
| the script exists on all hosts and they are reachable by ssh) |
|
|
| .. code:: shell |
|
|
| $ mlx.launch --hosts ip1,ip2,ip3,ip4 my_script.py |
| 3 array([4, 4, 4, ..., 4, 4, 4], dtype=float32) |
| 2 array([4, 4, 4, ..., 4, 4, 4], dtype=float32) |
| 1 array([4, 4, 4, ..., 4, 4, 4], dtype=float32) |
| 0 array([4, 4, 4, ..., 4, 4, 4], dtype=float32) |
|
|
| Consult the dedicated :doc:`usage guide<launching_distributed>` for more |
| information on using ``mlx.launch``. |
|
|
| Selecting Backend |
| ^^^^^^^^^^^^^^^^^ |
|
|
| You can select the backend you want to use when calling :func:`init` by passing |
| one of ``{'any', 'ring', 'mpi'}``. When passing ``any``, MLX will try to |
| initialize the ``ring`` backend and if it fails the ``mpi`` backend. If they |
| both fail then a singleton group is created. |
| |
| .. note:: |
| After a distributed backend is successfully initialized :func:`init` will |
| return **the same backend** if called without arguments or with backend set to |
| ``any``. |
|
|
| The following examples aim to clarify the backend initialization logic in MLX: |
|
|
| .. code:: python |
|
|
| # Case 1: Initialize MPI regardless if it was possible to initialize the ring backend |
| world = mx.distributed.init(backend="mpi") |
| world2 = mx.distributed.init() # subsequent calls return the MPI backend! |
|
|
| # Case 2: Initialize any backend |
| world = mx.distributed.init(backend="any") # equivalent to no arguments |
| world2 = mx.distributed.init() # same as above |
|
|
| # Case 3: Initialize both backends at the same time |
| world_mpi = mx.distributed.init(backend="mpi") |
| world_ring = mx.distributed.init(backend="ring") |
| world_any = mx.distributed.init() # same as MPI because it was initialized first! |
|
|
| Training Example |
| ---------------- |
|
|
| In this section we will adapt an MLX training loop to support data parallel |
| distributed training. Namely, we will average the gradients across a set of |
| hosts before applying them to the model. |
|
|
| Our training loop looks like the following code snippet if we omit the model, |
| dataset and optimizer initialization. |
|
|
| .. code:: python |
|
|
| model = ... |
| optimizer = ... |
| dataset = ... |
|
|
| def step(model, x, y): |
| loss, grads = loss_grad_fn(model, x, y) |
| optimizer.update(model, grads) |
| return loss |
|
|
| for x, y in dataset: |
| loss = step(model, x, y) |
| mx.eval(loss, model.parameters()) |
|
|
| All we have to do to average the gradients across machines is perform an |
| :func:`all_sum` and divide by the size of the :class:`Group`. Namely we |
| have to :func:`mlx.utils.tree_map` the gradients with following function. |
|
|
| .. code:: python |
|
|
| def all_avg(x): |
| return mx.distributed.all_sum(x) / mx.distributed.init().size() |
|
|
| Putting everything together our training loop step looks as follows with |
| everything else remaining the same. |
|
|
| .. code:: python |
|
|
| from mlx.utils import tree_map |
|
|
| def all_reduce_grads(grads): |
| N = mx.distributed.init().size() |
| if N == 1: |
| return grads |
| return tree_map( |
| lambda x: mx.distributed.all_sum(x) / N, |
| grads |
| ) |
|
|
| def step(model, x, y): |
| loss, grads = loss_grad_fn(model, x, y) |
| grads = all_reduce_grads(grads) # <--- This line was added |
| optimizer.update(model, grads) |
| return loss |
|
|
| Utilizing ``nn.average_gradients`` |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ |
|
|
| Although the code example above works correctly; it performs one communication |
| per gradient. It is significantly more efficient to aggregate several gradients |
| together and perform fewer communication steps. |
|
|
| This is the purpose of :func:`mlx.nn.average_gradients`. The final code looks |
| almost identical to the example above: |
|
|
| .. code:: python |
|
|
| model = ... |
| optimizer = ... |
| dataset = ... |
|
|
| def step(model, x, y): |
| loss, grads = loss_grad_fn(model, x, y) |
| grads = mx.nn.average_gradients(grads) # <---- This line was added |
| optimizer.update(model, grads) |
| return loss |
|
|
| for x, y in dataset: |
| loss = step(model, x, y) |
| mx.eval(loss, model.parameters()) |
|
|
|
|
| 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. |
|
|
| The simplest possible usage is the following which, assuming the minimal |
| example in the beginning of this page, should result in: |
|
|
| .. code:: shell |
|
|
| $ mlx.launch --backend mpi -n 2 test.py |
| 1 array([2, 2, 2, ..., 2, 2, 2], dtype=float32) |
| 0 array([2, 2, 2, ..., 2, 2, 2], dtype=float32) |
|
|
| The above launches two processes on the same (local) machine and we can see |
| both standard output streams. The processes send the array of 1s to each other |
| and compute the sum which is printed. Launching with ``mlx.launch -n 4 ...`` would |
| print 4 etc. |
|
|
| Installing MPI |
| ^^^^^^^^^^^^^^ |
|
|
| MPI can be installed with Homebrew, using the Anaconda package manager or |
| compiled from source. Most of our testing is done using ``openmpi`` installed |
| with the Anaconda package manager as follows: |
|
|
| .. code:: shell |
|
|
| $ conda install conda-forge::openmpi |
|
|
| Installing with Homebrew may require specifying the location of ``libmpi.dyld`` |
| so that MLX can find it and load it at runtime. This can simply be achieved by |
| passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun`` and it is |
| done automatically by ``mlx.launch``. |
|
|
| .. code:: shell |
|
|
| $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python test.py |
| $ # or simply |
| $ mlx.launch -n 2 test.py |
|
|
| Setting up Remote Hosts |
| ^^^^^^^^^^^^^^^^^^^^^^^ |
|
|
| MPI can automatically connect to remote hosts and set up the communication over |
| the network if the remote hosts can be accessed via ssh. A good checklist to |
| debug connectivity issues is the following: |
|
|
| * ``ssh hostname`` works from all machines to all machines without asking for |
| password or host confirmation |
| * ``mpirun`` is accessible on all machines. |
| * Ensure that the ``hostname`` used by MPI is the one that you have configured |
| in the ``.ssh/config`` files on all machines. |
|
|
| Tuning MPI All Reduce |
| ^^^^^^^^^^^^^^^^^^^^^ |
|
|
| .. note:: |
|
|
| For faster all reduce consider using the ring backend either with Thunderbolt |
| connections or over Ethernet. |
|
|
| Configure MPI to use N tcp connections between each host to improve bandwidth |
| by passing ``--mca btl_tcp_links N``. |
|
|
| Force MPI to use the most performant network interface by setting ``--mca |
| btl_tcp_if_include <iface>`` where ``<iface>`` should be the interface you want |
| to use. |
|
|
| 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 is 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 <launching_distributed>`. 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 |
|
|
| 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. |
|
|
| To validate your connection without configuring anything |
| ``mlx.distributed_config`` can also plot the ring using DOT format. |
|
|
| .. code:: shell |
|
|
| mlx.distributed_config --verbose --hosts host1,host2,host3,host4 --dot >ring.dot |
| dot -Tpng ring.dot >ring.png |
| open ring.png |
|
|
| 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. |
|
|
