Skip to main content

Crate nuvec_distributed

Crate nuvec_distributed 

Source
Expand description

Distributed backends for nuvec: MLX’s TCP ring.

init reads MLX’s launch environment (MLX_HOSTFILE, a JSON list with one entry per rank of "host:port" addresses, and MLX_RANK), connects the ring and makes it the default group of nuvec::distributed’s collectives. Every rank runs the same program; the processes may share one Mac (over loopback) or run on several (over Ethernet or Thunderbolt IP).

use nuvec_core::{array, distributed};

let group = nuvec_distributed::init(false)?; // size 1 without a hostfile
let x = array![1.0, 2.0];
let total = distributed::all_sum(&x, None); // the sum over every rank
println!("rank {} of {}: {:?}", group.rank(), group.size(), total.to_vec::<f32>()?);

Stability: no promise when this crate is depended on directly; depend on nuvec, whose re-exports are semver-tracked.

Structs§

Options
Options of init_with.

Functions§

init
Initialises the ring from MLX’s environment (MLX_HOSTFILE, MLX_RANK) and makes it the default group, as mx.distributed.init(). Without the environment: an error if strict, else the group of this process alone.
init_with
Connects rank rank of the ring over nodes (per rank its addresses, one connection per address, as MLX’s hostfile) and returns its group (not made the default). A single rank is the group of this process alone.