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>
28 lines
1.1 KiB
Docker
28 lines
1.1 KiB
Docker
FROM nvidia/cuda:12.6.3-devel-ubuntu24.04
|
|
|
|
# Install Python 3, pip, and CUDA runtime compiler (tinygrad compiles kernels via NVRTC)
|
|
RUN apt-get update && \
|
|
apt-get install -y --no-install-recommends \
|
|
python3 \
|
|
python3-venv \
|
|
python3-pip \
|
|
ca-certificates && \
|
|
rm -rf /var/lib/apt/lists/*
|
|
|
|
# Install tinygrad and numpy
|
|
RUN python3 -m pip install --no-cache-dir --break-system-packages \
|
|
tinygrad==0.12.0 \
|
|
numpy
|
|
|
|
# Copy the pipeline-parallel binaries and tinygrad worker
|
|
# Build context should be the workspace root:
|
|
# docker build -f examples/pipeline-parallel-inference/Dockerfile -t <tag> .
|
|
COPY examples/pipeline-parallel-inference/target/release/pp-gpu-node /usr/local/bin/pp-gpu-node
|
|
COPY examples/pipeline-parallel-inference/target/release/pp-smoke-run /usr/local/bin/pp-smoke-run
|
|
COPY examples/pipeline-parallel-inference/pp_tinygrad_worker.py /usr/local/share/pp_tinygrad_worker.py
|
|
|
|
# Enable CUDA backend for tinygrad (override with -e DEV=CPU for CPU runs)
|
|
ENV CUDA=1
|
|
ENV WORKER_SCRIPT=/usr/local/share/pp_tinygrad_worker.py
|
|
|
|
CMD ["pp-gpu-node"]
|