307 lines
14 KiB
Markdown
307 lines
14 KiB
Markdown
# Pipeline-Parallel Inference — Spec (N stages)
|
||
|
||
> This spec covers the N-stage generalisation of the pipeline-parallel
|
||
> example. For the deployment-shape rationale and where this rung sits
|
||
> in the broader buildout, see [`/docs/ROADMAP.md`](../../docs/ROADMAP.md)
|
||
> §3.1. For the test enumeration, see
|
||
> [`./TEST_SPEC.md`](./TEST_SPEC.md).
|
||
>
|
||
> Behavioural snapshot of the **2-stage** deployment, as it actually
|
||
> runs today, is preserved in
|
||
> [`../single-gpu-inference/BEHAVIOR.md`](../single-gpu-inference/BEHAVIOR.md)
|
||
> for cross-reference and in
|
||
> [`./BEHAVIOR.md`](./BEHAVIOR.md) (to be written once the N-stage path
|
||
> is deployed).
|
||
|
||
## 1. Goal
|
||
|
||
`pp-smoke-run --vastai --num-stages N --api-key <key>` runs a linear
|
||
chain of `N ≥ 2` GPU nodes that jointly serve one `InferenceRequest`
|
||
end-to-end. Acceptance test: `N = 5` on vast.ai, `llama3.2:1b`,
|
||
prompt `"Say hello"`, produces non-empty text within the 5-minute
|
||
budget. The localhost (`--seed`) path supports the same N values via
|
||
spawning N child processes.
|
||
|
||
`N = 1` is **not** supported by this example — single-node inference
|
||
is what `examples/single-gpu-inference` exists for, and the role-based
|
||
actor design here is built on the assumption that "first" and "last"
|
||
are distinct stages. `N = 2` is the minimum: first stage handles the
|
||
embed + first half; last stage handles the second half + head; no
|
||
middle stages.
|
||
|
||
## 2. What changes vs. the 2-stage MVP
|
||
|
||
| Area | 2-stage today | N-stage target |
|
||
|---|---|---|
|
||
| `pp-gpu-node` binary | hard-checks `NUM_STAGES == 2` (line 214) and dispatches to `run_stage_0` or `run_stage_1` | accepts any `N ≥ 2`; dispatches to a single `run_stage` function parameterised by `(stage, num_stages)` |
|
||
| Stage actors | `Stage0Actor` and `Stage1Actor` are separate types | one `StageActor` parameterised by **role** (`First` / `Middle` / `Last`) |
|
||
| Message bridges | three bridge types (`Stage0RequestBridge`, `Stage0NextTokenBridge`, `Stage1ActivationBridge`) bound to specific actor types | same three bridge shapes, but generic over the unified `StageMsg` |
|
||
| Orchestrator (`pp-smoke-run`) | spawns exactly two children, hands stage 1 stage 0's `PEER_DIRECT` | spawns `N` children in order; each non-first child receives its **predecessor's** addressing as `PEER_DIRECT` |
|
||
| Topology helpers | already generic (`stage_name`, `next_stage_name`, `prev_stage_name`) | unchanged |
|
||
| Worker (`pp_tinygrad_worker.py`) | accepts stage-0 ops (`embed_and_forward`, `decode_step`) and stage-N-1 ops (`forward_and_sample`); already param on `STAGE`/`NUM_STAGES` for layer range | adds `forward_range` op for middle stages; layer-range computation unchanged |
|
||
|
||
The worker contract is the **anchor of stability**. Anything that
|
||
doesn't have to change, doesn't.
|
||
|
||
## 3. Actor design — a single `StageActor`
|
||
|
||
### 3.1 Role enum
|
||
|
||
```rust
|
||
pub enum StageRole {
|
||
/// stage == 0. Owns: prompt entry, tokenizer, embed + first range.
|
||
/// Reacts to: `InferenceRequest` from network; `NextToken` from
|
||
/// the *last* stage (decode-loop feedback).
|
||
First,
|
||
|
||
/// 0 < stage < num_stages - 1. Owns: a middle block range, no
|
||
/// tokenizer, no sampler. Reacts to: `StageActivation` from prev.
|
||
Middle,
|
||
|
||
/// stage == num_stages - 1. Owns: last block range, output head,
|
||
/// sampler, detokenizer. Reacts to: `StageActivation` from prev.
|
||
/// Emits: `NextToken` to stage 0; `InferenceResponse` to
|
||
/// `reply_to` on terminate.
|
||
Last,
|
||
}
|
||
```
|
||
|
||
Role is set once at construction and never changes. The actor's
|
||
per-message logic branches on it. The role is computed by the binary
|
||
from `(STAGE, NUM_STAGES)` and is a pure function (see TEST_SPEC §3).
|
||
|
||
### 3.2 Unified incoming message
|
||
|
||
```rust
|
||
pub enum StageMsg {
|
||
Inference(InferenceRequest), // First only
|
||
Activation(StageActivation), // Middle and Last
|
||
NextToken(NextToken), // First only
|
||
Process(ProcessNotification), // all roles
|
||
SetNeighbors { prev: Option<ActorAddress>,
|
||
next: Option<ActorAddress>,
|
||
reply_to: Option<ActorAddress> },
|
||
Reset,
|
||
}
|
||
```
|
||
|
||
Compared to today, this collapses `Stage0Msg` and `Stage1Msg` into one
|
||
type. `prev` and `next` are `Option` because the chain ends are
|
||
asymmetric. `reply_to` is `Some` only for `Last`. An actor that
|
||
receives a variant outside its role's set logs and drops it — this is
|
||
defensive, not load-bearing, because the orchestrator wires bridges
|
||
matching the role.
|
||
|
||
### 3.3 Per-role message behaviour
|
||
|
||
| Role | On `Inference` | On `Activation` | On `NextToken` |
|
||
|---|---|---|---|
|
||
| **First** | `tokenize` → `embed_and_forward` → send `StageActivation { is_prefill: true }` to `next` | drop | `done=true` → drop; else `decode_step` → send `StageActivation { is_prefill: false, seq_len: 1 }` to `next` |
|
||
| **Middle** | drop | `forward_range` → send `StageActivation { … }` to `next`, preserving `request_id`, `position`, `seq_len`, `is_prefill` | drop |
|
||
| **Last** | drop | `forward_and_sample` → send `NextToken` to `prev` (= the first stage); on terminate, `detokenize` → send `InferenceResponse` to `reply_to` | drop |
|
||
|
||
Note: the autoregressive feedback edge is **last → first**, not
|
||
last → prev. This is the same shape as the 2-stage MVP (where prev
|
||
*is* first), and the right shape for arbitrary N — the first stage is
|
||
the only one that owns the embed step.
|
||
|
||
### 3.4 New worker op: `forward_range`
|
||
|
||
Middle stages need a worker op that takes a hidden state, runs blocks
|
||
`[lo..hi)`, and returns the new hidden state. Stage 0's
|
||
`embed_and_forward` already runs a block range; stage 1's
|
||
`forward_and_sample` runs a range and then samples. The new op is the
|
||
range-only variant.
|
||
|
||
```
|
||
→ {"op": "forward_range", "request_id": rid, "hidden_b64": "...",
|
||
"position": p, "seq_len": s}
|
||
← {"request_id": rid, "hidden_b64": "...", "seq_len": s}
|
||
```
|
||
|
||
In `pp_tinygrad_worker.py` this is the same code path as the prefix of
|
||
`forward_and_sample`, factored out and exposed as its own dispatch
|
||
case. Tinygrad's per-block KV cache continues to be allocated lazily
|
||
per block, so each middle stage owns the K/V for its own range
|
||
automatically.
|
||
|
||
## 4. Orchestrator changes (`pp-smoke-run`)
|
||
|
||
### 4.1 Seed mode
|
||
|
||
Today: spawn stage 0, read its `PP_GPU_NODE_ADDR`, spawn stage 1 with
|
||
that addressing as `PEER_DIRECT`. For N stages:
|
||
|
||
```
|
||
for i in 0..N:
|
||
spawn child i with PEER_DIRECT = (addr of child i-1, if i > 0)
|
||
read child i's PP_GPU_NODE_ADDR
|
||
```
|
||
|
||
Sequential, **not parallel.** Each subsequent child needs its
|
||
predecessor's direct addressing in `PEER_DIRECT`, so the spawn loop is
|
||
inherently a chain. Each child only needs its *predecessor's*
|
||
addressing (not all earlier addresses): the outbound dial from i → i-1
|
||
makes i-1's iroh stack learn i's source-socket reverse-path. The
|
||
seed/orchestrator is dialed by every child, so it always learns
|
||
everyone.
|
||
|
||
### 4.2 vast.ai mode
|
||
|
||
Today: `find_offer` twice (excluding the first offer's id), then
|
||
`create_pipeline_instances` with both offer ids. For N stages:
|
||
|
||
* `find_offer` N times, each call excluding all previously chosen
|
||
offer ids.
|
||
* `create_pipeline_instances` takes a slice of N offer ids, sets
|
||
`STAGE`, `NUM_STAGES`, `SEED_ADDR`, `SEED_RELAY` on each; in vast.ai
|
||
mode, **no `PEER_DIRECT` is passed** — relay-mediated NAT traversal
|
||
+ SWIM gossip handle peer discovery (this is already how vast.ai
|
||
mode works for 2 stages today, no change).
|
||
* `wait_for_running` once per contract, sequentially (same as today,
|
||
just `N` times instead of 2).
|
||
* `destroy_all_instances` is already a slice-aware helper.
|
||
|
||
### 4.3 Convergence target
|
||
|
||
The orchestrator waits for SWIM to report `alive >= N` peers (today:
|
||
`alive >= 2`, i.e. N=2). One-line change in both `run_seed` and
|
||
`run_vastai`.
|
||
|
||
### 4.4 Submitting the request
|
||
|
||
Unchanged: resolve `pp-entry`, build route, `send_to`,
|
||
`await_response`. The chain is invisible to the orchestrator past the
|
||
entry node.
|
||
|
||
## 5. `pp-gpu-node` changes
|
||
|
||
### 5.1 Boot sequence (generalised)
|
||
|
||
```
|
||
1. Parse STAGE, NUM_STAGES (any N ≥ 2; remove the == 2 check).
|
||
2. Determine role: First if STAGE==0; Last if STAGE==N-1; else Middle.
|
||
3. Build iroh driver, join seed.
|
||
4. Wait for SWIM convergence (>= 1 alive peer — unchanged).
|
||
5. Spawn the StageActor with the role + placeholder addresses.
|
||
6. Register the per-index name pp-stage-{STAGE} immediately so
|
||
neighbours can resolve us.
|
||
7. Wait for worker ready (unchanged).
|
||
8. Resolve neighbours by role:
|
||
First: resolve pp-stage-1 (next).
|
||
Note: the feedback target (NextToken from last) is *us* —
|
||
the last stage looks up pp-stage-0 itself.
|
||
Middle: resolve pp-stage-{stage+1} (next).
|
||
A middle does not send to prev; nothing to resolve there.
|
||
Last: resolve pp-stage-0 (the NextToken destination),
|
||
plus pp-orchestrator (for InferenceResponse).
|
||
9. Build IrohActorTransport routes for each resolved neighbour.
|
||
10. SetNeighbors on the actor with the resolved addresses.
|
||
11. Pump 100ms to drain the setup message.
|
||
12. Register pp-entry (First) / pp-exit (Last). pp-entry must be
|
||
registered last because the orchestrator may send InferenceRequest
|
||
as soon as it resolves — and First is not ready to forward until
|
||
pp-stage-1 is wired.
|
||
13. Enter main pump loop.
|
||
```
|
||
|
||
The step-8 "who do I resolve" matrix is the only real new logic. It
|
||
flows from the per-role message table in §3.3.
|
||
|
||
### 5.2 Name registry summary
|
||
|
||
| Stage | Registers | Resolves (and routes to) |
|
||
|---|---|---|
|
||
| First (0) | `pp-stage-0`, `pp-entry` | `pp-stage-1` |
|
||
| Middle (i) | `pp-stage-i` | `pp-stage-{i+1}` |
|
||
| Last (N-1) | `pp-stage-{N-1}`, `pp-exit` | `pp-stage-0`, `pp-orchestrator` |
|
||
|
||
This is exactly what `topology.rs::register_stage_names` already
|
||
emits — the topology helpers do not change.
|
||
|
||
## 6. Failure modes (deltas only)
|
||
|
||
The 2-stage failure table previously in this spec applies stage-by-stage.
|
||
N-stage-specific additions:
|
||
|
||
| Failure | Detection | Response |
|
||
|---|---|---|
|
||
| Fewer than N offers available | `find_offer` returns empty on attempt < N | Destroy any allocated instances, exit 1. |
|
||
| Any one of N instances never reaches running | Per-contract `wait_for_running` | Destroy *all* contracts, exit 1. |
|
||
| SWIM never converges to N alive peers | 180s timeout in vast.ai mode | Destroy all, exit 1. Fetch logs from each instance. |
|
||
| Any one of N workers never reports ready | Per-node `wait_for_worker_ready` (600s) | That stage exits 1; SWIM marks it dead; orchestrator times out on convergence and destroys all. |
|
||
| Middle stage drops a `StageActivation` (mid-decode) | Orchestrator's 600s response timeout | Destroy all, exit 1. Recovery is rung 6 (fault tolerance), out of scope here. |
|
||
|
||
The teardown invariant — "every code path that creates an instance
|
||
also destroys it" — generalises to "any creation failure rolls back
|
||
all already-created instances." `destroy_all_instances` already does
|
||
this; the loops above just have to call it on every error branch.
|
||
|
||
## 7. Test plan
|
||
|
||
The full enumeration of behavioural tests lives in
|
||
[`./TEST_SPEC.md`](./TEST_SPEC.md). Summary of tiers in landing order:
|
||
|
||
1. **Pure-Rust unit** — codec, topology helpers, role computation,
|
||
per-stage layer-range math.
|
||
2. **Python worker contract** — `forward_range` (new), all existing
|
||
ops verified against stub mode, layer-range math up to N=8.
|
||
3. **Stage actor (in-process, stub worker)** — `StageActor` in each
|
||
of its three roles plus role-mismatch defensive drops.
|
||
4. **Cluster transport (in-process, real iroh)** — N-node clusters
|
||
(N ∈ {2, 3, 4}) carry every message type intact across every hop.
|
||
5. **vast.ai client (mocked HTTP)** — N-instance create/destroy with
|
||
distinct `STAGE` env vars and rollback on partial failure.
|
||
6. **In-process N-stage integration (stub worker, real iroh)** —
|
||
full chain at N ∈ {2, 3, 4, 5}.
|
||
7. **Sliced-vs-full equivalence (real tinygrad on CPU, gated)** —
|
||
the load-bearing correctness test: N-stage output token sequence
|
||
matches single-process `Transformer.generate()` up to `max_tokens`.
|
||
Run at N=2 (carried forward), N=3, and N=4.
|
||
8. **Localised binary E2E (`pp-smoke-run --seed`)** — the
|
||
**pre-deploy gate**: runs the actual `pp-gpu-node` binaries
|
||
spawned by the actual `pp-smoke-run` binary in `--seed` mode at
|
||
N ∈ {2, 3, 5}.
|
||
9. **Manual smoke run on vast.ai at N=5** — gated, rare, the
|
||
acceptance test for the rung.
|
||
|
||
A change that breaks any tier 1–4 test must not pass review. Tier 8
|
||
must pass before any change is deployed. Tiers 5–7 are CI-gated;
|
||
real-tinygrad and equivalence are `#[ignore]` and run on demand or in
|
||
a slow lane.
|
||
|
||
## 8. Migration
|
||
|
||
Land in this order, each its own PR:
|
||
|
||
1. **Worker op `forward_range`.** Additive in
|
||
`pp_tinygrad_worker.py`. No call sites yet. Tested with Python
|
||
tests in `test_worker.py`.
|
||
2. **Unify `Stage0Actor` + `Stage1Actor` into `StageActor`.** Keep
|
||
message-shape behaviour identical for N=2. Delete the old types.
|
||
Existing `t_actor.rs` tests adapt to the unified API. **This PR
|
||
does not change the wire or the binary.**
|
||
3. **Add `Middle` role + actor logic** behind a constructor variant.
|
||
Add N=3 / N=4 integration tests against the stub worker. Binary
|
||
still rejects N != 2 to avoid premature deploys.
|
||
4. **Generalise `pp-gpu-node`** (drop the N==2 check; generalise the
|
||
resolve/route step). Update `t_binary.rs` to cover N=3 in seed
|
||
mode.
|
||
5. **Generalise `pp-smoke-run`** (spawn loop, find-offer loop,
|
||
convergence target). Add `--num-stages` flag, default 2.
|
||
6. **Land tier-8 N=5 localised E2E.** This is the pre-deploy gate.
|
||
7. **Manual vast.ai run at N=5.** No code change — just exercise the
|
||
path and capture the result in `BEHAVIOR.md`.
|
||
|
||
Each PR is independently mergeable and ships behind the existing
|
||
2-stage acceptance. Tier 8 at N=5 must pass before step 7.
|
||
|
||
## 9. Out of scope (this rung)
|
||
|
||
* Microbatching, multiple in-flight requests, streaming output.
|
||
* Dynamic stage assignment / rebalancing. Stages are decided at boot.
|
||
* Fault tolerance of any kind beyond "destroy everything on any
|
||
failure." Rung 6.
|
||
* Heterogeneous layer splits — every stage still owns
|
||
`total / N` blocks ± 1. The planner (rung 5) takes over splitting.
|
||
* `N = 1`. Use `examples/single-gpu-inference`.
|