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, no fit check), 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).
• CUDA compiles and self-reports via Runtime::check() but is unexercised on Apple hardware.
• 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, before anything else

TODO A worker dying deadlocks the coordinator. The worker side is fine — liveliness evicts it, a restart with the same --name reclaims its stage and rebuilds (~28 s), and KV self-heals since every prompt resets the caches. But the coordinator declares no liveliness subscriber and its per-token wait (result_sub.recv_async()) has no timeout, so it blocks forever on a result that will never come, and the serial generate loop never accepts another request. Verified by killing a worker mid-generation: after the worker had fully rejoined, a new request got no reply in 45 s, coordinator at 0.0% CPU. A healthy cluster, deadlocked; only a coordinator restart clears it.

The fix is small and the pattern already exists a few lines above it (the reset-ack): bound the wait with a select! deadline, and subscribe to distract/live/* so a Delete for a participating node aborts the generation.

TODO The planner has no fit check. It splits weights proportionally to advertised budgets and never verifies a shard fits the node it lands on — an oversubscribed node fails at load time rather than plan time. That was survivable while the weight figures were themselves ~8x wrong (they counted block-quant weights at dequantized size); now that they are accurate, the check is cheap.

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).