Skip to content

PyTorch runtime (migration)

The PyTorch training runtime is being introduced per the PyTorch migration epics. This page covers the dependency and process-hygiene rules established in E1.

Installation

PyTorch is the only training runtime (TensorFlow/Keras were decommissioned in the migration). Dependencies are split so the control plane stays light:

  • Control plane / CLI / leader (no ML deps): pip install -e . (or -r requirements.txt).
  • Fellow (training) nodes: pip install -e ".[runtime]" (or -r requirements-runtime.txt) — adds torch + transformers + accelerate.

One framework per process

The migration kept a one-framework-per-process guard so a process never mixes torch with another DL framework. Torch runtime entrypoints claim the framework before using torch:

from silent_swarm.runtime.framework_guard import ensure_single_framework
ensure_single_framework("torch")
import torch

ensure_single_framework raises FrameworkConflictError if the other framework has already been imported or claimed in the process. Run torch and TF workloads as separate processes (the fellow already runs training in its own process).

GPU / driver expectations

  • NVIDIA driver new enough for the installed CUDA wheels (torch>=2.4; verified locally with torch 2.12 + cu130).
  • Verify visibility: python -c "import torch; print(torch.cuda.is_available(), torch.cuda.device_count())".
  • Multi-GPU placement within a node uses accelerate device_map and/or the leader placement plan (equal/weighted by free VRAM); cross-node activation exchange remains the bespoke ZMQ pipeline.
  • Set CUDA_VISIBLE_DEVICES to pin a process to specific GPUs (the fellow already manages this for the Keras runtime and will for torch).

Selecting a runtime (E2)

The control plane resolves a runtime by name, never importing a framework directly:

from silent_swarm.runtime import get_runtime, available_runtimes
runtime = get_runtime("torch")   # lazy: torch is imported only here

available_runtimes() lists known names without importing any of them. Built-in runtimes are resolved lazily; third parties/tests can register_runtime(name, factory). The contracts live in runtime/interfaces.py (ModelRuntime, StageRuntime, Compressor, Evaluator).

Wiring the fellow's training execution to get_runtime(...) is E8; E2 delivers the swappable seam, not the live fellow rewrite.

Loading & introspection (E3)

The torch runtime loads any HF causal LM and introspects its decoder generically — no per-architecture weight porting:

runtime = get_runtime("torch")
runtime.load("HuggingFaceTB/SmolLM2-135M", dtype=torch.float16, device_map="auto")
topo = runtime.topology()   # blocks, embed_tokens, final_norm, lm_head, dims

introspect_decoder (runtime/torch/loader.py) finds the block ModuleList, input embedding, final norm, and head across GPT-2 (.h + ln_f) and LLaMA-family (.layers + RMSNorm) models. Sharded safetensors load natively.

Sharding & device placement (E4)

runtime/sharding/ plans which layers go on which GPU and builds the accelerate device map:

runtime = get_runtime("torch")
runtime.shard("gpt2", gpu_free_gb=[24.0, 16.0], mode="weighted")  # or "equal"
runtime.block_devices()   # ['cuda:0', ..., 'cuda:1', ...]

equal distributes layers evenly; weighted scales by free VRAM (same policy as the leader, whose gpu_slots plan can be consumed directly via spans_from_placement). The device-map builder is generic over architecture (GPT-2 transformer.h.{i} vs LLaMA model.layers.{i}) and co-locates tied embeddings. This is intra-node placement; running a stage's forward across nodes is E6.

Compression boundary (E5)

runtime/compression/ provides NoCompression, FixedQuantization (STE), and LearnedBottleneck as nn.Modules, plus compressor_from_spec(...) (same spec schema as the Keras path) and attach_compressor(block, compressor) (forward hook). Each module separates forward (differentiable, same-shape, for training) from encode/decode (the split bottleneck wire path that sends the narrow code — int8 is ~12× smaller than fp32). wire_nbytes(x) reports the payload size for bandwidth accounting.

Cross-node pipeline (E6)

runtime/pipeline/ runs two stages in separate processes that exchange the boundary activation forward and its gradient backward over a ZMQ PAIR channel, composed as a torch.autograd.Function:

# upstream (node 0)
from silent_swarm.runtime.pipeline import PipelineBoundary, ZmqPipeChannel, run_upstream_step
channel = ZmqPipeChannel("tcp://node1:5701", bind=False)
loss = run_upstream_step(stage0_module, channel, inputs, optimizer)

# downstream (node 1)
channel = ZmqPipeChannel("tcp://0.0.0.0:5701", bind=True)
run_downstream_step(stage1_module, channel, loss_fn, target, optimizer)

A single loss.backward() upstream trains the upstream params via the returned gradient; the downstream stage trains its own params locally. Verified to match a single-process baseline (gradients diff < 1e-5). Compressed payloads (E5 codes) serialize through the same channel with fewer bytes; byte-accounting callbacks feed the existing heartbeat metrics.

Transports (cutover P2): PipeChannel has two backends — ZmqPipeChannel (direct p2p) and RelayPipeChannel (through the leader relay, the tunnel-mode default; framework-free RelaySocket matches the Keras wire format). runtime/torch/distributed_stage.py:run_torch_stage runs a GPT-2 2-stage split (segmented forward + forwarded labels + the boundary, optional split bottleneck); a 2-process run matches the single-process baseline. Wiring the fellow's multi-stage subprocess and the 2-fellow cluster run is operator-gated (operator-gated in the cutover plan).

Multi-GPU per node + shard export (cutover P3): run_stage_blocks + plan_block_devices (E4) spread one stage's blocks across the node's GPUs by free VRAM (block_devices/embed_device/head_device on run_torch_stage); a 2-GPU segmented forward matches single-GPU. The boundary already sends the narrow code (split-compression on the wire). runtime/torch/checkpoint.py exports/reconstructs per-stage shards (state_dict by layer range) — reconstruction matches the original logits. The full 4-GPU / 2-node run is the operator cluster gate.

Relay wire format

RelaySocket frames are length-prefixed multipart: !I count, then !I len plus bytes per frame. decode_multipart validates every length against the bytes actually present and rejects trailing data. Slicing past the end of a bytes returns a short result rather than raising, so an unchecked decoder turns a truncated frame into a plausible-but-wrong tensor that only fails much later, somewhere else — a decode error at the boundary is far cheaper to diagnose.

Receivers long-poll .../pop with a timeout floored at RELAY_POLL_MIN_S (0.1s), which is the endpoint's own minimum: asking for less earns a 422 that reads as a transport failure rather than the deadline it actually is.

Worker subprocess I/O

Anything that reads a subprocess's stdout and stderr must drain both concurrently. Draining stdout to EOF first and reading stderr afterwards deadlocks as soon as the child writes more than a pipe buffer (~64 KB) of warnings — which torch and transformers do on an ordinary startup: the child blocks on stderr, stops writing stdout, and the parent waits for an EOF that can never arrive. fellow/torch_dispatch.py (training worker) and fellow/inference_dispatch.py (serving subprocess) each drain stderr on its own thread for the life of the process, keeping a bounded tail of it for error messages.

Those drain threads guard the process handle rather than asserting on it. A subprocess stopped while its drain thread is still starting is an ordinary race, and python -O strips asserts anyway — the stdout drain posts its EOF sentinel so a caller waiting on the event queue is not left hanging.

Stage 1 of a distributed serving pipeline is relay-driven and takes its only instruction (unload) from stdin, parsed as JSON rather than matched as a substring — a model id or prompt containing the word used to stop the stage. A failed relay session backs off (RELAY_RETRY_MIN_S doubling to RELAY_RETRY_MAX_S) on a stop_event.wait(), so an unreachable relay does not pin a core and an unload arriving mid-backoff is still acted on immediately. A clean session resets the ladder.

Fine-tuning & calibration (E7)

runtime/finetune/ provides freeze_backbone(model, [compressor]), a calibrate(...) loop (optional amp= / grad_checkpointing=), perplexity(...), and an optional lora.apply_lora(...) (PEFT, lazily imported). Freeze the backbone and calibrate only the inserted compressors, or pass all parameters for full fine-tuning. The GPT-2 untrained→train→recovered perplexity result is reproduced end-to-end (gated test).

Evaluation & experiment harness (E8)

runtime/eval/metrics.py provides perplexity and top1_accuracy (reused by calibration). The end-to-end before/after compression experiment lives at exp/scripts/torch_gpt2_split_compression.py — the torch counterpart of the Keras hf_split_compression.py. Torch GPT-2 perplexity (50.8) matches the Keras path (50.45).

Control-plane integration & decommission (gated)

The runtime seam is get_runtime("torch") (E2). Rewiring the live fellow training execution to dispatch through it — against the leader's job/peers/relay protocol — and then removing the TF/Keras stack (E9) is deliberately not done yet: the deployed distributed path still runs on TF, and deleting it before the torch path is integrated and validated in the real cluster would break the system. The decommission is gated on the torch path being validated in the real cluster first.

Test markers

pytest.ini registers markers used by the migration:

  • torch — needs the PyTorch runtime
  • gpu — needs at least one CUDA GPU
  • slow — downloads a model from HuggingFace

Because of the one-framework-per-process rule, the torch runtime tests live under tests/unit/runtime/torch/ and auto-skip when TensorFlow is loaded (i.e. during the full Keras suite). Run them on their own, in a fresh process:

pytest tests/unit/runtime/torch                      # fast, no downloads
RUN_HF_DOWNLOAD_TESTS=1 pytest tests/unit/runtime/torch   # + real-model parity

The Keras suite (pytest tests/unit) skips them automatically.