Meet Dis-Tract 0.1.0

[czoli1976]

Dis-tract splits a model by layer across workers, aiming at running a model too big for one box: each worker loads only its own contiguous slice of layers plus that slice's weights, keeps that slice's KV resident, and only the residual crosses the wire. Opening this as a draft to try and to settle one design question (below) — it is not proposed for merge in this shape, and the too-big-for-one-box payoff is not yet demonstrated (see the next section).

Follow-up to #2465.

It lives in examples/dis-tract, next to causal_llm, and is in members but not default-members — so a bare cargo build does not pull its zenoh tree in.

What has actually been tested — read this before the numbers

Everything below is two worker processes on ONE 16 GB Apple-silicon machine, over loopback. It has never run on two physical machines. That matters more than it sounds:

• The premise is unproven. The point is running a model too big for one box, but Qwen3-8B-q40ef16 (4.29 GB) fits that machine easily. Splitting it 2210/2183 MiB shows the mechanism — each worker loads only its own layers' weights — not the payoff. Nothing here demonstrates a model that could not have run single-node.
• The wire cost is not measured. Loopback makes the per-token residual hop essentially free, so "the split is free" is a localhost result. On a real LAN it buys a round-trip per token per stage boundary; at ~43 ms/step even a 1 ms hop is ~2%, and a slow link would be far worse. Only the residual crosses (never the KV), so the payload is small — but small is not zero.
• Both stages share one GPU, so they serialise on it. On two machines each stage would own its GPU, which cuts the other way; the two effects have not been separated.
• Discovery is bypassed. macOS loopback multicast is unreliable, so workers bootstrap from a fixed tcp/127.0.0.1:7447 endpoint. Zenoh scouting — the thing that would find peers on a real LAN — is therefore untested here.

So: the sharding, the KV residency, the wire protocol, and the parity are real and exercised. The distributed claim is not yet earned. Two 16 GB machines and a model that needs both (Qwen3-32B, ~17.6 GB) is the experiment that would earn it, and I have not run it.

The result worth reporting: the split is free (on one machine)

A 2-stage split runs at whole-model speed. Qwen3-8B-q40ef16, M-series, Metal, measured with distract-shardbench in a single process, so the only variable is how the model was built:

config	ms/step (min of 14)	tok/s
whole model (load_model)	42.5	23.5
full-range shard, unsplit	42.4	23.6
2-shard chain	43.0	23.3

Token-identical across all three; splitting costs ~1% in-process, with no wire between the stages. (Min rather than mean: the means carry occasional scheduler outliers on a laptop, e.g. the unsplit shard's mean swings 48-58 ms across runs while its min is stable.) A live 2-worker cluster — both workers on the same machine, zenoh over loopback — does 21-22 tok/s, TTFT ~190 ms, each worker holding 2.2 GB of a 4.29 GB model.

That number was not free. The 2-shard chain was originally 8.4 tok/s and I assumed the pipeline bubble or the KV round-trip. Both were wrong: an unsplit full-range shard was equally slow, which pointed at the shard builder, not the split. load_shard built its model with the low-level nnef().translate(&proto, ..), which — unlike the api/rs loader (model_for_path(p)?.into_decluttered()?) — does not declutter. Undecluttered, the graph was 5364 nodes instead of 1690; MetalTransform's rewrite rules only match decluttered patterns, so 73 MultiBroadcastTo stayed on the host instead of becoming GpuMultiBroadcastTo, plus 146 stray host Slices. CPU was unaffected (271 vs 306 ms), so a CPU-only A/B hides it completely. One .into_decluttered()? fixed it.

That translate() silently returns an undecluttered model — no error, no warning, ~2.8x on GPU only — may be worth a docs note or a guard at the GPU transforms; happy to file separately if useful.

How to try it

Models are the q40ef16 NNEF exports CI publishes. The coordinator is the zenoh router, so start it first; workers and dashboard are clients and retry until it is up.

MODEL=/path/to/Qwen3-8B-q40ef16.nnef.tgz

cargo run --release -p tract-distributed --bin distract-llm -- --model "$MODEL" --workers 2

# each worker loads ONLY its own layers' weights — the full model is never materialised
cargo run --release -p tract-distributed --bin distract-worker -- --name node-a --backend metal --mem-mb 4096
cargo run --release -p tract-distributed --bin distract-worker -- --name node-b --backend metal --mem-mb 4096

# optional: live cluster view + chat box on http://127.0.0.1:8088
DISTRACT_TOKENIZER=/path/to/tokenizer.json cargo run --release -p tract-distributed --bin distract-dashboard

--backend cpu works identically; a heterogeneous cpu + metal cluster was the original demo. CPU currently needs #2477 or Qwen decodes NaN; Metal is fine on main (#2476, #2472, #2428 merged).

Reproduce the table with distract-shardbench <model> metal, and check device-vs-host op placement with distract-metalaudit <model> <n_layers>.

Docs

• examples/dis-tract/README.md — what it is, how to run it, what is and is not tested.
• examples/dis-tract/docs/FAQ.md — why the cluster's tok/s is half a node's, why a shard can be slow on Metal and identical on CPU, why the KV never crosses the wire, how shards are assigned (memory only, and why a budget is a claim rather than a reservation), whether it serves one request or many (one), whether it works off a LAN (no).
• examples/dis-tract/docs/shard-cache.md — a design, not implemented: how a worker could get only the tensors it needs, from a shared path or over the wire, and cache them verified. Includes the measurements that shape it — zenoh carries the 334 MB embedding in one payload but nothing at 1 GiB, and SHA-256 verifies a shard in 0.9 s vs MD5's 3.3 s because ARM accelerates it.
• examples/dis-tract/docs/upstream-partition-api.md — the partition primitive.

The design question

You suggested on #2465 that the split be a registered ModelTransform. I could not make that work, and I think the conflict is real rather than a matter of effort: a ModelTransform operates on an already-loaded TypedModel, so the full model must be materialised first — which is precisely what a too-big-for-one-machine split has to avoid. This crate instead prunes the NNEF graph AST and reads only the shard's .dat tensors (EXO's approach), which needs loader internals api/rs does not expose.

So I suspect the upstreamable primitive is not the splitter but a public partial-load API for NNEF — load a graph subset plus only its tensors. Given that, Dis-tract becomes an ordinary causal_llm-style example on the public API and most of what is objectionable below disappears. Which way would you like it to go?

Where it stands — honest limits

• Not on the public API. 35 tract_core/tract_nnef/tract_transformers import sites, against causal_llm's single use tract::. So it does not meet the bar the examples set. That is the blocker the design question decides.
• Never run on more than one machine — see the section above. The distributed claim is unearned.
• The coordinator still materialises the full model to compute the layer weight profile, then drops it. So even once it runs on two boxes, the coordinator needs one that fits the whole model — which for the interesting case (a model too big for any single node) is a hard blocker, not a wart. The profile should come from the graph AST, as the shard loader already does.
• shard_graph.rs hardcodes torch2nnef naming (model_model__{N}_inputLayernorm_...). Fine for the published q40ef16 Qwen/Llama exports; not general.
• api/rs escape hatch: this branch adds Model::typed_ref / into_typed_model (#[doc(hidden)]), because a Model cannot otherwise be split or measured. Smallest thing that unblocks it — reshape as you prefer.
• Greedy decode, one prompt at a time, KV reset per prompt (no multi-turn memory).
• A node's memory budget is a claim about the wrong pool, and nothing reserves it. It is system RAM (available_memory() * --mem-frac, or --mem-mb) sampled at join, whatever the backend. On Apple silicon that is roughly right, since unified memory means a Metal buffer really does come from the same pool. On a discrete GPU it is simply wrong — a box with 128 GB of RAM and a 16 GB card advertises ~100 GB, the plan fits it, and the shard then has to land in VRAM. Two co-located workers also double-count the same host RAM under the default --mem-frac; passing --mem-mb side-steps it, which is why the launch script always does. The planner's fit check compares against these numbers, so it is only as good as they are.
• CUDA compiles and self-reports via Runtime::check() but is unexercised on Apple hardware — the budget issue above is one concrete consequence.
• 13 binaries, only 4 of which are the product (llm, worker, dashboard, gen). The other 9 are diagnostics that earned their keep finding the declutter bug (shardbench, metalaudit, probe, shard-run, cmpconst, dump, eval, shard, trace); say the word and I will cut them down to the two worth keeping.

Transport is zenoh (scouting + pub/sub + a queryable for assignment) — the same stack EXO uses, on crates.io, Apache-2.0. Tensor parallelism is a documented hook, not implemented.

🍍

[czoli1976]

Some directions this could go, roughly in the order I'd argue for them. Most came out of building it; a couple are already designed and measured in examples/dis-tract/docs/.

Correctness — both now fixed in this PR

A worker dying deadlocks the coordinator. Fixed. The coordinator now watches the liveliness tokens the workers already declared, so zenoh retracting a dead node's token ends the in-flight generation at once, and the per-token wait is bounded anyway for a stage that is merely stuck. It abandons that generation with its partial tokens and an error rather than retrying — the sequence is unrecoverable, but the next prompt resets every stage's KV, so the cluster keeps serving. Verified by killing a worker mid-generation: generation aborted at turn 65: node node-metal-2 left the cluster, then restarting it and getting a normal answer with no coordinator restart. Previously that left the coordinator at 0.0% CPU forever with a healthy rejoined worker idle beside it.

The planner has no fit check. Fixed. A split that would hand a node more weight than it advertised now fails at plan time, naming the stage, what it would hold and what the node offered, instead of one worker discovering it ~40 s in. Weights are a lower bound (KV and runtime measured 1.2-1.6x on an 8B model), so it catches a node that cannot possibly fit its shard, not one that merely might not — deliberately, since a headroom factor would have been a guess. The single-node path returned before any check at all and is now covered too.

What that check is still only as good as: the budget itself. A node's mem_budget is available_memory() * --mem-frac (or --mem-mb) sampled at join, ~40 s before the shard loads, and nothing reserves it. Two co-located workers double-count the same host RAM under the default frac — NodeCaps carries the hostname, so the coordinator could divide a host between its nodes, but does not. And it is system RAM whatever the backend: fine on Apple silicon where memory is unified, wrong on a discrete GPU, where a 128 GB box with a 16 GB card advertises ~100 GB and the shard still has to fit VRAM. metal-rs exposes has_unified_memory / recommended_max_working_set_size and CUDA has cuMemGetInfo, so a per-backend budget is the real fix; I have no NVIDIA hardware here, so I would rather not write that part blind.

Making the premise true

TODO The coordinator materialises the whole model to plan. It loads it only to count layers and build a per-layer weight profile, then drops it — but that means it needs a box that fits the model, which is exactly what the too-big-for-one-node case does not have. So the headline claim is blocked by the one component that does not need the weights at all. The profile could come from the graph AST, which shard_graph already parses for pruning. This is the single highest-value item here.

TODO Then actually run it on two machines. Everything so far is two processes on one box over loopback, on a model that fits it. Two 16 GB machines and Qwen3-32B (~17.6 GB) is the experiment that earns the word "distributed", and it would also put a real number on the per-token wire cost that loopback hides.

Throughput

TODO Micro-batching. With a single stream only one stage is ever busy: stage 1 idles while stage 0 computes token n. Each node is ~50% idle in a 2-stage split and ~1/N for N stages, and that does not improve with more nodes. Keeping several sequences in flight is what makes pipeline parallelism pay, and it is the difference between a latency demo and a serving system. Single-stream latency is already fine (splitting costs ~1%).

TODO Drop the double optimise on the load path. The worker optimises a clone of its shard purely to classify I/O slots, discards it, then hands the raw model to prepare() which optimises again. Shard building is ~40 s of the ~49 s startup, so this is the hot spot, and it needs no protocol change.

TODO A profile-guided planner. Placement is memory-proportional only: a Metal node and a CPU node with equal budgets get equal layers, and since every token traverses both, the pipeline runs at the slower one's pace. Today you bias it by lying about budgets. Weighting by measured throughput and link cost is the obvious next step.

Deployment

TODO Ship shards to workers, cache them per node. Today AssignSpec carries a path, so every worker needs the whole 4.29 GB export on its own disk and decompresses all of it to reach the 1.8 GB it owns — a bare node cannot join, and a model too big for a node's disk cannot run even though only its RAM is the real constraint. Designed in docs/shard-cache.md, with the measurements that shape it:

• the old blocker does not apply — nothing needs re-serializing (graph.nnef is 0.91 MB and ships whole; .dat files are already a wire format)
• zenoh carries the 334 MB embedding in one payload intact, but nothing at 1 GiB — so chunk-per-tensor is the design, whole-shard is not
• SHA-256 verifies a shard in 0.9 s vs MD5's 3.3 s (ARM accelerates it), so a restart costs ~0.9 s of hashing instead of a transfer

with the source as a resolution order — cache, then a shared path, then the wire — so a LAN with a NAS never touches the wire and cross-site nodes need no mount, from one config. Not implemented.

TODO An NNEF directory instead of a .tgz. Gzip is a stream, so "read only my shard's tensors" still decompresses the whole archive. tract already loads a directory natively, and the model's .dat files are already per-tensor with the layer in the filename — so this mostly means not using the archive.

Further out

Tensor parallelism (the step channel already carries a generic tensor vector, so a collective needs no new message type); multi-turn context (the KV reset per prompt is a deliberate PoC choice, not a constraint); sampling beyond greedy (argmax is what makes the token-parity check against a single-machine reference possible, so it would want to stay available).

[kali]

Sorry for being so long. Here is a couple of thoughts, focusing on API usabillity for now, more than actual evaluation of performance of the distribution. I do not know yet if "dis-tract" should be a tract consumer, or a tract optional runtime, or something else. But it's a good exercise to figure out the API. And I hear you, it's not there yet. Let me summarize the main snags you hitted in my own words:

1. Loading the model is all or nothing, all the weights get in memory before we have a chance to split and prune the rest. I think a possible answer to this is just lazy loading the weights (when delivery allow, if the model is a http stream, this will not work. So file system). This is something I have played a bit with before (maybe there is a stale PR ?), but now we have a cleaner framework for this with the exotics TensorStorage. I designed it with a mind to load cuda models without the weights getting in main memory, since cuda can load data from files now.
2. node names. as a "cheap" alternative to model walking, did you try using the edge labels to find the cut points ? They may be more stable than nodes name (and we may want to give them a bit of love so they are more stable).
3. I think transformers will end up being folded in core, which should somewhat force us to reconsider how to deal with the transformers detectors.

As a next step, are you interested in looking at the lazy loading story ? I think I'll spend a bit more time on recovering the various performance drift the bench refactoring has surfaced, so I'd prefer not to add another dis-traction.

[czoli1976]

Rebased on main and pushed. All three framings were useful; two were directly actionable.

1. Lazy weights. Found the stale PR — #1386 lazy-values. The seam is right and the
output type is wrong, which I think is what you're getting at with TensorStorage.
LazyConstProvider is the correct hook, but LazyConst is a stateful op that re-reads from
disk on every eval: fatal for a weight touched each decode step, and since it stops being a
Const, declutter and the block-quant fusion no longer see it — on q40 you'd lose the packed
paths. Two smaller things: it's directory-only (the tar branch passes None, matching your
stream caveat), and LazyDatLoader still calls read_tensor just to keep dt_shape, so the
first pass streams everything through memory anyway; the fact should come from the .dat
header.

So: have the provider yield a Const backed by file storage instead of an op that re-reads.
Read once, stays a const, same mechanism you want for CUDA. Happy to take #1386 that way if
the shape sounds right — it would delete much of the 700-odd lines of NNEF-text machinery in
shard_graph.rs, which exists to read only a shard's .dat files.

2. Edge labels. Tried, works — cut points now come from cache depth over outlet_labels
rather than node names, and they survive the passes that run before a shard is built. Caveat
on "giving them some love": it's still string parsing, just on stabler strings
(contains("cache"), then rsplit('_') for the layer). What would actually help is
structured metadata on the edge — ask "is this a KV slot, which layer?" without agreeing on a
spelling.

3. transformers into core. Small exposure: one call site (transformers_detect_all +
unfold-kv-cache). The fragile part isn't the detector, it's shard_graph.rs hardcoding
torch2nnef naming, which breaks wherever transformers lives.

One question. Micro-batching is in (per-sequence KV, interleaved steps; concurrent output
byte-identical to serial). Gain is small, and why is the interesting part — Qwen2.5-7B, cpu
stage 0 + metal stage 1, two concurrent sequences vs the same two back-to-back:

cut	stage 0 (cpu)	stage 1 (metal)	ceiling (Σtᵢ)/max(tᵢ)	measured
14	514.2 ms	44.6 ms	1.09x	1.12x
5	261.7 ms	65.2 ms	1.25x	1.26x

Both at the pipeline ceiling, so nothing is left in the scheduler. memory_weighted_cuts
balances bytes — right for fitting, wrong for throughput once stages sit on backends of
different speed, since an equal split idles the fast one. Moving the cut 14 → 5 by advertised
budgets alone, no code change, nearly halved wall-clock. Balancing time would mean nodes
advertising throughput as well as memory: worth doing, or is memory-only the right scope here?

Still a PoC, still two processes on one box. #2495 means the CPU path no longer NaNs, so the
docs no longer point at an unmerged PR.

[czoli1976]

Pushed a fix for something the README was simply wrong about. It claims only the residual
crosses the wire; I finally measured, and seven tensors did.

Placement asked cache depth which block a node belongs to. Depth flows along the residual,
so the query/key/value projections a block computes from its predecessor's residual took
the predecessor's depth — they were computed one stage early and shipped to the stage
that needed them. Segmenting by the residual instead (a node belongs to the last block
whose incoming residual reaches it) puts them back where they are used.

Per decode step on Qwen2.5-7B, cut at layer 14: 7 wire tensors -> 4, constant payload
16.5 KB -> 7.5 KB. Prefill of a 1000-token prompt drops ~8.8 MB. Greedy decode stays
token-identical to the single-machine reference and prefill cosine is 1.000000.

Two notes. The larger effect may not be the bytes at all — this moves ~9 MB of matmul off
the boundary stage, which matters whenever that stage is the bottleneck. And it shifts
those weights between shards, so the planner's prediction and the shard's real contents can
now disagree; measured drift is 0.4%, and wire_cost.rs pins it so a future placement
change cannot quietly invalidate the fit check.

Still crossing, and the more interesting one: the attention mask, in f32, sized S+P. At
32k context that is ~125 KB per token and growing, against a flat 7 KB residual. It is a
pure function of position, but each stage cannot recompute it because the shared tables
read past-length from caches the other stage owns — RoPE from one, the mask from another.
Normalising those reads per stage is the next piece.

[czoli1976]

Follow-up to the last one: the wire is now a single tensor.

I had assumed the shared tables could not be rebuilt per stage, because they read the
past-length from arbitrary caches — RoPE from one, the mask from another — so neither
stage would have what it needed. That is true of the graph as exported. It is not true of
the model after declutter and the transformer passes: walking their cones, all three reach
no Source at all, only constants and arithmetic over the sequence dimensions.

rotaryEmb_cosTo0   10 nodes   sources: []
rotaryEmb_sinTo0   10 nodes   sources: []
unsqueeze1.1       25 nodes   sources: []     <- the mask

So there was nothing to normalise. Keeping them off the frontier is enough: each stage's
extract_subgraph pulls the cone in and rebuilds them, a few dozen ops.

Per decode step on Qwen2.5-7B, cut at layer 14:

	before	after
tensors crossing	7	1
bytes/token, no context	16.5 KB	7.0 KB
bytes/token, 32k context	141.5 KB	7.0 KB
prefill, 1000-token prompt	19.9 MB	6.8 MB

The mask was f32 and S+P wide, so it was ~88% of the payload at 32k and growing — the
wire scaled with the conversation. It no longer does.

Parity is unchanged: prefill cosine 1.000000, and greedy decode token-identical to the
single-machine reference over 8 steps. That was the thing worth checking, since both stages
now build the causal mask independently and any disagreement would show up as divergence a
few tokens in.

wire_cost.rs asserts one wire tensor per stage at three different cut points, so this
cannot quietly regress — it would otherwise only be visible to someone running a long
conversation over a real network.

The README, FAQ and event trace all claimed only the residual crossed. That was wrong when
written; it is now true.