Module distributed 

Distributed groups and collective operations, as MLX’s mx.distributed.

A Group is a set of processes with a rank and a size. Collectives (all_sum, all_max, all_min, all_gather, sum_scatter, send, recv) are lazy array operations over a group: they run when their result is evaluated, compose with grad, vmap and compile as MLX defines, and run on the CPU over the arrays’ shared memory. On a GPU stream they run on the group’s worker thread once the GPU work producing their inputs completes, and the GPU work reading their outputs is committed once they are done, so neither the evaluating thread nor the GPU waits for the network. Without a configured group, the default group has size 1 and every collective is the identity, so single-process code runs unchanged.

The transport is a Communicator; nuvec-distributed provides MLX’s TCP ring and JACCL backends and installs one as the default group.

Structs

Group

A group of processes (MLX’s Group): cheap to clone.

Enums

AllReduce

The reduction of Communicator::all_reduce.

CollectiveOp

A collective operation’s kind, with its parameters (Primitive::Collective).

Traits

Communicator

The transport of a Group: every rank of the group calls the same operations in the same order. Buffers are row-contiguous elements; a failure names the peer.

Functions

all_gather

Every rank’s x concatenated along the first axis in rank order (a scalar gathers to a vector), as MLX’s all_gather.

all_max

The element-wise maximum of x over the ranks of group (MLX’s all_max).

all_min

The element-wise minimum of x over the ranks of group (MLX’s all_min).

all_sum

The element-wise sum of x over the ranks of group (the default group if None), as MLX’s all_sum: every rank gets the same bits.

default_group

The default group: the one set by set_default_group, else this process alone.

recv

An array of shape and dtype received from rank src of group, as MLX’s recv. A sender’s other dtype or shape is an error when it is evaluated.

recv_like

recv of x’s shape and dtype (MLX’s recv_like).

send

Sends x to rank dst of group when the result (which is x) is evaluated, as MLX’s send. MLX’s ring backend sends to its neighbours only.

set_default_group

Makes group the default group of the collectives (what nuvec-distributed’s init does).

sum_gradients

x unchanged, whose gradient is summed over group (MLX’s sum_gradients): the input of a layer whose output is sharded across the group, so each rank’s gradient covers every rank’s slice of the output.

sum_scatter

Rank r gets the sum over group of every rank’s rows r·k..(r+1)·k of x, whose first axis is size·k, as MLX’s sum_scatter (the reduce-scatter half of an all-reduce). A group of size 1 returns x.