|
|
.. _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. |
|
|
|
