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, plus bitsandbytes for weight quantization.

bitsandbytes is imported lazily: a fellow without it runs full-precision jobs normally and fails a quantized one with an explicit message rather than an ImportError.

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.pipe.split.torch.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).

Compute capability and dtype¶

Not every GPU can run every dtype, and the difference decides where a job can be placed. pipe/split/precision.py holds the rules; like the quantization specs it is torch-free, so the leader can apply them without a training runtime.

Compute capability	Architecture	Trains in	Example
8.0 and above	Ampere and newer	bfloat16, float16, float32	RTX A5000, A4000
7.0 - 7.9	Volta, Turing	float16, float32	GTX 1660 Super (7.5)
below 7.0	Pascal and older	float32 only	GTX 1050 Ti, GTX 1080 (6.1)

bfloat16 arrived with Ampere. It has the same width as float16 but spends more of it on the exponent, which is what keeps training stable without a loss scaler - so it is preferred wherever it exists, and its absence on Turing is a hard limit rather than a slow path.

The leader picks one dtype per job, because only it sees every node:

An explicit model.dtype in the job config wins, and any GPU that cannot run it is excluded from the placement.
Otherwise the best dtype every ready GPU supports is used. Adding a GTX 1660 Super to a cluster of A5000s therefore moves the whole job to float16 rather than failing on the 1660 at step 1. float32 runs everywhere and is the last resort: a Pascal card beside a Turing one (no half-precision format in common) runs the whole job in float32, so both pipeline stages compute in the same format - at twice the memory of float16 on the newer card.
A GPU that reports no capability is not excluded. An empty compute_cap means the fellow could not read one, not that the card is old, and excluding those would strand a cluster running an older fellow build.

The chosen dtype travels to every fellow as resolved_dtype in the job config, so all stages agree. A job that cannot be placed records placement_blocked_reason on itself naming the rejected cards:

placed 1 of 2 required workers; 1 GPU(s) excluded by dtype bfloat16:
fellow-2 GPU 0 (GTX 1660 SUPER, compute 7.5) cannot run bfloat16

Such a job stays PENDING rather than failing, because a capable node joining later is a real resolution.

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.

Hybrid models. Qwen3.5 (and Qwen3-Next) interleave three linear-attention (Gated DeltaNet) layers with every full-attention one; each block says which it is in layer_type. The segmented forward hands the causal mask only to full-attention blocks (a linear-attention block takes a padding mask, and pipeline segments carry no padding), and DecoderAdapter.new_cache(blocks) gives each serving stage a cache with the right layer per block - recurrent state or keys and values - laid out by the block's index in the whole model. A stage of linear-attention layers only builds no causal mask at all. The split may fall anywhere; tests/unit/runtime/torch/test_arch_adapter.py holds the forward, gradient and cached-decode parity for several cut points. AutoModelForCausalLM gives the text-only Qwen3_5ForCausalLM from the multimodal checkpoint. Its config keeps the decoder's numbers under text_config (loader.decoder_config reads through it wherever a layer count or head count is needed), and its weights are stored as model.language_model.* beside a vision tower and an MTP head: the stage shard loader renames stored keys with the same plain renames from_pretrained applies (transformers.conversion_mapping), so a stage reads exactly its own decoder tensors and nothing of the rest. It needs transformers 5.8 or newer (the floor in requirements-runtime.txt; a fellow gets it on swarm update). Without the optional flash-linear-attention and causal-conv1d kernels the linear layers run transformers' plain torch path: correct, slower.

Weight quantization (quantized-lora E1)¶

Quantized weights are what let a model larger than the cluster's VRAM be trained at all. Pass a spec to load or shard:

runtime = get_runtime("torch")
runtime.load("Qwen/Qwen3-8B", quantization={"bits": 8})
runtime.shard("Qwen/Qwen3-32B", gpu_free_gb=[24.0, 16.0], quantization={"bits": 4})

The spec is a small dict, normalized by pipe/split/quantization.py:

Key	Values	Meaning
`bits`	`8`, `4`	Quantization width. `16`, `32`, `None` mean no quantization.
`method`	`"int8"`, `"nf4"`, `"fp4"`, `"none"`	Alias for `bits` (+ `quant_type`), so job configs can read like the compression specs.
`quant_type`	`"nf4"` (default), `"fp4"`	4-bit data type.
`double_quant`	`true` (default)	Quantize the 4-bit block scales too.
`compute_dtype`	`"bfloat16"` (default), `"float16"`, `"float32"`	Dtype the dequantized 4-bit matmul runs in.
`threshold`	`6.0` (default)	LLM.int8() outlier threshold.
`cpu_offload`	`false` (default)	Required when the device map places any module off-GPU.
`skip_modules`	list, or omitted	Modules kept in full precision. Omitted means "the output head, resolved by introspection"; `[]` defers to transformers.

Omitting quantization loads full-precision weights exactly as before, so nothing changes for existing jobs.

The output head stays in full precision by default. Quantizing it saves one matrix and measurably hurts quality; full_precision_modules resolves its name by introspecting a meta-device model built from the config, so this works for any architecture and downloads no weights. When the head is tied to the input embedding they share one tensor, so skipping the head protects both.

Estimating whether a model fits¶

bytes_per_param and estimate_weights_gb are deliberately torch-free, so the leader can answer "does this fit?" without a training runtime:

from silent_swarm.pipe.split.quantization import estimate_weights_gb

estimate_weights_gb(32e9, None, dtype="bfloat16")   # 59.6 GB - does not fit 40 GB
estimate_weights_gb(32e9, {"bits": 8})              # 31.6 GB - fits
estimate_weights_gb(32e9, {"bits": 4})              # 15.5 GB - fits comfortably

These are weight-only estimates including bitsandbytes' per-block scales. Activations, gradients, optimizer state and the KV cache are all on top, so fitting here is necessary but not sufficient. shard uses the same arithmetic as a preflight check and refuses a placement whose weights alone exceed the free VRAM offered, rather than downloading the model and hitting an OOM part-way in.

Weight quantization shrinks VRAM and is unrelated to the boundary compression below, which quantizes the activation crossing a stage boundary and shrinks bandwidth. The two are independent and compose.

Sharding & device placement (E4)¶

pipe/split/plan.py plans which layers go on which GPU (torch-free, so the leader uses the same planner) and pipe/split/torch/device_map.py 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)¶

Compression is two concerns in two places:

Trainable compressors - NoCompression, FixedQuantization (with a straight-through estimator) and LearnedBottleneck - are nn.Modules in pipe/split/torch/compression/, with compressor_from_spec(...) and attach_compressor(block, compressor) (forward hook). Their forward is what training sees.
Stateless codecs - int2...int8 (one float32 scale, codes packed two per byte up to 4 bits), tokint2...tokint8 (the same with one float32 scale per token, for a bottleneck's code) and fp16 - live in pipe/wire/codec.py and turn a TensorFrame into a smaller one. A module's encode/decode/wire_nbytes delegate to them, and codec_from_spec maps the same spec strings to a codec. The int codec quantizes by exactly the rule FixedQuantization.forward simulates, so a decoded frame is bit-for-bit what the model trained on.

What a distributed job sends

A 2-stage job applies the compressor as part of the model. With fixed_intN that yields a quantize-dequantized activation: floats that take only 2**N values times one scale.

Wire 1 sends those floats as they are.
Wire 2 (when every fellow in the job declares wire-v2) re-quantizes them to N-bit codes plus the scale. That is lossless, so the peer decodes the identical floats and training is bit-for-bit the same - a test trains both and compares every parameter - at a fraction of the bytes (int8: about a quarter of float32). The sender verifies each round trip and sends the plain floats if one were ever not exact.

A learned bottleneck's code is narrower and quantized per token, so wire 2 carries it as N-bit codes plus one scale per token, the same lossless way; with code_bits: 16 it is float16 and crosses as it is. The split sender rounds its code with the same straight-through quantizer the whole module uses (D8), so a split boundary trains what it was configured with.

Each stage builds the whole bottleneck and keeps its half, so both halves must come from one initialisation: build_stage draws a learned bottleneck's weights from a private generator with a fixed seed. From the process RNG, the two machines (or two threads) drew two different bottlenecks, and the same job trained differently on every run. The calibration below then replaces those weights on both sides with one projection.

The learned bottleneck¶

LearnedBottleneck(hidden, k, code_bits=8) is two linear maps and a rounding step, nothing else:

code   = round_per_token(down(x), code_bits)     # k values per token, N-bit + one scale
x_hat  = up(code)                                 # back to the hidden width

The default is k = half the hidden width (a multiple of 64) with 8-bit codes: a quarter of the bytes of 16-bit activations. A quarter of the width ("bottleneck_dim" set by hand) with 8-bit codes is an eighth. code_bits takes 2-8, or 16 for float16 codes; quantized + quantization_bits mean the same as code_bits, and "quantized": false alone means float16 codes.

Why it changed. The first bottleneck wrapped each projection in a LayerNorm and started from random weights. A LayerNorm rescales every token to one norm, which throws away the magnitudes the next block depends on. A hidden state's outlier channels sit in the thousands beside values under one. A random projection also starts at chance. On the 14B run the boundary took top-1 from 56.5 % to 0.6 % before training, and training only got it back to 41 %. On Qwen3-0.6B cut at its middle layer (exp/spikes/compression/, which has the full tables):

Boundary	Bytes vs 16-bit	Top-1 kept, untrained	After 150 glue-only steps
LayerNorm bottleneck, random start, quarter width (retired)	1/4	3 %	28 %
Linear, principal components, half width, 8-bit code (default)	1/4	99 %	102 %
Linear, principal components, quarter width, float code (`code_bits: 16`)	1/4	97 %	100 %
Linear, principal components, quarter width, 8-bit code	1/8	96 %	-
Every token renormalised, nothing compressed	1	71 %	-

Models trained with the retired bottleneck (their compressor.pt holds down_norm / up_norm) are refused at load with "retrain the model".

Calibration. A projection to k of d dimensions keeps the most of a hidden state when it projects onto the activations' first k principal components, so every learned boundary starts there (pipe/split/torch/compression/calibration.py, run by runtime/torch/distributed_stage.calibrate_boundaries). After the base evaluation and before step 1, on every stage of the job:

Stage 0 runs calibration_tokens (default 16384, a boundary setting) tokens of its training data forward, without gradients.
The stage before each learned boundary accumulates the mean and covariance of its output (float64 on the CPU: nothing but a d x d matrix is kept), takes the leading k eigenvectors and sets its sender to them: down(x) = (x - mean) @ basis.
It sends the next stage the mean and the basis (in frames of 256 rows, 5 MB for a 5120-wide model), which sets its receiver to up(code) = code @ basis.T + mean. Both halves start as one projection.
A middle stage whose outgoing boundary is learned needs the activations too: the stage before it forwards them raw. A two-stage job moves only the basis.

A resumed job skips this (its snapshot holds the trained halves), and a job with a learned boundary is never re-split after a node loss: the bottleneck belongs to its boundary. Because the stage after the boundary must take part, the leader holds such a job until every fellow in it declares calibrated-bottleneck (swarm update), and loads a served model with one only onto such fellows.

What serving sends

Distributed serving applies the boundary the model was trained with. A learned bottleneck's halves load from each shard's compressor.pt. A fixed_intN boundary has no weights and is served from the manifest's compression alone: stage 0 quantizes the hidden state it sends. Before, serving skipped it and sent full-width hidden states.

When both serving fellows declare serve-codec (and speak wire 2), the leader sets transport.boundary_codec, and a quantized boundary crosses as N-bit codes plus a scale - for fixed_int8, a quarter of float32. Otherwise stage 0 quantizes in the graph and sends floats, so a peer that cannot decode codec frames never receives one. The two paths generate the same tokens; a test runs both.

Cross-node pipeline (E6)¶

Two stages run in separate processes and exchange the boundary activation forward and its gradient backward, composed as a torch.autograd.Function (pipe/split/torch/boundary.py) over a channel (pipe/split/torch/channel.py):

# upstream (node 0)
from silent_swarm.pipe.split.torch.boundary import run_downstream_step, run_upstream_step
from silent_swarm.pipe.split.torch.channel import ZmqPipeChannel
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. The channel's link counts payload bytes; the worker reports them with each step and the fellow sends them to the leader as act_bytes / grad_bytes, which the job metrics turn into act_mb_s / grad_mb_s.

Frames and steps: each tensor crosses as a TensorFrame - a JSON header with shape and dtype, then the raw bytes, the same two parts the pipeline has always sent. bfloat16 goes as its raw 2-byte values (dtype "bfloat16"); before this, a bf16 activation could not be sent at all, so 2-stage jobs on an all-Ampere cluster (which resolves to bf16) failed at the first step. Every forward frame carries the upstream's step and the reply echoes it; an upstream that receives a reply for another step raises instead of applying a stale gradient. Older peers ignore the field and their step-less frames are accepted.

More than two stages (XL1). SplitSpec.boundaries holds all N - 1 cuts and SplitSpec.layer_range(i, n) returns stage i's blocks for any N; boundary_compression holds one compression profile per boundary (compression_at(i)). StageSpec.receive_boundary / send_boundary name the boundary a stage reads from and writes to, so a middle stage has both. A stage of such a pipeline reaches its neighbours through two routes: StageAssignment.prev_route / next_route from the leader, which become WorkerSpec.prev_transport / next_transport (route_to_prev() / route_to_next()). Every pair of neighbours gets relay channels of its own ("{up}-{down}-act", "{down}-{up}-grad"). A middle stage binds the link to the stage before it and dials the stage after it, so one port pair per fellow still suffices. A two-stage job is sent exactly what it was before: the new keys stay unset.

The middle stage (XL2). build_stage builds any stage of the pipeline. The first stage holds the embeddings and the last one the final norm and the head. A middle stage holds neither, but has two compressor halves on its bundle: the receiver of the boundary it reads from and the sender of the one it writes to. Each learned bottleneck is seeded with its boundary index, so both of its stages build the same one, and calibrated before step 1 (above). train_stage runs a middle stage between two channels. Per step it forwards the labels, runs receiver -> blocks -> sender, and sends the result through its own PipelineBoundary. When the reply comes back it backpropagates and answers the stage before it with the same loss and the gradient of the activation it received:

stage 0 --labels, code--> stage 1 --labels, code--> stage 2 (loss)
        <--loss, grad---          <--loss, grad---

So every link speaks the two-stage wire format, and wire v2's codec applies to each link by its own boundary's profile. The held-out evaluation passes run through the middle the same way, and every stage reports the last stage's numbers. With a frozen backbone, a stage whose boundaries have no learned bottleneck has nothing to train and only relays. In the executor, a middle stage binds the link to the stage before it and dials the one after (prev_transport / next_transport), and reports the link it sends activations on. Its shard carries compressor-recv.pt and compressor-send.pt; the end stages keep compressor.pt. The fellow uploads them with the stage's shard (POST /api/v1/models/{id}/shards/{i}/compressor/{name}, only those three names, before the weights): until then they stayed on the fellow that trained them, and a model served anywhere else ran its boundaries without the trained compressor. Serving reads them with weights_only. The servable manifest records the trained boundaries and each one's compression. Three- and four-stage runs match a single-process run's loss and every block's gradient (tests/unit/runtime/torch/test_n_stage_training.py).

Micro-batched schedules (XL4). With micro_batches = M > 1, train_stage splits each step's batch into M micro-batches and runs the op list pipe/split/schedule.stage_schedule gives its stage. GPipe runs every forward, then every backward. 1F1B runs N - i - 1 warm-up forwards on stage i, then one forward and one backward in turn, then the remaining backwards, so stage i holds at most N - i micro-batches' activations. A forward sends the micro-batch on. The last stage computes each micro-batch's loss, scaled by 1 / M, backpropagates it at once and replies. A backward takes its reply, backpropagates it into the accumulating gradients, and on a middle stage answers the stage before with the gradient of what it received. Each stage takes one optimizer step per step. Frames carry a micro id beside the step and replies are matched by (step, micro), so a reply to the wrong micro-batch raises (Session.recv_reply(micro=...); the upstream closes the step with end_step). Both schedules answer micro-batches in ascending order on every link and never wait on a send, so they run without deadlock. The accumulated gradient equals the full batch's, for two to four stages and both schedules (tests/unit/runtime/torch/test_n_stage_training.py). recompute wraps each block in torch.utils.checkpoint while training.

Long sequences. A training stage switches its model to the swarm_sdpa attention implementation (pipe/split/torch/attention.py, registered in transformers' AttentionInterface and AttentionMaskInterface), so the layers keep building no mask. On CUDA without Flash Attention (compute capability below 8.0), a call with grouped-query heads expands keys and values to the query heads and calls SDPA without enable_gqa: CUDA takes GQA only in the Flash and math kernels, so the expansion is what lets the memory-efficient kernel run instead of math's seq x seq scores. Everywhere else it is transformers' own SDPA path. The last stage's loss (runtime/torch/head_loss.py) scores loss_chunk_tokens(vocab) tokens at a time, each chunk under torch.utils.checkpoint, in fp32, and returns the mean in the stage's compute dtype - the value that crosses the wire back. The held-out evaluation scores chunk by chunk too. pipe/split/memory.py prices all of it (Workload.efficient_attention, chunked_loss, recompute; StageMemory.peak_gb, the one block a card runs again, paid once per card in plan_balanced_gpu_layers). The leader uses those prices only when every fellow declares long-context, and turns train.recompute on when a job fits only with it (auto_settings.choose_recompute).

Because the micro-batches split the batch, a schedule never holds more activations than one whole batch did, so the leader's memory check needs no change. The relay buffers 64 frames per channel by default (SWARM_RELAY_QUEUE_MAX), which is why micro_batches stops at 32. The leader applies micro-batches only when every fellow declares microbatch, and records the result in the job's pipeline_schedule (with the bubble fraction).

Mixture-of-experts models (XL9). Mixtral, Qwen3-MoE and gpt-oss split, train and serve through the same generic adapter as dense models: a segmented forward equals the whole model's, three-stage training matches a single process's loss, and a ring generates what model.generate does (tests/unit/runtime/torch/test_moe.py). Two things are specific to them:

Loading. transformers keeps a layer's experts fused in two 3-D weights (experts.gate_up_proj, experts.down_proj), while a checkpoint stores them one by one (experts.<e>.w1.weight ...). The shard loader follows transformers' own conversion rules (_CheckpointMap). It stacks the experts (MergeModulelist) and joins gate and up (Concatenate), one fused weight at a time, and downloads only the files holding the stage's experts. Any other conversion falls back to the offloaded load.
Quantizing. bitsandbytes replaces only nn.Linear, so a "4-bit" MoE stage used to keep its experts, nearly all of its weights, in 16 bits. Each fused weight now becomes a moe_quant.ExpertStack: every expert's matrix in NF4 (or int8). Indexing an expert, which is how the eager experts forward reads one, dequantizes only that expert for its matmul. The model's experts are switched to that eager forward. LoRA trains attention by default, so the router and the experts stay frozen and no load-balancing loss is needed.

The leader balances a pipeline of more than two stages by parameters (plan.cost_balanced_split) when its layers differ in size, such as dense layers before MoE ones. Its memory check scales the MLP's activations by the experts a token is routed to (memory.moe_ffn_scale), an estimate the XL0 spike measures.

Transports: a PipeChannel encodes tensors and hands the bytes to a silent_swarm.pipe.link link - ZmqLink (direct p2p) or RelayLink (through the leader relay, the tunnel-mode default). ZmqPipeChannel and RelayPipeChannel build the matching link. Every link passes one shared contract suite (tests/unit/pipe/link/contract.py). Building a stage, then training it. The model surgery is one call, with no optimizer, loop or channel in it:

from silent_swarm.pipe.split import Placement, SplitSpec, StageSpec
from silent_swarm.pipe.split.torch.stage_builder import build_stage

spec = StageSpec(
    stage_index=0,
    split=SplitSpec(split_layer=12, compression={"method": "learned_bottleneck", "learned_dim": 64},
                    lora={"enabled": True, "r": 8}),
    placement=Placement(device="cuda:0", block_devices=["cuda:0", "cuda:1"]),
)
bundle = build_stage(model, spec)   # this stage's blocks, placed; its compressor half;
bundle.trainable_parameters         # exactly what it optimizes (never empty)

runtime/torch/distributed_stage.py trains a bundle: train_stage(bundle, channel, tokens=...) is the AdamW loop (forwarded labels, the boundary, (loss, gradient) back), and run_torch_stage(...) keeps the keyword interface the executor uses - build the spec, build the stage, train it. The spec and the builder are the same API as the swarmpipe library's L1. A per-block device list shorter than the stage's blocks is an error; it used to leave the extra blocks on the CPU. 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.

Loading only one stage's slice¶

Each node used to load the whole model and then move its own slice onto the GPU. That is nearly free when the requested dtype matches the checkpoint's, because safetensors is memory-mapped and untouched pages never become resident. It costs a full second copy the moment a conversion is needed - and E2 made that the normal case for a heterogeneous cluster, since a pool containing a Turing card resolves float16 against bfloat16 weights.

load_stage_shard materialises only this stage's tensors:

model, topology = load_stage_shard(
    "Qwen/Qwen3-8B",
    layer_start=0, layer_end=18, is_first_stage=True, is_last_stage=False,
    dtype=torch.float16,
)

Measured on a 1.19B model, host RSS for one stage:

requested dtype	full load	shard load
matches the checkpoint	1.11 GB	1.10 GB
needs a conversion	4.45 GB	2.23 GB

Stage output is byte-identical either way - this is a pure memory optimization, and the tests assert equality rather than closeness.

Quantized weights load the same way (XL3). Before anything is read, the empty model's linear layers become bitsandbytes ones. This is transformers' own replace_with_bnb_linear, with the job's skip list (by default the output head) left in full precision. Each owned weight is then quantized as it is read, exactly as from_pretrained quantizes it: a GPT-2 Conv1D weight is turned around first, and moving the new Params4bit / Int8Params to the device is what quantizes it. One tensor is in full precision at a time, so a stage never holds its 16-bit copy. On a fellow with CUDA each weight is quantized on the card and its packed bytes moved back to the host until the stage is placed (quantize_device). bitsandbytes' CPU kernels are built for newer vector instructions than an older host may have: a GTX 1050 Ti fellow died on them without a word. Every torch subprocess enables faulthandler, and the fellow names the signal a serving process was killed by, so a native crash is reported where it happened. Its output is bit-identical to a full from_pretrained quantized load, for NF4 and int8, rotary models and GPT-2 alike. A quantized stage comes back in eval mode, as from_pretrained hands one back, since bitsandbytes' 8-bit layer computes differently in training mode.

Only the files that hold this stage's tensors are downloaded. The loader first fetches a sharded checkpoint's model.safetensors.index.json, then exactly the files whose tensors the stage needs (owned_files). A stage of a 120B model fetches its share, not 240 GB, and keeps no offload folder. The leader shows each stage's share in the job plan's memory.downloads, next to the fellow's free disk (swarm train --plan prints it). It is estimated from the checkpoint's headers: the stage's layers, plus the embeddings on the first stage and the head on the last. Before downloading, ensure_disk_for asks the Hub for the sizes of exactly those files, skips the ones already in the local cache (try_to_load_from_cache), and raises with the numbers when the rest plus 512 MB will not fit the cache's filesystem. Offline or gated, it checks nothing. The leader leaves a node out on the same arithmetic (dispatch.plan_job), from the free disk the fellow reports at its cache (fellow/hardware.model_cache_dir).

The old offloaded strategy remains as a fallback. When a checkpoint does not name a tensor the way the model does (the manual path's NotLoadableByName), a quantized stage loads through from_pretrained with the non-owned modules sent to disk through build_stage_device_map. That still reads every tensor on its way to the offload folder, and it runs only when the worker was given an offload_dir.

The device map has to be a complete partition. build_device_map places a whole model across a node's GPUs and its default-plus-overrides shape is harmless there, because every value is a GPU index. build_stage_device_map emits one entry per module, because a "disk" parent with an on-device child makes accelerate's dispatch fail with ValueError: weight is on the meta device.

Two sharp edges, both load-bearing. rotary_emb belongs to both stages - the downstream one recomputes the rotary context from absolute positions - and with tied weights the last stage must load the input embedding it never runs, because the checkpoint stores no lm_head.weight and the head aliases the embedding. Without that rule the manual path silently produces NaN. This is the one place owned_param_keys cannot be reused unchanged, which is why loading has its own stage_load_keys: exporting deliberately does not duplicate that tensor into both shards.

init_empty_weights leaves buffers materialized, which is why a rotary model's non-persistent inv_freq - absent from the checkpoint - needs no special handling.

LoRA across the two stages¶

run_torch_stage(..., lora={"enabled": True, "r": 8}) trains adapters instead of whole parameters on both stages of a split model - the mode that makes a model larger than the cluster's optimizer budget trainable across two machines.

Adapters are injected into the decoder blocks in place, so the segmented forward runs the wrapped modules without knowing about it. PEFT is handed the whole model, but each stage optimizes only the adapters inside the blocks it owns, reusing the scoping the non-LoRA paths already had. The two stages therefore train disjoint halves of one adapter set.

Both stages resolve target_modules from the same architecture, so they agree by construction. They do not agree on the random lora_A initialisation of the half they do not own, and that is harmless: lora_B starts at zero, so an untrained adapter is exactly an identity and merges as a no-op.

A stage that ends up with nothing to train now fails instead of running. A target_modules that matches nothing on this stage's blocks gives an optimizer over an empty list, which steps every iteration, reports a loss, and leaves the weights untouched.

Shards record how they were trained. Each stage writes a lora_fingerprint into its shard manifest - a hash of the effective rank, scaling, dropout and resolved targets - and reconstruct_from_shards refuses to merge halves whose fingerprints differ:

Refusing to merge stage shards trained with different LoRA settings
(stage-0=6f1c…, stage-1=none). Retrain the stages with one adapter configuration.

Merging mismatched halves would otherwise produce a model that loads cleanly and is quietly wrong. Shards written before the field existed carry no fingerprint and are treated as unknown rather than as a mismatch.

Relay wire format¶

RelayLink bodies 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)¶

pipe/split/torch/finetune/ provides freeze_backbone(model, [compressor]) and lora.apply_lora(...) (PEFT, lazily imported) - what trains is part of the split. The calibrate(...) loop (optional amp= / grad_checkpointing=) is in runtime/torch/finetune.py. The GPT-2 untrained→train→recovered perplexity result is reproduced end-to-end (gated test).

Which parameters train¶

A job runs in one of four modes, chosen by finetune_mode from the job config:

Mode	Requested by	Trains	Rough optimizer cost
`full`	the default	every parameter	~16 bytes per model parameter
`compressor_only`	`model.freeze_backbone` with a compression boundary	only the inserted bottleneck	~4 bytes per model parameter
`lora`	`model.lora.enabled`	low-rank adapters, plus the compressor if present	the model at ~2 bytes per parameter, plus ~16 per adapter parameter
`none`	`model.freeze_backbone` with a `fixed_intN` boundary	nothing - the quantizer has no parameters	none: no data is read and no step runs

full is out of reach for anything large: AdamW's two moments and the fp32 master copy mean a 32B model needs around 512 GB. lora is what makes such a model trainable, because the frozen weights carry no optimizer state at all - only the adapters do, and there are a few hundred times fewer of them.

lora takes precedence over freeze_backbone. Asking for adapters and getting a frozen backbone that trains nothing but a compressor would be a silent no-op.

Held-out evaluation of a pipeline job (`train.eval_batches`)¶

A pipeline job measures itself on data it did not train on, three times:

Pass	When	Measures
`base`	before step 1, boundary compressor bypassed	the pretrained model as loaded (LoRA's `B` starts at zero)
`calibrated_boundary`	before step 1, through the compressor (only with one), after a learned bottleneck's calibration	what the boundary costs before training
`trained`	after the last step	what training won back

Each pass reports the mean token cross-entropy (loss), its perplexity, the top1 next-token accuracy and the number of tokens scored. The upstream stage takes train.eval_batches batches (training's batch_size x seq_len) from the dataset's validation split - WikiText-103's by default - and sends them through, forward only and with dropout off; the downstream stage scores them and replies with the summed loss and the hit count, so both stages end up with the same numbers (runtime/torch/distributed_stage.py). A bucket file's validation split is its last rows, one in twenty (data/file_dataset.HELD_OUT_EVERY), and its training split is every row before them, so the job never trains on what it is scored on. A file of fewer than 20 rows has no validation split and evaluates nothing rather than score on its own training data; the upstream still opens every pass, with 0 batches, so its peer never waits. Until this split existed a bucket file's "validation" was its first tokens, the ones the first steps trained on, and a run's trained pass looked better than it was.

Both stages must run the same passes, so the leader sets eval_batches - 16 unless the job says otherwise, 0 to switch it off - only when every fellow in the job declares pipeline-eval. The fellow reports the result with COMPLETED; the leader keeps it on the job as config.evaluation (known passes and numbers only - it comes from a node), swarm jobs show prints it as a table, and the exported model's manifest carries it under training.evaluation.

The training buffer is sized to the schedule too: the upstream fetches at least num_steps x batch_size x (seq_len + 1) tokens, because the batches are fixed windows over that buffer and a run that outgrew dataset_config.real_min_tokens (200k by default) used to stop early - 2,000 steps of 4 x 128 need about a million.

Half-precision storage of a frozen backbone (`model.weight_dtype`)¶

A lora or compressor_only stage never updates its base weights, so it does not need them in float32 - only its arithmetic does. With model.weight_dtype set to "float16" (or "bfloat16"; swarm train --weight-dtype float16), the stage builder converts this stage's frozen nn.Linear / nn.Embedding (and GPT-2 Conv1D) weights before they move to the GPU (pipe/split/torch/half_storage.py) and gives each layer a forward that upcasts on the fly. The weights take half the VRAM; activations, gradients, the LoRA adapters, norms and the boundary compressor stay in the job's compute dtype.

It exists for cards that cannot train in half precision: a GTX 10xx (Pascal) runs float16 math at a small fraction of its float32 speed, and pipeline stages train without loss scaling, so a mixed Pascal/Turing cluster computes in float32. Storing in float16 still halves the model - Qwen3.5-4B goes from 16.8 GB to 8.4 GB of weights, which is what lets it train with LoRA over two GTX 1050 Ti and two GTX 1660 SUPER (about 19 GB free in all).

The upcast is a custom autograd function, not weight.float(): with the plain version autograd would keep a float32 copy of every weight until backward. The function saves the float16 weight and upcasts again in backward, in row chunks for a large matrix (a 248k-token head), so the transient copy stays small. Against a float32 model with the same weights rounded to float16 the loss and the adapter gradients agree to float32 round-off; a bf16 checkpoint converts to float16 exactly, so for Qwen3.5 they are identical.

The option is per stage and changes nothing on the wire, but a fellow from before it would load the weights in full where the leader's memory check counted half: the leader holds such a job until every fellow in it declares the half-weights capability (swarm update). A full fine-tune ignores it (its weights train and need their float32 copy), and so does a single-fellow job, which loads the whole model onto its GPU before the backbone is frozen.

none is how to load a model and add compression without training it: the job loads the model, puts the quantizer on the boundary and exports the servable model, with the boundary in its manifest so serving applies it. Both stages of a distributed job derive this from the same SplitSpec.trains_nothing, so both skip the dataset, the link and the steps. Before, a single-node job tokenized its whole dataset to run zero steps, and a distributed one failed with "Stage 1 has no trainable parameters".

swarm train --config job.json --pretrained Qwen/Qwen3-8B \
    --compression fixed_int4 --freeze-backbone

A learned bottleneck has weights. It starts calibrated (above), so a short compressor_only run fine-tunes it rather than training it from nothing.

{
  "model": {
    "pretrained": "Qwen/Qwen3-8B",
    "lora": {"enabled": true, "r": 8, "alpha": 16, "dropout": 0.05}
  }
}

target_modules may be given explicitly; omitted, it is resolved from the model's own structure by default_target_modules, which reads the leaf module names inside the decoder blocks (all of them: a hybrid model's first block may be linear attention only) and matches them against the known attention-projection groups:

Family	Adapted
Llama, Qwen, Mistral, Gemma	`q_proj`, `k_proj`, `v_proj`, `o_proj`
GPT-NeoX, Falcon	`query_key_value`
GPT-2	`c_attn` (fused qkv; `c_proj` is skipped because the MLP reuses the name and PEFT matches on the suffix)
MPT	`Wqkv`
Phi-3	`qkv_proj`
Qwen3.5 (hybrid)	`q_proj`, `k_proj`, `v_proj`, `o_proj`, plus the linear-attention `in_proj_qkv`, `in_proj_z`, `out_proj`
Qwen3-Next (hybrid)	the attention group plus `in_proj_qkvz`, `out_proj`

An unrecognised architecture falls back to every two-dimensional weight in the block rather than failing, and GPT-2's Conv1D projections are matched on weight shape rather than on nn.Linear, so the resolution is not tied to a model type string.

What a LoRA job exports¶

Two ways, chosen by model.lora.export (or swarm train --lora-export):

Mode	Shard holds	Needs PEFT to serve	Size of a 32B fine-tune
`merged` (default)	the whole model, adapters folded in	no	~64 GB
`adapters`	the adapters only	yes	a few hundred MB

Merged is the default because it needs nothing at serving time: the shard has the model's own parameter names and the inference path never learns that LoRA was involved.

Over a quantized base the adapters always ship on their own (QLoRA-style), whatever export says, and the shard carries no base weights at all (runtime.torch.executor.base_is_quantized). A stage quantized as it loads holds its frozen weights packed in 4 or 8 bits: PEFT cannot fold an adapter into such a layer - it picks its bitsandbytes-aware LoRA layer only for a model loaded through transformers' own quantization, so the merge added a full-size delta to the packed buffer and a 14B run failed at its export, after 2,000 steps, with "size of tensor a (13107200) must match the size of tensor b (5120)" - and the packed bytes are nothing a shard's reader could take either. Serving quantizes the same checkpoint itself and keeps the adapters beside it, so it computes exactly what was trained. The fellow uploads the adapter files (POST /api/v1/models/{id}/shards/{i}/adapter/{name}, only adapter_config.json and adapter_model.safetensors) before the shard's weights, whose last chunk is what makes the leader list the model as ready.

adapters ships the delta. A LoRA run freezes the pretrained weights, so a merged servable is mostly bytes that did not change. In this mode the shard's weights.pt is written empty - weights_included: false in its manifest - and the adapters go beside it:

<model_id>/shards/stage-0/
  weights.pt                        # empty; the base weights did not change
  manifest.json                     # weights_included: false, lora_fingerprint
  adapter/adapter_model.safetensors
  adapter/adapter_config.json

reconstruct_from_shards handles both. It skips a shard with no weights, then looks for adapter/ directories, unions them, and attaches them to the base model pulled from base_model. A pipeline stage saves only the adapters for the blocks it ran, so reassembly is exactly that union - and because the halves must describe one adapter set, reconstruction refuses to combine halves that disagree on rank, scaling or dropout. Pass merge_adapters=False to keep them separable instead of folded in.

The servable manifest records which way it went, under training.lora, along with the rank and the resolved target modules.

Evaluation & experiment harness (E8)¶

runtime/torch/eval.py provides perplexity and top1_accuracy (reused by calibration). Datasets (WikiText, uploaded files, HF datasets) are tokenized by runtime/torch/data/. 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¶

pyproject.toml ([tool.pytest.ini_options]) 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.