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, 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 wait for the GPU work producing their inputs). 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 installs it 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.