Distributed Communication

Distributed Communication#

MLX provides a distributed communication package using MPI. The MPI library is loaded at runtime; if MPI is available then distributed communication is also made available.

Group

An mlx.core.distributed.Group represents a group of independent mlx processes that can communicate.

is_available()

Check if a communication backend is available.

init([strict])

Initialize the communication backend and create the global communication group.

all_sum(x, *[, group, stream])

All reduce sum.

all_gather(x, *[, group, stream])

Gather arrays from all processes.

send(x, dst, *[, group, stream])

Send an array from the current process to the process that has rank dst in the group.

recv(shape, dtype, src, *[, group, stream])

Recv an array with shape shape and dtype dtype from process with rank src.

recv_like(x, src, *[, group, stream])

Recv an array with shape and type like x from process with rank src.