This repo now has two parts:
edgecolab/: current implementation (model-family runner + adapter architecture)previous_implementations/: archived earlier scripts and source patches
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>"Install dependencies only:
pip install -r requirements.txtThen run with module commands from repo root:
python -m edgecolab --help
python -m edgecolab.local_ring --help
python -m edgecolab.example_local_stage --help- 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--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.
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 32Exit 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.
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.
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 untimedlocal_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
eonly onCTRL_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.layersorder
Each barrier emits a window telemetry record (barrier_ms, prestage_ms,
commit_ms, n_acks, aborted, decision).
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 killkill 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".
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.
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.
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.
--baseline reference # default: the single-process M0 oracle
--baseline noshift # run the ring twice, with and without the planUse 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.
Every stage binds its own address and connects to its peers':
--peer_hosts 10.0.0.1,10.0.0.2,10.0.0.3 # per-stage addresses, in stage order
--ctrl_host 10.0.0.1 # the HEAD's address, not your own--peer_hosts defaults to --recv_host repeated, which is correct only on
loopback — there a stage's own address and its peer's are the same string, so a
channel that confuses the two still works. Off-loopback that default makes
every non-head stage reach for the head's control listener at its own IP and
hang, so --ctrl_host hard-warns when it is omitted with a non-loopback bind.
Launch a real cluster with:
tools/launch_cluster.sh --hosts 10.0.0.1,10.0.0.2,10.0.0.3 --cuts 6,11 --tokens 200 --remote-dir ~/Distributed_LLM --python ~/venv/bin/pythonOne run_id and one set of ports for every stage — hand-typing three terminals
is where --max_new_tokens mismatches come from, and those wedge the ring. It
kills and verifies stale listeners before binding; a leftover process
holding a port produces a diagnostic that points at the wrong node.
127.0.0.0/8 is entirely loopback, so three stages can run on three different addresses on one machine — which makes bind and connect addresses genuinely distinct:
python tools/check_multinode.py # model-free: addressing, channels, parity
python tools/check_equivalence.py --peer_hosts 127.0.0.1,127.0.0.2,127.0.0.3 ...A 127.0.0.1-only run does not exercise this path.
Every stage reports its transformers/torch/python/dtype/CUDA versions — plus
L4T, JetPack and nvpmodel -q power mode on Jetson — into stage_meta and the
HELLO roster. The head refuses to start if the enforced fields disagree, because
a version skew across nodes produces exactly the output the bit-exactness tests
exist to catch, with no hint which node is at fault. Power mode and L4T are
recorded but never enforced: on a heterogeneous testbed they are supposed to
differ, and that asymmetry is the experiment. Downgrade with
--no_strict_parity.
python tools/derive_link_params.py traces/Regresses recv_wait_ms[s] - compute_ms[s-1] on payload size — no active
prober, since the ring already supplies a sample per token. On loopback this
is degenerate: recv_wait absorbs peer scheduling delay, which on a
contended host dwarfs the transfer, and the tool flags the result as
NOT A LINK MEASUREMENT rather than reporting it. Real numbers need real
Ethernet.
Two backends behind one interface, emitting the identical obs dict the DACI
schemes already read, so the controller cannot tell them apart:
--drift replay --drift_trace drift.jsonl --drift_seed 77
--drift liveSame flags on laptop and Jetson, so moving to the testbed needs no code change.
obs has exactly four keys — theta_obs, q_cmp_obs, q_mem_obs,
link_obs. Do not add more. A real sensor with no analogue here belongs in
telemetry; the point of the contract is that decide_runtime runs unmodified
against either backend.
Applying the drift without reporting it, or reporting without applying, both look like working replay:
- Applies it —
phi_nbecomes a sleep of(phi_n - 1) x compute_ms, injected inside the measured region socompute_msand the head's TPOT both move. Measured:phi=1.5on the middle stage takes it from 26.1 to 41.9 ms while head and tail are unchanged, and head TPOT follows 1.24x. - Reports the noisy reading, never ground truth — Gaussian sensor noise
with
sigma^2 = 0.5 degC^2. Ground truth goes to the trace only. Feeding the controller clean values tests a system nobody deploys, socheck_drift_source.pyfails if the noise is zero.
Seeded per (run_id, stage_idx): same seed reproduces identical token ids
and identical obs sequences, which is the precondition for comparing schemes
in M6.
The replay backend reads a file, so edgecolab never imports the simulator
(and no Jetson needs it installed). Schema is documented in
tools/make_drift_trace.py, which also generates five profiles for testing:
python tools/make_drift_trace.py --out drift.jsonl --profile one_slow_node
# steady | thermal_ramp | one_slow_node | oscillating | link_degradeA trace shorter than the run holds its last record and flags exhausted
rather than wrapping — wrapping would make a short trace look like a long
stationary run.
Sensors are read by a background thread; the barrier reads the cache. Measured
0.002 ms per sample against a 2 ms budget. A synchronous nvidia-smi would
cost tens of ms — an order of magnitude over budget, which would make the drift
source the overhead it exists to measure. obs piggybacks on the CTRL_ACK
the barrier already sends: no new socket, and the sample is aligned to the
window boundary where the decision is made.
Verification status: the laptop path (psutil + nvidia-smi) is verified.
The Jetson path is written but has never run — grep UNVERIFIED-JETSON in
edgecolab/sensors.py before trusting any Jetson number.
python tools/check_drift_source.py # model-free: contract, determinism,
# noise, injection, budget, exhaustionpython tools/bench_release.py --device cuda # does gc.collect scale with heap?
python tools/bench_ln.py --device cuda --cold # what is L_n really?bench_release answers no: heap rises 68% for a +0.15 ms change in gc_ms.
The cost is the interpreter's object graph (~360k tracked objects, of which the
model is under 1% — tensors are C-level and not individually tracked), so it is
a near-constant ~100-130 ms per reconfiguration rather than something that
grows with model size. It is reported as its own release_ms component so a
published overhead figure can exclude it.
bench_ln refuses to print an L_n number without a cold page cache,
which needs Linux and root. The warm number is an mmap'd memcpy, not a disk
read. It does report that the weight read is 4.5x slower under a co-running
compute load, which matters: if L_n is not stable, the cost model's weight
term needs a drift factor of its own.
Replaces the cost model's FLOPs-per-second term with measured coefficients:
python control/calibrate.py --device cuda --out configs/devices_measured.json
python control/calibrate.py --only decode,swap --contexts 128,2048,8192Fits t_block^dec ~= c0 + c1*(P+t) for decode and t_block^pf ~= d0 + d1*P + d2*P^2 for prefill, and measures L_n, H_swap, link alpha/beta,
contention and thermal response. phi keeps its exact meaning as a
dimensionless degradation factor, which is what makes the swap safe.
Sweep ranges matter more than they look. Timing a single 1 ms decode step
on a boost-clocking GPU measures jitter: the per-block time came out
non-monotone in context and the fit was meaningless at r2=0.13. The harness
times a batch of 24 steps and divides, which lifts it to r20.9. Likewise d2
fits negative — unphysical — if prompts are shorter than ~512, because the
quadratic term does not exist yet at that scale. The defaults are set where the
term being fitted is actually observable; shrink them and you will fit noise.
Measured on a laptop the KV term is invisible below ~2k context and reaches
+74% of c0 at 8192, which is the argument for the affine form: c0, the
memory-bandwidth-bound weight-streaming term, dominates at edge context
lengths.
Every emitted record carries "calibrated": false plus a warning naming each
reason the values are invalid off-testbed — c0/c1 are the wrong GPU,
L_n is a warm mmap'd memcpy, alpha/beta are loopback memcpy, and a laptop
has no thermal knee to fit. Two fields are deliberately left null rather than
fabricated: rho (the co-runner saturates rather than loads — the record says
so) and theta_th/gamma/nu (no throttling observed).
configs/devices_measured.json in this repo is a laptop-generated schema
example, not calibration data. Its schema is reconstructed from the plan's
field names because the simulator's configs/devices.json is not in this
workspace; --schema_from PATH diffs the two key-for-key once it is.
meta-llama/Llama-3.2-1Bis 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.