Ray collective library is bundled with the released Ray wheel. Besides Ray, users need to install either `pygloo <https://github.com/ray-project/pygloo>`_
or `cupy <https://docs.cupy.dev/en/stable/install.html>`_ in order to use collective communication with the GLOO and NCCL backend, respectively.
..code-block:: python
pip install pygloo
pip install cupy-cudaxxx # replace xxx with the right cuda version in your environment
To use these APIs, import the collective package in your actor/task or driver code via:
..code-block:: python
import ray.util.collective as col
Initialization
^^^^^^^^^^^^^^
Collective functions operate on collective groups.
A collective group contains a number of processes (in Ray, they are usually Ray-managed actors or tasks) that will together enter the collective function calls.
Before making collective calls, users need to declare a set of actors/tasks, statically, as a collective group.
Below is an example code snippet that uses the two APIs ``init_collective_group()`` and ``declare_collective_group()`` to initialize collective groups among a few
remote actors. Refer to `APIs <#api-reference>`_ for the detailed descriptions of the two APIs.
results = ray.get([w.compute.remote() for w in workers])
Note that for the same set of actors/task processes, multiple collective groups can be constructed, with ``group_name`` as their unique identifier.
This enables to specify complex communication patterns between different (sub)set of processes.
Collective Communication
^^^^^^^^^^^^^^^^^^^^^^^^
Check `the support matrix <#collective-primitives-support-matrix>`_ for the current status of supported collective calls and backends.
Note that the current set of collective communication API are imperative, and exhibit the following behaviours:
* All the collective APIs are synchronous blocking calls
* Since each API only specifies a part of the collective communication, the API is expected to be called by each participating process of the (pre-declared) collective group.
Upon all the processes have made the call and rendezvous with each other, the collective communication happens and proceeds.
* The APIs are imperative and the communication happends out-of-band --- they need to be used inside the collective process (actor/task) code.
An example of using ``ray.util.collective.allreduce`` is below:
# Create two actors A and B and create a collective group following the previous example...
A = Worker.remote()
B = Worker.remote()
# Invoke allreduce remotely
ray.get([A.compute.remote(), B.compute.remote()])
Point-to-point Communication
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
``ray.util.collective`` also supports P2P send/recv communication between processes.
The send/recv exhibits the same behavior with the collective functions:
they are synchronous blocking calls -- a pair of send and recv must be called together on paired processes in order to specify the entire communication,
and must successfully rendezvous with each other to proceed. See the code example below:
# An anti-pattern: the following code will hang, because it does instantiate the recv side call
ray.get([A.do_send.remote(target_rank=1)])
Single-GPU and Multi-GPU Collective Primitives
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
In many cluster setups, a machine usually has more than 1 GPU;
effectively leveraging the GPU-GPU bandwidth, such as `NVLINK <https://www.nvidia.com/en-us/design-visualization/nvlink-bridges/>`_\ ,
can significantly improve communication performance.
``ray.util.collective`` supports multi-GPU collective calls, in which case, a process (actor/tasks) manages more than 1 GPU (e.g., via ``ray.remote(num_gpus=4)``\ ).
Using these multi-GPU collective functions are normally more performance-advantageous than using single-GPU collective API
and spawning the number of processes equal to the number of GPUs.
See the API references for the signatures of multi-GPU collective APIs.
Also of note that all multi-GPU APIs are with the following restrictions:
* Only NCCL backend is supported.
* Collective processes that make multi-GPU collective or P2P calls need to own the same number of GPU devices.
* The input to multi-GPU collective functions are normally a list of tensors, each located on a different GPU device owned by the caller process.
An example code utilizing the multi-GPU collective APIs is provided below:
# Note that the world size is 2 but there are 4 GPUs.
num_workers = 2
workers = []
init_rets = []
for i in range(num_workers):
w = Worker.remote()
workers.append(w)
init_rets.append(w.setup.remote(num_workers, i))
a = ray.get(init_rets)
results = ray.get([w.allreduce_call.remote() for w in workers])
results = ray.get([w.p2p_call.remote() for w in workers])
More Resources
--------------
The following links provide helpful resources on how to efficiently leverage the ``ray.util.collective`` library.
*`More running examples <https://github.com/ray-project/ray/tree/master/python/ray/util/collective/examples>`_ under ``ray.util.collective.examples``.
*`Scaling up the Spacy Name Entity Recognition (NER) pipeline <https://github.com/explosion/spacy-ray>`_ using Ray collective library.
*`Implementing the AllReduce strategy <https://github.com/ray-project/distml/blob/master/distml/strategy/allreduce_strategy.py>`_ for data-parallel distributed ML training.
