Pipeline Parallelism¶

Table of Contents¶

• Concept
• Zero-Config Startup
• Layer Partitioning
• Activation Transfer Between Stages
• Micro-Batching and GPipe Schedule
• Pipeline Bubbles
• Per-Stage KV Cache Management
• Transport Backends
• Configuration
• Composing with Tensor Parallelism (2D PP × TP)
• LoRA with PP

Concept¶

Pipeline Parallelism (PP) splits a model vertically by layers: each GPU (pipeline stage) is responsible for a contiguous range of transformer layers. The stages are arranged in a pipeline so that when stage i finishes processing a batch and passes the resulting activations to stage i+1, it can immediately start processing the next micro-batch.

This is different from Tensor Parallelism, where every GPU participates in every layer. In PP, each GPU only runs the layers assigned to its stage.

graph LR
    subgraph "Pipeline (3 stages, 4 micro-batches)"
        subgraph S0["Stage 0 (layers 0-15)"]
            S0_MB0[MB 0]
            S0_MB1[MB 1]
            S0_MB2[MB 2]
            S0_MB3[MB 3]
        end

        subgraph S1["Stage 1 (layers 16-31)"]
            S1_MB0[MB 0]
            S1_MB1[MB 1]
            S1_MB2[MB 2]
            S1_MB3[MB 3]
        end

        subgraph S2["Stage 2 (layers 32-47)"]
            S2_MB0[MB 0]
            S2_MB1[MB 1]
            S2_MB2[MB 2]
            S2_MB3[MB 3]
        end

        S0_MB0 -->|activations| S1_MB0
        S0_MB1 -->|activations| S1_MB1
        S0_MB2 -->|activations| S1_MB2
        S0_MB3 -->|activations| S1_MB3

        S1_MB0 -->|activations| S2_MB0
        S1_MB1 -->|activations| S2_MB1
        S1_MB2 -->|activations| S2_MB2
        S1_MB3 -->|activations| S2_MB3
    end

The stages overlap: while stage 1 is processing micro-batch 0, stage 0 is already processing micro-batch 1, so no stage sits idle.

The repository currently exposes two PP entrypoints:

• mlxcel generate for single-machine, in-process pipeline execution on supported text families (Llama, Mistral, Mixtral, DeepSeek V3, Llama 4 Scout text tower, GPT-OSS, Gemma 3, Gemma 4, Qwen3, Qwen3.5, Qwen3-Next, GLM4 family, and the hybrid Mamba + Transformer families Jamba and Nemotron-H).
• mlxcel serve for coordinator-driven distributed serving.

Stage-local execution is intentionally split into:

• shared stage helpers for token embedding, per-layer execution, final norm, and lm-head emission
• family-specific backends that only know how to load and execute their own block type

This keeps the transport protocol, worker loop, and remote runtime unchanged as new text families are added.

Supported families¶

Each supported family carries a stable textual name that mlxcel-server advertises during cluster handshake (supported_families() in src/distributed/pipeline/stage_executor/mod.rs). A coordinator and its stage workers exchange these names at connection time; a cluster whose workers advertise a subset that does not cover the coordinator's declared family will fail to start with pipeline capability mismatch rather than silently corrupt generation.

Representative notes:

• Gemma 4 keeps the transport and runtime protocol unchanged by packing downstream per_layer_inputs into the stage hidden-state payload.
• Gemma 4 also has KV-shared layers, so representative 2-stage splits must keep each shared source layer in the same stage as its consumer. For gemma-4-e2b-it-4bit, the validated split is 0-12,13-34.
• The current representative GLM4-family validation uses glm4-flash-4bit, which is detected as GLM4 MoE Lite and therefore exercises the explicit prefill causal-mask path used by the full-model loader.
• GLM MoE DSA reuses the shared DeepSeek V3.2 stage-model loader path, but the repository does not currently ship a local glm5-4bit checkpoint for representative real-model parity.
• DeepSeek V3 drops the last num_hidden_layers entry (the MTP trailer) when counting transformer blocks; the pipeline partitioner sees num_hidden_layers - 1 layers and must be given a split inside that range. sanitize_weights runs once on the full weight map before partial loading so the split kv_b_proj decomposition and stacked MoE expert weights are consistent on every stage.
• Llama 4 Scout exposes only the text language tower through pipeline parallelism. Vision encoder stage splitting is out of scope for this rollout. The stage executor uses the legacy TransformerBlock::forward(x, cache, mask) path with the standard KVCache so the transport wire protocol stays unchanged; the full-model iGQA ChunkedKVCache optimisation is not used in pipeline mode.
• Jamba and Nemotron-H are hybrid Mamba + Transformer architectures. Every Mamba / Mamba2 block carries persistent conv_state + ssm_state that is updated per-token; splitting a single SSM block across stages is not supported because the wire protocol only transmits hidden states. LayerFilter::layer_range is block-granular, so stage boundaries can land between any two decoder blocks (attention ↔ mamba is fine) but never inside a block. Both hybrid executors load the full model on every stage and execute only the layer range assigned to the local stage; partial-weight loading for hybrid models is outside the current implementation. The correctness invariants on SSM stage boundaries are independent of this execution mode.
• Kimi K3 (kimi_k3) is not a supported pipeline family. Its full 93-layer checkpoint occupies about 1.4 TB at 4 bits, more than any single Apple Silicon host holds, but there is no stage executor for it, so --pp-size and --pp-layers refuse a kimi_k3 model. The partition planner prices its dense, MLA, and MoE layers separately. The family runs single-process only, and a layer-truncated local copy (a lower text_config.num_hidden_layers in config.json) is the single-host check for loading, shapes, and finite logits; its output text is not meaningful.

Zero-Config Startup¶

Zero-config bring-up starts a multi-machine pipeline from CLI flags alone, with no hand-written cluster TOML and no per-peer --node-id. The coordinator plans the cluster and pushes one stage assignment (stage index, model identity, layer range, cluster topology) to every --pp-peer process; a peer does not need --distributed-config or --node-id. The flags are available on mlxcel-server and the matching mlxcel serve subcommand:

The coordinator assigns stage indices deterministically (peer addresses are sorted lexicographically before partitioning), so rerunning the bring-up with the same set of peers produces a byte-identical TOML that is safe to commit to version control for reproducibility. The underlying planner is implemented in src/distributed/cluster_init.rs and the TOML is consumed by the existing ServerStartupConfig distributed resolution path, so the zero-config path shares runtime code with the manual TOML path. It is not a parallel prototype.

Cluster Secret and Transport Protection¶

Every pipeline role (--pp-auto, --pp-peer, and a --distributed-config with pipeline_parallel_size >= 2) and every disaggregated role refuses to start without a cluster secret. Generate it once and copy it to every host:

openssl rand -hex 32 > cluster.secret

The secret is read from --cluster-secret-file, then MLXCEL_CLUSTER_SECRET_FILE (a path), then MLXCEL_CLUSTER_SECRET (the value), then [cluster] secret_file in the cluster TOML. It is used for three things: a mutual HMAC-SHA256 handshake on every TCP connection, per-connection keys (HKDF-SHA256) for the AES-256-GCM encryption and authentication of every frame after the handshake, and the signature on discovery beacons.

The handshake is cluster protocol version 5. A node running a build older than v0.8.0 cannot join a v0.8.0 cluster, so upgrade every node together. The server's accept, at-capacity, or refusal verdict is itself authenticated, so a party on the path cannot turn one verdict into another. Frames carry an implicit sequence number, so a replayed, reordered, dropped, or injected frame fails authentication and closes the connection.

Example: LAN bring-up with static seeds¶

# Machine B (peer): listens for its stage assignment on 192.168.1.11:19001.
mlxcel-server --model ./models/llama-3.2-1b-4bit \
  --pp-peer \
  --cluster-secret-file ./cluster.secret \
  --cluster-control-addr 192.168.1.11:19001 \
  --host 0.0.0.0 --port 18081 --no-warmup

# Machine C (peer): listens on 192.168.1.12:19001.
mlxcel-server --model ./models/llama-3.2-1b-4bit \
  --pp-peer \
  --cluster-secret-file ./cluster.secret \
  --cluster-control-addr 192.168.1.12:19001 \
  --host 0.0.0.0 --port 18082 --no-warmup

# Machine A (coordinator): plan the cluster and push one assignment to each peer.
mlxcel-server --model ./models/llama-3.2-1b-4bit --alias llama-auto-pp \
  --pp-auto 2 \
  --cluster-secret-file ./cluster.secret \
  --cluster-peers 192.168.1.11:19001,192.168.1.12:19001 \
  --cluster-control-addr 192.168.1.10:19000 \
  --cluster-config-out ./.mlxcel/cluster.toml \
  --parallel 2 --max-batch-size 2 --pp-micro-batch-size 2 \
  --host 0.0.0.0 --port 18080

The coordinator runs no stage itself, so --pp-auto 2 needs two --pp-peer processes. The assignment exchange has three phases: each peer accepts or rejects its offer, the coordinator sends a commit only after every peer accepted, and it sends an activate only after every peer committed. A peer writes its cluster TOML and starts its stage on activate. If any peer rejects, nothing starts anywhere, and a peer whose checkpoint does not match the offered model (config hash, weight file names and sizes, layer count) refuses with a typed error. The coordinator starts serving HTTP after all peers committed.

Example: Opt-in mDNS discovery¶

When peers cannot be enumerated ahead of time, the coordinator can wait for peer announcements over UDP broadcast instead:

# Each peer announces itself with signed UDP beacons on the default beacon port.
mlxcel-server --model ./models/llama-3.2-1b-4bit \
  --pp-peer \
  --cluster-secret-file ./cluster.secret \
  --cluster-discovery=mdns \
  --cluster-name studio-pp

# Coordinator waits until 2 peers have announced, then plans the topology.
mlxcel-server --model ./models/llama-3.2-1b-4bit \
  --pp-auto 2 \
  --cluster-secret-file ./cluster.secret \
  --cluster-discovery=mdns \
  --cluster-name studio-pp \
  --cluster-config-out ./.mlxcel/cluster.toml

mDNS discovery is off by default. When it is enabled, the coordinator accepts only beacons signed with the cluster secret for its --cluster-name, so unrelated broadcast traffic from other applications on the same LAN is ignored.

Dry run and operator errors¶

--dry-run prints the resolved topology summary and the path of the emitted TOML without starting the model or the HTTP listener:

$ mlxcel-server -m ./models/llama-3.2-1b-4bit \
    --pp-auto 2 \
    --cluster-peers 192.168.1.11:19001,192.168.1.12:19001 \
    --dry-run
Cluster: mlxcel-cluster
  TP size: 1, PP size: 2
  Nodes (3):
    - coordinator @ 127.0.0.1:19000 [hybrid]
    - stage-0 @ 192.168.1.11:19001 [pipeline_stage]
    - stage-1 @ 192.168.1.12:19001 [pipeline_stage]
Emitted cluster TOML at: /home/op/.mlxcel/cluster.toml

The dry-run exits with a non-zero status code when:

• --pp-auto N is set but fewer than N peers were resolved (static or mDNS).
• The coordinator control address collides with the HTTP listen address.
• A TCP port in the generated topology is already bound by another process.

Failure messages include actionable remediation hints (which flag to change, which subnet/firewall rule to inspect, which fallback to use).

Determinism contract¶

The cluster TOML is guaranteed to be byte-identical across reruns with the same topology. Specifically:

• Peers are sorted lexicographically by textual address before stage assignment, so --cluster-peers A,B and --cluster-peers B,A produce the same TOML.
• Node keys are emitted in a fixed order (id, address, role, stage, rank), and the coordinator entry always appears before pipeline stages.
• The header comment is a static string. It does not embed the coordinator's hostname, PID, or clock.

The determinism guarantee is covered by plan_cluster_emits_byte_identical_toml, plan_cluster_reordered_peers_produce_same_topology, and render_deterministic_toml_is_stable in src/distributed/cluster_init_tests.rs.

Migration from the manual TOML path¶

The manual TOML / scripts/benchmark_pipeline_remote_rollout.sh write-config workflow is not removed. The zero-config flags compose additively: an operator who already has a working cluster TOML can keep using --distributed-config and ignore --pp-auto. Mixing the two in a single invocation is rejected at startup with an explanatory error.

Layer Partitioning¶

Layer partitioning is implemented in src/distributed/pipeline/partition.rs (core types and validators), partition_balance.rs (byte-accurate constrained balancer), partition_profile.rs (model-aware ModelProfile builder), and partition_quality.rs (estimated-vs-actual memory reporting).

Auto Partitioning (default)¶

The auto-partitioner minimises the maximum per-stage byte load across stages, subject to four hard constraints. The balancer is byte-accurate and constraint-aware, so the default path works for MoE models and for Gemma 4 / Gemma 4 VLM without operators hand-specifying --pp-layers.

Algorithm:

1. Read config.json and build a ModelProfile with real per-layer byte weights and adjacency constraints. MoE layers are counted with their expert blocks; Gemma 4 KV-shared consumer layers carry the double-wide MLP; quantisation bit-width scales the estimate.
2. Reserve embedding on stage 0 and lm_head on the last stage. Remaining stages see layer-only byte budgets.
3. Binary-search the smallest T (max stage byte load) for which a K-stage split exists that respects every per-stage budget and never cuts a LayerAdjacencyGroup. Stages with tighter budgets receive fewer bytes; spilled weight flows to roomier stages.
4. Reconstruct layer ranges for the feasible T. Assign embedding to stage 0 and lm_head to the last stage, with contiguous layer ranges to minimise activation transfer.

Model-aware per-layer byte accounting¶

A uniform layer_param_bytes per layer would skew MoE families: expert-heavy layers would be estimated the same as the dense layers beside them, the balancer would hand them to the wrong stage, and the heavy stage would run out of memory at load time. The profile builder instead recognises:

• Mixtral / Qwen 2.5 MoE / GLM 4 MoE / GLM 4 MoE Lite / GLM MoE DSA / GPT-OSS / Phi-MoE / Exaone MoE: every layer is MoE; per-layer bytes include all num_local_experts expert weight tensors.
• DeepSeek V3 / Nemotron-H: first_k_dense_replace determines which early layers are dense; moe_layer_freq controls which later layers are MoE. The profile marks each layer accordingly.
• Llama 4 (text tower and VLM): interleave_moe_layer_step selects every step-th layer as MoE with intermediate_size experts; the dense layers use intermediate_size_mlp.
• Jamba: expert_layer_period / expert_layer_offset mark MoE layers; the Mamba blocks themselves stay stage-local.
• Gemma 4 / Gemma 4 VLM: num_kv_shared_layers + layer_types emit KV-shared adjacency groups (see below); use_double_wide_mlp doubles the MLP footprint on each consumer layer.
• Quantised models: the profile honours the quantization.bits field (4-bit, 8-bit, etc.) with a small overhead for scales/biases.

KV-shared layer adjacency (Gemma 4)¶

Gemma 4's last num_kv_shared_layers decoder blocks do not compute their own keys and values. Instead each of them borrows the cache from an earlier layer of the same layer_types[i] value. Splitting between a source layer and its consumer would strand the cache on the wrong stage and break generation.

The profile builder records every (source, consumer) pair as a LayerAdjacencyGroup; the balancer refuses to drop a stage boundary in the interior of any such group. As a result, mlxcel generate --pp-size 2 and mlxcel-server --pp-auto 2 produce a valid plan by default for gemma-4-e2b-it-4bit.

Manual Partitioning (still supported)¶

Layer ranges can be specified explicitly with --pp-layers:

# 3-stage pipeline; stage 0 gets layers 0-19, stage 1 gets 20-39, stage 2 gets 40-59
mlxcel-server --model ./models/llama-60b \
  --pp-layers "0-19,20-39,40-59" \
  ...

The same syntax is available in the CLI generate path:

mlxcel generate \
  --model ./models/llama-3.2-1b-4bit \
  --prompt "Hello" \
  --max-tokens 64 \
  --temp 0 \
  --no-chat-template \
  --pp-layers "0-7,8-15"

When is manual override still useful?

• Reproducible experiments: you need byte-for-byte identical assignments between runs regardless of future profile-builder changes.
• Heterogeneous fleets: you want a specific stage to hold a specific range, independently of the memory estimate.
• Debugging: you want to verify that a specific boundary reproduces or fixes a bug.

Manual plans still pass through adjacency validation; attempting to cut a Gemma 4 KV-shared group surfaces an explicit error pointing at the source / consumer pair that would break.

Partition-quality report¶

The partitioner emits a PartitionQualityReport that records the estimated per-stage byte load, the overall imbalance ratio (max / min), and any warnings issued while honouring the adjacency constraints. In server mode, enabling verbose logging with --verbose makes mlxcel-server log one multi-line block per startup:

pipeline partition quality: 2 stages, max=12500000 bytes, min=11900000 bytes, imbalance=1.5x
  stage 0 device 'local-stage-0' layers=13 estimated=12500000 actual=unknown
  stage 1 device 'local-stage-1' layers=13 estimated=11900000 actual=unknown

The actual column is populated once each stage's weights have been loaded, so operators can see when the byte-cost estimator drifts from reality (the drift=... suffix). The balancer also emits an explicit warning at partition time when an adjacency constraint has forced a significantly imbalanced plan (max/min above ~1.5x):

pipeline auto-partition produced an imbalanced plan: stage byte loads range from
80 to 120 (1.50x), which may hurt throughput. Heaviest stage 1 carries
adjacency group 0..7 (synthetic forced-join).

That warning is a hint to rebalance hardware, not an error. Generation still proceeds.

Stage Assignment¶

Each StageAssignment records:

• stage_index: 0-based stage number.
• device_id: matches the node ID in the cluster configuration.
• layer_range: half-open range [start, end) of layer indices.
• has_embedding: true for stage 0.
• has_lm_head: true for the last stage.
• estimated_memory_bytes: estimated total memory for this stage.

Activation Transfer Between Stages¶

When a stage finishes processing a micro-batch, it sends the output activations to the next stage via the transport layer.

Activation tensors are serialized using the TensorProtocol wire format (src/distributed/tensor_protocol.rs) and sent over TCP or Thunderbolt connections. The transfer happens asynchronously with the next micro-batch forward pass on the sending stage.

sequenceDiagram
    participant S0 as Stage 0
    participant NET as Network
    participant S1 as Stage 1

    S0->>S0: Forward MB 0 (layers 0-15)
    S0->>NET: Send activations MB 0
    S0->>S0: Forward MB 1 (layers 0-15)
    NET->>S1: Receive activations MB 0
    S1->>S1: Forward MB 0 (layers 16-31)
    S0->>NET: Send activations MB 1
    S1->>S1: Forward MB 1 (layers 16-31)

The activation transfer module is in src/distributed/pipeline/activation_transfer.rs.

Micro-Batching and GPipe Schedule¶

A full request batch is split into micro-batches before entering the pipeline. This is the key mechanism that keeps all stages busy simultaneously.

gantt
    title GPipe Pipeline Schedule (3 stages, 4 micro-batches)
    dateFormat  s
    axisFormat  %S

    section Stage 0
    MB 0 :a0, 0, 1s
    MB 1 :a1, after a0, 1s
    MB 2 :a2, after a1, 1s
    MB 3 :a3, after a2, 1s

    section Stage 1
    bubble :crit, b_bubble0, after a0, 1s
    MB 0 :b0, after b_bubble0, 1s
    MB 1 :b1, after b0, 1s
    MB 2 :b2, after b1, 1s
    MB 3 :b3, after b2, 1s

    section Stage 2
    bubble :crit, c_bubble0, after a0, 2s
    MB 0 :c0, after c_bubble0, 1s
    MB 1 :c1, after c0, 1s
    MB 2 :c2, after c1, 1s
    MB 3 :c3, after c2, 1s

The GPipeSchedule emits ScheduleAction values that drive the execution loop:

Micro-batch splitting is in src/distributed/pipeline/micro_batch.rs.

Suggested Micro-Batch Size¶

suggested_micro_batch_size(batch_size, num_stages) returns the recommended number of sequences per micro-batch to maximize pipeline utilization:

• With num_stages stages and a total batch of batch_size sequences, the optimal number of micro-batches is num_stages (to fill all bubble slots).
• Each micro-batch has ceil(batch_size / num_stages) sequences.

Pipeline Bubbles¶

A pipeline bubble is a time slot where a stage has no work to do because it is waiting for activations from the previous stage.

In a k-stage pipeline with m micro-batches, the bubble fraction is:

To minimize bubbles: - Use more micro-batches (higher m). With m >> k, the bubble fraction approaches zero. - Use smaller micro-batches to allow faster filling of the pipeline.

For example, a 4-stage pipeline with 16 micro-batches has a bubble fraction of 3/18 ≈ 16.7%. With 32 micro-batches, the fraction drops to 3/34 ≈ 8.8%.

graph LR
    subgraph "Bubble fill-in with more micro-batches"
        B[More micro-batches] --> LB[Lower bubble fraction]
        B --> HT[Higher throughput]
        LB --> HT
    end

Per-Stage KV Cache Management¶

Each pipeline stage maintains its own KV cache for the layers it owns.

The cache manager (src/distributed/pipeline/cache_manager.rs) is responsible for:

• Allocating and releasing KV cache entries per sequence per stage.
• Tracking cache capacity per stage to prevent OOM.
• Coordinating cache eviction when a sequence is preempted.

Because each stage only holds caches for its own layer range, the total KV cache memory is distributed across stages rather than concentrated on a single node. This is one of the memory advantages of pipeline parallelism for long-context workloads.

Transport Backends¶

The cross-stage transport is selected through the transport_backend field on [cluster] in the cluster TOML or through the matching CLI surface. Three backends are shipped; all three implement the same Transport trait, so the activation-transfer consumer in src/distributed/pipeline/activation_transfer.rs is unchanged regardless of the backend in use.

Selection matrix¶

Example:

[cluster]
name = "pp-cluster"
pipeline_parallel_size = 2
transport_backend = "rdma"

[[nodes]]
id = "stage-0"
address = "192.168.1.11:9000"
role = "pipeline_stage"
stage = 0

[[nodes]]
id = "stage-1"
address = "192.168.1.12:9000"
role = "pipeline_stage"
stage = 1

RDMA-aware backend¶

The RDMA-aware backend is implemented in src/distributed/rdma_transport.rs and sits on top of the existing TCP core. On bind() it probes the OS for a zero-copy primitive:

• Linux: io_uring registered buffers. The probe first inspects /proc/sys/kernel/io_uring_disabled (for the kernel knob shipped in 6.6+) and then attempts a lightweight io_uring_setup syscall to confirm the syscall is reachable and not blocked by seccomp / LSM.
• macOS: kqueue + writev batched send with registered memory regions. When a Thunderbolt Bridge interface with a routable IPv4 address is available, the wrapper attaches on top of the existing Thunderbolt transport so registered-buffer sends ride the validated interface. Otherwise it attaches on top of plain TCP.
• Other platforms: no accelerated path is attempted. The transport falls back to TCP with the reason os=<family>, driver=unsupported.

Peers exchange a protocol-version handshake at connect() time (RDMA_PROTOCOL_VERSION). If the peer advertises 0 (not RDMA-aware) or a mismatched version, the local side records a peer-level fallback and keeps running on the TCP semantics.

Silent-but-logged fallback¶

Every fallback is reported once, on a single line, with a structured reason. Operators should grep their log stream for the prefix rdma_transport: falling back to TCP. Every such line carries a reason= key naming the fault in OS / driver / peer-version / capability-mismatch form:

rdma_transport: falling back to TCP (reason=os=linux, driver=io_uring_missing: kernel does not export io_uring_setup (ENOSYS))
rdma_transport: peer negotiation fallback (peer=10.0.0.1:9000, reason=os=linux, peer_version=99: incompatible with local RDMA protocol version 1)

The coordinator keeps running after any fallback. mlxcel generate / mlxcel-server never crash because of a missing zero-copy primitive. The structured metrics exposed via RdmaTransport::acceleration(), fallback_reason(), peers_negotiated(), and peer_fallbacks() make the negotiation state easy to plumb into the distributed metrics collector.

Benchmarking¶

A 2-node loopback activation-transfer benchmark lives in src/distributed/bench_activation.rs and is exercised by tests/rdma_transport_bench.rs. Run:

cargo test --release --test rdma_transport_bench \
    rdma_vs_tcp_activation_transfer -- --ignored --nocapture

The harness prints a markdown table with p50/p95/p99 one-way latency and bandwidth for TCP vs RDMA across four payload sizes that mirror real activation shapes (RPC-sized handoff, decode-step activation, prefill-sized activation, and worst-case stage hand-off). Use this generated table as the local baseline for transport changes.

Configuration¶

The PipelineConfig requires at least 2 stages and at least 1 sequence per micro-batch. An error is returned at startup if these constraints are violated.

Example: 4-stage pipeline with the zero-config path, which uses four --pp-peer machines and a coordinator that runs no stage:

# Each of the four peers (one command per machine, own address in --cluster-control-addr).
mlxcel-server --model ./models/llama-70b \
  --pp-peer \
  --cluster-secret-file ./cluster.secret \
  --cluster-control-addr 192.168.1.11:19001

# Coordinator: plans the cluster and pushes one stage assignment to each peer.
mlxcel-server --model ./models/llama-70b \
  --pp-auto 4 \
  --cluster-secret-file ./cluster.secret \
  --cluster-peers 192.168.1.11:19001,192.168.1.12:19001,192.168.1.13:19001,192.168.1.14:19001 \
  --cluster-config-out ./.mlxcel/cluster.toml

The manual multi-command workflow uses the generated cluster TOML and the standard distributed flags:

# Stage node using a shared cluster TOML
mlxcel-server --model ./models/llama-70b \
  --distributed-config /shared/.mlxcel/cluster.toml \
  --node-id stage-0 \
  --node-role pipeline_stage \
  --peers 192.168.1.11:8080,192.168.1.12:8080,192.168.1.13:8080

# Repeat with each stage's node-id on the other nodes

Composing with Tensor Parallelism (2D PP × TP)¶

Pipeline Parallelism composes with Tensor Parallelism to form a 2D mesh. Each node owns a specific (pp_stage, tp_rank) intersection: PP splits the model vertically across stages, and within each stage TP shards the layer weights across ranks.

See also: Tensor Parallelism → Composing with Pipeline Parallelism (2D PP × TP) for the TP-side perspective on the same mesh.

When to use 2D parallelism¶

Use PP × TP when a single stage of a PP-only topology still exceeds one node's memory, for example when serving a Llama-70B-class model on four medium-sized GPUs as a 2×2 (2 PP stages × 2 TP ranks) mesh. PP splits the model into two halves; TP further shards each half across two devices so each device owns a quarter of the weights.

Current limitations:

• Uniform TP width per stage. Every stage must use the same tensor_parallel_size. Heterogeneous widths (e.g., stage 0 at tp=4, stage 1 at tp=2) are a future extension.
• No 3D parallelism. Expert-parallelism (PP × TP × EP) is tracked separately.

Role encoding¶

Nodes participating in a 2D mesh use the pipeline_tensor_parallel role (aliases: pp_tp, 2d). Each such node sets both stage and rank on its TOML entry. The following invariants are enforced at validation:

• exactly pipeline_parallel_size × tensor_parallel_size nodes carry the PPTP role
• every (stage, rank) pair in the grid is present exactly once
• non-PPTP nodes cannot claim stage or rank

The role encoding is exposed in Rust as the NodeRole::PipelineTensorParallel variant and the PpTpCoord coordinate on the cache side.

Transport traffic classes¶

The 2D routing layer shares a single connection pool across every node in the mesh, but it distinguishes two traffic classes so that the right handler processes each message:

Both classes are selected via PipelineTensorParallelRouter, which routes on (preferred_stage, preferred_rank, traffic_class) and returns the single (stage, rank) peer that should receive the message. Broadcasting to every rank on a stage (for all-reduce) or every stage at the same rank (for pipelined PP activations) uses NodeRegistry::nodes_at_stage or nodes_at_rank directly.

Cache coordination¶

Each (pp_stage, tp_rank) slot owns its own per-layer, per-rank KV shard. Admission of a sequence must be coherent across the entire grid: if a single rank cannot fit the sequence, the stage would run a sharded forward pass with only some ranks caching the result and silently produce wrong output.

The coordinator uses coordinated_2d_admission which performs a two-step check: a dry-run on every slot followed by a commit only if every slot accepts. On rejection nothing is committed anywhere.

use mlxcel::distributed::pipeline::{
    coordinated_2d_admission,
    CacheAdmissionRequest,
    PpTpAdmissionOutcome,
    PpTpCoord,
};

let mut grid = vec![
    (PpTpCoord::new(0, 0), &mut mgr_s0_r0),
    (PpTpCoord::new(0, 1), &mut mgr_s0_r1),
    (PpTpCoord::new(1, 0), &mut mgr_s1_r0),
    (PpTpCoord::new(1, 1), &mut mgr_s1_r1),
];
let req = CacheAdmissionRequest::new(42, prompt_len);
match coordinated_2d_admission(&mut grid, &req)? {
    PpTpAdmissionOutcome::Admitted => {
        // Stage coherent; proceed.
    }
    PpTpAdmissionOutcome::Rejected { at, reason } => {
        tracing::warn!(%at, ?reason, "2D admission rejected");
    }
}

Cluster TOML (2×2 example)¶

[cluster]
name = "llama-70b-2x2"
pipeline_parallel_size = 2
tensor_parallel_size = 2
transport_backend = "rdma"

# Stage 0, TP rank 0
[[nodes]]
id = "s0-r0"
address = "192.168.1.10:8080"
role = "pipeline_tensor_parallel"
stage = 0
rank = 0

# Stage 0, TP rank 1
[[nodes]]
id = "s0-r1"
address = "192.168.1.11:8080"
role = "pipeline_tensor_parallel"
stage = 0
rank = 1

# Stage 1, TP rank 0
[[nodes]]
id = "s1-r0"
address = "192.168.1.20:8080"
role = "pipeline_tensor_parallel"
stage = 1
rank = 0

# Stage 1, TP rank 1
[[nodes]]
id = "s1-r1"
address = "192.168.1.21:8080"
role = "pipeline_tensor_parallel"
stage = 1
rank = 1

CLI entry point¶

mlxcel generate \
  --model ./models/llama-70b-4bit \
  --pp-size 2 \
  --tp-size 2 \
  -p "Hello" -n 50

The validator in src/commands/generate.rs accepts the combination and emits a readable error on inconsistent topologies (e.g., when the emitted cluster TOML has fewer nodes than pp_size × tp_size).

Parity report¶

A small Llama-class proxy is used to verify that --pp-size 2 --tp-size 2 produces the same greedy-decoded tokens as a single-device reference on a fixed seed. Use this check as the reference validation for 2D PP × TP token parity.

LoRA with PP¶

Pipeline Parallelism composes with LoRA adapters: the --adapter flag that drives single-process adapter inference also drives adapter inference across pipeline stages, using the same adapter directory layout and the same rank / scaling semantics. Operators do not need a second tool chain or a second adapter format to run an adapter across multiple stages. The current implementation supports single-adapter composition on the Llama family; multi-adapter stacking and runtime hot-swap are not supported.

Stage-local adapter loading¶

Each pipeline stage loads the base model weights for its assigned layer range and then composes the LoRA delta into those base weights before the stage executor is materialized. The per-stage loader walks the adapter's adapters.safetensors (or the fallback filename adapter_model.safetensors) file via the mlxcel-core filtered loader and skips any adapter tensor whose layer index falls outside the stage's assigned range. Out-of-range adapter tensors are never taken into the stage's Rust-side weight map, so they are neither buffered nor discarded. The saving is in both heap traffic and in MLX-side array construction.

The fusion itself is shared with the non-PP path: both call into the same fuse_lora_weights_into primitive, so there is exactly one implementation of the LoRA W_fused = W_base + scale * (lora_b @ lora_a) math and exactly one place that negotiates the mlx-lm vs. HuggingFace PEFT shape conventions.

graph LR
    A[adapter_config.json] --> C[AdapterConfig]
    B[adapters.safetensors] --> D[load_stage_adapter_weights]
    C --> E[apply_stage_lora_adapter]
    D --> E
    S[Stage layer range] --> D
    F[Base weight map for stage] --> E
    E --> G[Fused stage weight map]
    G --> H[Stage executor]

CLI usage¶

# Two-stage PP + LoRA on a Llama checkpoint.
mlxcel generate \
  --model ./models/llama-3.2-1b-4bit \
  --adapter ./models/llama-3.2-1b-4bit-lora \
  --pp-size 2 \
  -p "Hello" \
  -n 64 \
  --temp 0

The same flag and adapter directory layout also work with mlxcel-server when a pipeline-parallel in-process runtime is configured.

Parity contract¶

For a fixed prompt, fixed seed, and --temp 0, a PP + LoRA run must match the decoded token stream of the single-process --adapter run bit-for-bit. This contract is exercised by the ignored integration test tests/pipeline_lora_composition_real_models.rs, which runs a published LoRA on top of a Llama checkpoint across two stages and asserts strict equality against the single-process adapter output.

Unit coverage in src/distributed/pipeline/partial_loading.rs independently verifies:

• the stage-local adapter tensor filter (layer-range inclusion, embedding and lm_head ownership, language_model.model.layers. VLM-style prefixes);
• that load_stage_adapter_weights skips, rather than buffers and discards, adapter tensors that belong to other stages;
• that the stage-filtered base + stage-filtered adapter composition produces the same fused tensors as the full-fusion-then-filter path on the stage's layer range.

Scope and guards¶

The first stage executor to add adapter support is LlamaStageExecutor. Other families route through a guard (ensure_no_adapter) that bails loudly if an adapter is passed for a model family whose PP stage executor does not implement the composition step. This prevents silently dropping the adapter on the floor. Adding a new family to the supported matrix is a three-line change: accept an adapter_path: Option<&Path> on the family load signature, call apply_stage_lora_adapter between the weight load and the filter_weight_map call, and remove the family from the guard.

Currently out of scope:

• Multi-adapter stacking across a PP cluster. Only one adapter may be loaded at stage initialization.
• Runtime hot-swap of adapters on an active PP cluster. Adapters are resolved and fused once during stage initialization.
• Tensor Parallelism + LoRA + PP. The existing TP validator still rejects the adapter path; adapter support for 2D TP x PP composition is not supported by the current implementation.
• Remote-coordinator PP with adapters. The remote pipeline topology explicitly bails on adapter paths today; adapter composition covers the in-process server runtime and the CLI mlxcel generate path.

Pipeline Validation Coverage¶

Multi-host pipeline-parallel behavior is covered by two validation tiers. This section summarizes the guarantees operators can rely on; exact workflow files, script names, and test module paths are intentionally kept out of the user manual because they are maintainer-facing details.

The logical tier is designed to run without model weights. It validates that the automatic partitioner assigns fewer layers to memory-constrained stages and that long prompts are rejected coherently when any stage lacks KV-cache capacity. This guards against over-admission that would otherwise surface later as an out-of-memory failure or a stage-incoherent cache state.

The real-model tier is intentionally opt-in because it consumes dedicated multi-host hardware and local model weights. Use it before changing transport behavior, activation handoff, stage scheduling, or anything that can affect token parity across a remote pipeline cluster.

Elastic repartitioning (experimental)¶

The partition plan is fixed at cluster startup by default. An opt-in elastic flow lets a running cluster redistribute layers across stages without a full restart.

Enabling the flow¶

Pass --enable-elastic-pp to mlxcel-server on the coordinator process. The flag is strictly additive. Clusters that do not pass it behave exactly as before.

mlxcel-server -m models/llama-3.2-1b-4bit \
  --distributed-config cluster.toml --node-id coordinator \
  --parallel 2 --pp-micro-batch-size 2 \
  --enable-elastic-pp \
  --elastic-pp-drain-timeout 120 \
  --elastic-pp-pressure-fraction 0.92 \
  --elastic-pp-cool-down 30

Triggers¶

Two trigger kinds are currently supported:

Predictive policies based on forecast load are not supported.

Drain semantics¶

When a trigger is accepted, the coordinator walks the state machine

Idle → Draining → Rebalancing → Resuming → Idle

driving each step:

1. Draining: the coordinator fans out BeginDrain to every stage, the server rejects new sequences with a 503-equivalent panic while the drain flag is set, and the coordinator polls probe_stages until every stage reports in_flight_requests == 0 or the --elastic-pp-drain-timeout elapses.
2. Rebalancing: the coordinator calls the byte-accurate balancer (partition_balance::balance_layers) with the current per-stage memory budgets, reuses existing adjacency constraints (Gemma 4, Jamba), and validates the new plan with partition::validate_partition.
3. Resuming: each stage reloads layer weights via partial_loading with the new LayerFilter. The KV cache is rebuilt from scratch. Any sequence that was drained is reissued by the client (HTTP retry); no user work is lost because drain completed before rebalancing.

Failures at any step drive the coordinator to the Failed terminal state. The previous layer ranges stay active; drain is released on a best-effort basis so the cluster can continue serving. Retrying requires another explicit trigger.

Observability¶

Every state transition emits a structured RepartitionEvent that feeds the metrics subsystem. See the Observability section below for the exact metric names and trace events exposed via /metrics and --debug-pp-trace.

Limitations¶

• Cross-node layer migration during in-flight requests is not supported. The current flow drains first, then repartitions. True zero-downtime migration is not supported by the current implementation.
• Peer-to-peer weight streaming is out of scope. Stages reload from the original safetensors on disk.
• Only the pipeline axis is elastically reshaped. 2D (pp_stage, tp_rank) repartition across both axes is not supported.
• There is no HTTP admin endpoint. explicit triggers are wired via the in-process coordinator API for now.

Observability¶

The Prometheus /metrics endpoint exposes four metric families for pipeline parallelism, along with a chrome-tracing JSON writer for offline analysis and attribution-aware OOM diagnostics.

/metrics endpoint¶

Enable with mlxcel-server --metrics or mlxcel-server --metrics-port <PORT>. The endpoint is currently multiplexed onto the main HTTP listener; passing a separate metrics port emits a warning. The endpoint exposes the existing server-level request counters plus the pipeline-parallel families below:

reason is one of memory, sequence_cap, already_cached, rejected_by_stage. trigger is one of explicit, memory_pressure. outcome is one of completed, aborted, failed, or progress (intermediate state transitions).

Chrome-tracing mode (--debug-pp-trace)¶

mlxcel-server --debug-pp-trace <path> enables the PpTracer. Every scheduler action is buffered in memory and the final JSON file is written on graceful shutdown (or explicit flush). Open the file in chrome://tracing or https://ui.perfetto.dev for timeline inspection.

Recorded events:

• batch_arrival: scheduler sees a new request batch.
• stage (duration): wraps a stage's forward pass; the tid column groups by stage_index.
• activation_send / activation_recv: one pair per hand-off, with bytes and direction arguments.
• admission_reject: includes the offending reason, required_bytes, and available_bytes.

Trace files are idempotent: the first successful flush wins, and Drop will not double-write.

OOM diagnostics¶

KV cache admission rejections surface the stage identity and KV occupancy through coordinated_admission_with_attribution. The returned AdmissionDiagnostic carries:

• stage_index: which stage rejected.
• reason: enum variant (with metric_label() producing the low-cardinality short name used by Prometheus).
• used_memory_bytes and budget_bytes: the offending stage's KV accounting at the moment of rejection.
• active_sequences: current occupancy.

The server error path formats this into a human-readable diagnostic instead of a generic 500.

Grafana dashboard¶

A typical Grafana dashboard for these metrics includes:

1. Per-stage busy fraction (timeseries).
2. Mean bubble ratio (stat panel with red/orange/green thresholds).
3. Activation transfer latency p95 and p99 per stage pair.
4. Rejection rate by (stage, reason).
5. Cumulative repartition events by (trigger, outcome).

Import it against a Prometheus data source that scrapes mlxcel-server --metrics. Operators supply their own Grafana stack. The dashboard JSON is the only deliverable shipped.