Distributed_LLM

This repo now has two parts:

• edgecolab/: current implementation (model-family runner + adapter architecture)
• previous_implementations/: archived earlier scripts and source patches

Install From GitHub Repo

Install directly from this repository (no PyPI required):

pip install "git+https://github.com/hzhou10cs/Distributed_LLM.git@main"

For local development from a cloned repo:

pip install -e .

To lock to an exact revision:

pip install "git+https://github.com/hzhou10cs/Distributed_LLM.git@<tag-or-commit>"

Run From Source (No Package Install)

Install dependencies only:

pip install -r requirements.txt

Then run with module commands from repo root:

python -m edgecolab --help
python -m edgecolab.local_ring --help
python -m edgecolab.example_local_stage --help

Current Entry Points

• Generic CLI:

edgecolab \
  --model_name meta-llama/Llama-3.2-1B \
  --model_family auto \
  --layer_begin -1 \
  --layer_end 6 \
  --recv_host 127.0.0.1 --recv_port 5600 \
  --send_host 127.0.0.1 --send_port 5601

• Local 3-process ring:

edgecolab-local-ring \
  --model_name meta-llama/Llama-3.2-1B \
  --model_family auto \
  --layer_end1 6 \
  --layer_end2 11 \
  --host 127.0.0.1 \
  --base_port 5600 \
  --max_new_tokens 100

• Single-process 1-device example:

edgecolab-example \
  --model_name meta-llama/Llama-3.2-1B \
  --role full \
  --prompt "Once upon a time, a wizard lived in a tower." \
  --max_new_tokens 100

Determinism and telemetry (M0)

--temperature 0 selects greedy decoding, the only reproducible mode. Anything comparing two runs — different partitions, before/after a reconfiguration — must use it, or the runs diverge for reasons unrelated to what is being tested.

Pass --trace_dir to write machine-readable JSONL telemetry alongside the human-readable device_*.log. One file per stage, one record per token:

edgecolab-local-ring \
  --model_name meta-llama/Llama-3.2-1B-Instruct \
  --layer_end1 6 --layer_end2 11 \
  --max_new_tokens 100 --temperature 0 \
  --trace_dir traces/

Fields: compute_ms is this stage's forward pass alone; recv_wait_ms is time blocked in recv (upstream compute plus link, i.e. how long this stage was starved, not its own cost); send_ms is the blocking sendall; tpot_ms is a full trip around the ring and is only meaningful on the head stage.

Equivalence check

Splitting the model must not change what it generates. This compares a single-process greedy run against the 3-stage ring:

python tools/check_equivalence.py \
  --model_name meta-llama/Llama-3.2-1B-Instruct --max_new_tokens 32

Exit code 0 means every token id matched. A mismatch at index 0 points at prefill (causal mask or rotary positions across a stage boundary); a later mismatch points at per-stage KV cache handling during decode. Run this after any change to partitioning, caching, or the adapter.

The reference run deliberately uses a different loader from the ring (--ref_loader hf vs --loader shard), so the two sides stay independent: if both used the shard loader, a bug in it would cancel out and the test would pass on two identically-wrong runs.

Strict weight residency (M1)

Each stage materialises only the blocks it owns. The skeleton is built on torch.device("meta"), pruned to the stage's layer range, and only then filled from the checkpoint's safetensors with load_state_dict(assign=True).

--loader shard   # default: strict residency
--loader hf      # legacy: load the whole model, then prune
--dtype auto     # default: the checkpoint's own dtype (bf16 for Llama-3.2)

Measured on Llama-3.2-1B-Instruct, bf16, CUDA, cuts (6, 11):

stage	blocks	resident	before (hf)
head	6	1197 MiB	4714 MiB
middle	5	580 MiB	4714 MiB
tail	5	1081 MiB	4714 MiB

The split is not thirds. tie_word_embeddings=True on Llama-3.2, and the checkpoint stores no lm_head.weight at all — the tail sources it from the embedding tensor, so that ~501 MiB table is resident on both head and tail.

Two things the loader asserts rather than assumes, because both fail silently otherwise: that no tensor is left on meta after loading (non-persistent buffers such as rotary inv_freq are computed at init and never appear in the checkpoint), and that every declared key prefix matched at least one tensor.

Per-stage load timing lands in stage_meta.shard_load. Note that on CPU safetensors are memory-mapped, so weights_ms there measures mmap setup, not I/O — use the CUDA numbers when calibrating real load bandwidth.

load_blocks materialises individual blocks by global index, for migration. It takes a CheckpointIndex (open it once per stage with open_checkpoint) and a block_factory from the adapter. Both exist for the same reason: constructing a full 16-layer meta skeleton per call to keep one block cost ~44 ms against ~0.8 ms of actual reading. CheckpointIndex caches metadata only — the snapshot path, config, and safetensors key map. Weight caching stays off deliberately; caching weights would measure the cache instead of the load.

The factory must return a fresh meta block each call, and load_blocks asserts it. A block handed back from a retained skeleton would be filled in place by load_state_dict(assign=True), leaving every block ever migrated materialised inside that skeleton — passing every short test and running out of memory later.

Control plane and windowed execution (M2)

Decode runs in windows of W tokens. At each window boundary the stages run a two-phase barrier over a separate control socket, star-connected to the head. The data ring is untouched and keeps its blocking, untimed semantics.

--window 20            # W: decode tokens per control window
--plan noop            # coordinator on the head (noop | forced_shift)
--ctrl_port 5700       # head's control listener -- SAME value on every stage
--ctrl_timeout 5.0     # control sockets only; the data ring stays untimed

local_ring derives --ctrl_port from base_port + 100 automatically. For a manual 3-terminal run you must pass the same port to all three stages: it is the head's listener, not a per-stage port.

Control never shares the data ring. Multiplexing it there deadlocks — a stage that commits a new boundary stops producing packets its neighbour is still blocking on, and no data socket has a timeout. It would also make Ω unmeasurable, since measured overhead would absorb however long stages happened to be blocked.

The cut is atomic because stages pre-stage before they ACK: a stage answers "ready" only once it physically holds the state the new configuration needs. A lost message stalls; it never half-applies. Invariants asserted in code:

• a stage applies epoch e only on CTRL_COMMIT{e}; a timeout is an abort
• the head commits only after every stage ACKs
• epochs increase monotonically, so retries are idempotent
• after commit, cache layer order must still match core.layers order

Each barrier emits a window telemetry record (barrier_ms, prestage_ms, commit_ms, n_acks, aborted, decision).

Control-plane tests

These run in ~3 s each with no model loaded, which is what keeps control bugs distinguishable from pipeline bugs:

python tools/check_control_plane.py            # happy, nack, kill
python tools/check_control_plane.py --only kill

kill terminates a stage mid-barrier; the survivors must stall with a diagnostic rather than hang or emit tokens. Note that a dead peer on localhost produces RST/EOF, not a timeout — both are treated as "peer state unknown, do not proceed".

Measured barrier cost

Llama-3.2-1B-Instruct, CPU, W=20, 200 tokens:

measurement	median	what it is
isolated control plane	0.76 ms	pure protocol
tail-side barrier, live ring	0.77 ms	same, confirmed in situ
head-side barrier, live ring	66.3 ms	protocol + pipeline drain

The head reaches the barrier first and waits for the in-flight token to drain through the other stages; that wait overlaps with useful work. The tail waits for nobody, so its number is the protocol cost. Across 3 repetitions the end-to-end effect of 9 barriers was smaller than run-to-run variance.

Boundary migration (M3)

Blocks move between adjacent stages mid-generation, without changing a single output token.

# 200 tokens, shift b_1 up by one block at token 100, verify against no-shift
python tools/check_equivalence.py --model_name meta-llama/Llama-3.2-1B-Instruct \
  --max_new_tokens 200 --window 20 --forced_shift_at 100 --shift_k 1

# same matrix in one process, no sockets -- isolates migration from transport
python tools/check_migration.py --tokens 40 --shift_at 20

--shift_k is positive to pull blocks backwards from the higher stage and negative to push them forwards; --shift_boundary selects which b_s moves (b_0 and b_S are fixed — the embedding and lm_head never migrate).

What moves, and from where:

• weights are read from local disk by the shard loader. Every node has the whole checkpoint, so this is a disk read, not a transfer. Pulling them from the donor would inflate the measured reconfiguration cost.
• KV comes from the donor over a dedicated migration channel — one bidirectional link per adjacent pair, separate from both the control star (8 KiB cap) and the unidirectional data ring.
• weights must be resident before KV is installed, and the donor releases nothing until COMMIT, which is what makes an abort recoverable.

A multi-megabyte send blocks once the socket buffer fills if the peer is not reading, and during pre-staging both peers are active. Resolved with one sender thread per outgoing transfer, joined before the ACK — confined to the barrier, sharing no state.

Measured cost

1-block shift on Llama-3.2-1B at token 100, CPU, 200 tokens:

term	stage	value
weight load (local disk)	receiver	4.1 ms (1.9 construct + 2.1 read)
KV transfer	receiver	20.4 ms
commit (pointer swaps)	both	0.2–0.5 ms
release (gc.collect)	donor	103.5 ms

Two caveats that matter more than the numbers:

• The weight term is a warm-page-cache memcpy, not a disk read. safetensors are memory-mapped, and the measured read runs at ~150 GB/s. On a device with cold cache and slower storage it will be much larger. These numbers establish the mechanism, not the magnitude.
• The receiver's KV milliseconds are an upper bound. 20 ms to move 226 KiB is not transfer time; it absorbs inter-stage skew, because the receiver finishes its weight load in ~4 ms and then waits for the donor to reach the barrier. The bytes are exact; do not fit a bandwidth to the milliseconds.

KV is 2 × num_key_value_heads × head_dim × dtype_bytes per block per token — 2048 B for 1B, 4096 B for 3B — asserted in telemetry, because a mismatch would most likely mean the whole cache was serialised instead of one layer's, which would still "work".

Releasing blocks needs gc.collect(). nn.Module graphs are cyclic, so a dropped block with zero referrers still occupies the device until a cycle collection runs. It is the dominant term above, it scales with total heap size rather than with how much was dropped, and it is therefore a runtime artifact rather than a property of reconfiguration — release_ms is reported separately so it can be excluded. empty_cache() follows it: on unified-memory devices, reserved-but-unused bytes starve the sibling stages.

Comparing against the right baseline

--baseline reference   # default: the single-process M0 oracle
--baseline noshift     # run the ring twice, with and without the plan

Use noshift for anything involving migration, and always on CUDA past ~60 tokens. There, the ring and the single-process reference diverge even with no migration at all: at token 61 of the default prompt the top-2 logits are 24.375 vs 24.250, one bf16 ULP apart at that magnitude, so argmax flips for reasons unrelated to partitioning. Both paths are internally deterministic, and CPU agrees to 200/200.

Notes

• meta-llama/Llama-3.2-1B is gated; first download still requires access/auth.
• After full cache exists locally, local-only loading paths are used automatically.
• Pin installation to a Git tag or commit hash to keep behavior tightly bound to GitHub source.