Zachery Aaron Shores-Chmielewski
a2ec548134
Extend the single-GPU example into a two-node pipeline-parallel run that splits llama3.2:1b across two rented vast.ai GPUs and closes the autoregressive decode loop over iroh.
- topology: add linear-chain helpers where each stage derives its neighbours locally from `STAGE`/`NUM_STAGES`, registering `pp-entry`/`pp-exit`/`pp-stage-{i}` SWIM names
- messages: add `StageActivation` (bf16 hidden-state hand-off carrying position/seq_len/is_prefill) and `NextToken` (sampled-token feedback with a `done` flag) that close the autoregressive loop between stage 0 and stage 1
- stage_actor: add `Stage0Actor` (tokenize -> embed_and_forward -> prefill activation; decode_step on each NextToken) and `Stage1Actor` (forward_and_sample -> NextToken back; emit InferenceResponse on EOS/max_tokens)
- vastai: fork the client and add `create_pipeline_instances` (rents one instance per stage, threading `STAGE`/`NUM_STAGES`, best-effort destroys on partial failure) and `destroy_all_instances`
- pp_tinygrad_worker.py: per-stage worker slicing `model.blk[start:end]` in stub and real (GGUF) modes, plus new `pp_gpu_node`/`pp_smoke_run` binaries and ROADMAP/SPEC/TEST_SPEC docs
- reuse: build on the single-GPU example's iroh transport and process bridge unchanged; add actor/codec/topology/integration test suites
Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
46 lines
1.5 KiB
Bash
Executable file
46 lines
1.5 KiB
Bash
Executable file
#!/usr/bin/env bash
|
|
# docker-gpu-node.sh — shim that pp-smoke-run can spawn instead of the
|
|
# pp-gpu-node binary directly. Boots one pp-gpu-node container per stage
|
|
# on the host network so iroh can dial without NAT.
|
|
#
|
|
# Required env (forwarded by pp-smoke-run):
|
|
# STAGE, NUM_STAGES, SEED_ADDR, SEED_DIRECT, MAX_TOKENS
|
|
# Optional env (forwarded if present):
|
|
# MODEL, PP_WORKER_STUB, PEER_NODE_ID, PEER_DIRECT
|
|
#
|
|
# Configurable via this shim:
|
|
# PP_IMAGE — image tag (default: swactor-pp-gpu:latest)
|
|
# PP_CONTAINER_PREFIX — name prefix (default: pp-stage)
|
|
# PP_DEV — tinygrad device override (default: CPU)
|
|
# PP_CACHE_DIR — host path to tinygrad cache (default: $HOME/.cache/tinygrad)
|
|
set -euo pipefail
|
|
|
|
IMAGE="${PP_IMAGE:-swactor-pp-gpu:latest}"
|
|
PREFIX="${PP_CONTAINER_PREFIX:-pp-stage}"
|
|
DEV="${PP_DEV:-CPU}"
|
|
CACHE_DIR="${PP_CACHE_DIR:-$HOME/.cache/tinygrad}"
|
|
NAME="${PREFIX}-${STAGE}"
|
|
|
|
# Idempotent cleanup of any stale container with the same name.
|
|
docker rm -f "$NAME" >/dev/null 2>&1 || true
|
|
|
|
# Ensure cache dir exists so the volume mount doesn't create a root-owned dir.
|
|
mkdir -p "$CACHE_DIR"
|
|
|
|
exec docker run --rm \
|
|
--name "$NAME" \
|
|
--network host \
|
|
-e STAGE \
|
|
-e NUM_STAGES \
|
|
-e SEED_ADDR \
|
|
-e SEED_DIRECT \
|
|
-e MAX_TOKENS \
|
|
-e MODEL \
|
|
-e PP_WORKER_STUB \
|
|
-e PEER_NODE_ID \
|
|
-e PEER_DIRECT \
|
|
-e DEV="$DEV" \
|
|
-e WORKER_SCRIPT=/usr/local/share/pp_tinygrad_worker.py \
|
|
-v "$CACHE_DIR":/root/.cache/tinygrad \
|
|
"$IMAGE" \
|
|
sh -c '[ -n "$DEV" ] && unset CUDA; exec pp-gpu-node'
|