Module dist 

Ship-graph distributed execution (build the worker once, run any model). Ship-graph distributed execution — build the worker once, run any model.

A worker is a generic executor: it compiles and runs a graph the coordinator ships at runtime, so no model is baked into the worker binary and there is no per-node / per-model recompile. Build one worker binary per architecture (Node::from_env().connect() + [serve_stage]), deploy it to every node, and drive any model from a coordinator.

The coordinator — which owns the model’s graph builders (e.g. a crate in rlx-models) — partitions the model into stages and ships each worker a [StageSpec]: the serialized subgraph, its I/O node names, where to fetch each weight, and a device directive. Weights are resolved locally by a caller-provided closure (GGUF / safetensors / HF — rlx core stays model-agnostic), so only specs (KB) and activations cross the wire — never the weights.

// Generic worker binary — the SAME build runs any model:
let mut stage = dist::recv_stage(&group, |uri| load_weights(uri))?;
let x = dist::recv_activation(&group, group.rank() - 1)?;
let y = stage.run(&x);
dist::send_activation(&group, /*next*/ 0, &y)?;

Structs

BackendDivergence

One backend’s worst divergence from the CPU oracle for a given graph run.

DataRef

One data tensor the worker feeds to a graph input every step. Resolved node-locally from uri (a file:// raw-f32 blob, or gguf/safetensors), then row-sharded: this rank owns samples [shard_start, shard_start+len), each elem floats wide.

StageSpec

A self-describing unit of work for a worker: everything needed to run one stage of a model, with nothing implied by convention.

TrainMetrics

Per-worker training result.

TrainSpec

A self-describing training job for a generic worker: everything needed to train a model the worker has no code for.

WeightCache

Parse-once, serve-many weight cache. A GGUF file is parsed a single time and then serves any number of its tensors (f32 or packed) — the right shape for a worker whose stage pulls many tensors (q/k/v/o/ffn) from one model file. Created per stage inside recv_stage; also usable standalone.

WeightRef

Where one parameter’s data comes from. The worker resolves the uri itself, so large weights stay node-local.

WorkerStage

A compiled stage sitting on this worker, ready to run activations.

Functions

backend_divergence

Run graph on the CPU oracle and on every other backend live on this node, reporting each one’s worst divergence from CPU. Call it on each cluster node to localize where distributed numbers drift — and whether that drift is bounded round-off (suspect == false) or a candidate kernel bug (suspect == true).

pull_shards

Worker: receive the pushed init params + data shards from the master (rank 0), write them to node-local files under dir, and rewrite spec.params/spec.data to point at those files (data offset reset to 0). After this the generic run_train resolves everything locally exactly as in the shared-FS case.

push_shards

Master: push worker rank everything it needs but can’t resolve locally — the init params (training-from-scratch has no node-local weights, so the WeightRef file:// URIs only exist on the master) and each data shard (only this rank’s samples). Use when the worker has no shared filesystem. Pair with pull_shards on the worker.

recv_activation

Receive an activation tensor from from.

recv_stage

Worker: receive a StageSpec from the coordinator (rank 0), resolve its weights via resolve (uri -> f32), pick the directed device, and compile. This is the generic worker’s setup — the same code runs any model’s stage.

recv_train

Worker: receive a TrainSpec from the coordinator (rank 0).

report_backend_divergence

Print a backend_divergence report and return true if ANY backend looks suspect (relative error past tol_rel) — a convenient one-liner for a CLI / CI gate on each node. A green report means the cluster’s cross-backend drift is bounded round-off; a red one names the backend + output to investigate.

resolve_weight_bytes

Load the quantized bytes of a weight as-is (no dequant), for the on-device-quant path. gguf://<path>#<tensor> returns the packed block bytes; file://<path> returns the raw file. Feed to set_param_typed(_, _, U8) behind a DequantMatMul graph.

resolve_weight_uri

Resolve a weight tensor to f32 from its source URI. Supported schemes: gguf://<path>#<tensor> — dequantize a GGUF tensor (any K-quant) safetensors://<path>#<tensor> — read a safetensors tensor (F32/F16/BF16) file://<path> — raw little-endian f32 array

run_train

Run one worker’s full data-parallel training from a shipped TrainSpec — the generic worker’s whole job, with no model code.

send_activation

Send an activation tensor to to (a worker, or the next pipeline stage).

serve_stage

Worker convenience: set up the stage, then run one activation cycle — receive from the previous rank (or an external feed for rank 1), run, and forward to the next rank (or back to the coordinator if last). Returns the device it ran on. For a serving loop, call recv_stage once and drive WorkerStage::run + recv_activation/send_activation yourself.

serve_stage_uri

Same as serve_stage but with no custom resolver — a worker binary that serves GGUF / safetensors / raw-f32 models out of the box (recv_stage’s built-in WeightCache handles those schemes; anything else warns).

ship_stage

Coordinator: ship a stage spec to worker rank.

ship_train

Coordinator: ship a training job to worker rank.

uri_resolver

Built-in resolver for serve_stage_uri: resolve_weight_uri, logging and returning empty on failure so a missing shard surfaces loudly rather than panicking mid-collective.