Pick pipeline or tensor, size memory

mkdir -p ~/mlx-cluster && cd ~/mlx-cluster
uv init --bare --python 3.11
uv venv
uv add mlx mlx-lm huggingface_hub

# PyTorch is optional in this course. On Apple silicon it uses MPS, on other machines the CPU.
uv add torch
uv run python -c "import torch; print('mps' if torch.backends.mps.is_available() else 'cpu')"

uv run python -c "import mlx.core as mx; print(mx.__version__, mx.distributed.is_available())"
sw_vers -productVersion        # 26.2 or newer for RDMA over Thunderbolt

Why you would do this

On 15 October 2026 Apple published Local and Distributed Large LLM Inference with MLX and RDMA on Mac Studio. The model in the post is Muse Spark Open, a mixture of experts model with 816B total parameters and 42B active. Apple says it needs more than 1.6 TB in BF16 and about 493 GB in 4-bit NVFP4, for the weights alone. Add memory for the work during inference and it does not fit on one Mac Studio with 512 GB. Apple split it across Mac Studios with M5 Ultra and measured 36.73 tokens per second on four of them at an 8,192-token context.

The business reason: a team can run a frontier-class open model on its own desk, for private code or private documents, with no cloud bill per token. Apple measured about 300 W peak for one M5 Ultra Mac Studio. The cost is setup work and new kinds of failure. This course is about both.

Two ways to split a model

• Pipeline parallel splits the model by depth. Mac A runs the first group of layers, then sends the result to Mac B, which runs the next group. Only the output between groups crosses the cable.
• Tensor parallel splits each layer by width. Every Mac holds a part of every layer. For every token and every layer the Macs exchange inputs and outputs, so there are many small messages.

Apple's numbers at an 8,192-token context (from Tables 2 and 3 of the post):

What the post says about when each one wins:

• Pipeline keeps decode speed about the same as one Mac, because layers still run one after another. Its job is to make a model fit. It is simple and sends little data.
• Pipeline can make the first token faster on very long prompts, up to about 1.8 times at 65,536 tokens, because MLX cuts the prompt into chunks and the Macs work on different chunks at the same time. At 8,192 tokens there was no gain.
• Tensor makes decode and the first token faster as you add Macs. It needs very low latency between Macs, which is why Apple used RDMA over Thunderbolt (the JACCL backend). Dense models scaled better than the MoE model.
• With only two Macs and a very long prompt, tensor can lose to pipeline on prefill. Apple saw this with Muse Spark Open at 65,536 tokens. With four Macs, tensor was ahead again.

Check which split your model supports

Not every model supports both. In the mlx-lm source, sharded_load() checks hasattr(model, "shard") for tensor and hasattr(model.model, "pipeline") for pipeline. If you ask for a split the model does not have, loading stops with an error. We hit this on our test box:

ValueError: The model does not support pipelining but a pipeline_group was provided

This script runs the same check without downloading the weights:

"""Check if mlx-lm can split a model with --pipeline, tensor parallel, or both.

Downloads only config and code files, not the weights.
Uses the same checks as mlx_lm.utils.sharded_load.
"""
import sys
from pathlib import Path

from huggingface_hub import snapshot_download
from mlx_lm.utils import load_model

repo = sys.argv[1]
path = Path(snapshot_download(repo, allow_patterns=["*.json", "*.py"]))
model, config = load_model(path, lazy=True, strict=False)

pipeline = hasattr(model, "model") and hasattr(model.model, "pipeline")
tensor = hasattr(model, "shard")
print(f"{repo} ({config.get('model_type')})")
print(f"  tensor parallel (default): {'yes' if tensor else 'no'}")
print(f"  pipeline (--pipeline):     {'yes' if pipeline else 'no'}")
if not (pipeline or tensor):
    print("  this model cannot be split by mlx-lm")

Output on our Linux test box (mlx 0.32.4, mlx-lm 0.32.2):

$ uv run python check_split_support.py mlx-community/Qwen3-0.6B-4bit
mlx-community/Qwen3-0.6B-4bit (qwen3)
  tensor parallel (default): yes
  pipeline (--pipeline):     no
$ uv run python check_split_support.py mlx-community/Qwen3-30B-A3B-4bit
mlx-community/Qwen3-30B-A3B-4bit (qwen3_moe)
  tensor parallel (default): yes
  pipeline (--pipeline):     yes

In mlx-lm, tensor parallel is the default when the model supports it. Add --pipeline to mlx_lm.benchmark, mlx_lm.chat, or mlx_lm.server to use pipeline instead.

Size the memory on each Mac

Total memory is not enough. Each Mac must hold its own part, plus room for the KV cache and buffers. Two details from the mlx-lm source matter here:

• Pipeline splits layers by count, not by memory. The low ranks get any extra layers, and rank 0 (the first host in your hostfile) runs the last layers. Every rank also keeps the embeddings, the final norm, and the output head. The command-line tools do not let you choose a custom split.
• Tensor splits each layer evenly across ranks. A Mac with less memory still gets the same share.

The second point bites people with mixed Macs. In exo #2430, a 69.6 GB model was placed with tensor sharding on a 36 GB MacBook Pro and a 64 GB Mac mini because it fit in their combined memory. Each Mac got about 34.8 GB. The 36 GB Mac could not hold that in GPU memory, its runner aborted with signal 6, and Exo restarted it in a loop. The other Mac then showed misleading JACCL errors.

This script reads tensor sizes from the safetensors headers on Hugging Face (no weight download) and compares each Mac's share with its budget:

"""Estimate weight memory per Mac for pipeline and tensor splits.

Reads tensor sizes from the safetensors headers on Hugging Face.
It does not download the weights.
"""
import argparse
import re
from collections import defaultdict

from huggingface_hub import get_safetensors_metadata

LAYER = re.compile(r"(?:^|\.)layers\.(\d+)\.")
GiB = 2**30


def tensor_bytes(info):
    start, end = info.data_offsets
    return end - start


def layer_sizes(repo):
    meta = get_safetensors_metadata(repo)
    per_layer = defaultdict(int)
    shared = 0
    for file_meta in meta.files_metadata.values():
        for name, info in file_meta.tensors.items():
            m = LAYER.search(name)
            if m:
                per_layer[int(m.group(1))] += tensor_bytes(info)
            else:
                shared += tensor_bytes(info)
    layers = [per_layer[i] for i in range(max(per_layer) + 1)]
    return layers, shared


def pipeline_split(n_layers, n_nodes):
    # Same rule as mlx_lm/models/pipeline.py: even split, low ranks get extra
    # layers, and rank 0 runs the LAST layers.
    base, extra = divmod(n_layers, n_nodes)
    split = [base + (1 if r < extra else 0) for r in range(n_nodes)]
    ranges = []
    for r in range(n_nodes):
        start = sum(split[r + 1:])
        ranges.append((start, start + split[r]))
    return ranges


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True)
    p.add_argument("--nodes", type=int, required=True)
    p.add_argument("--working-set-gib", type=float, nargs="+", required=True,
                   help="max_recommended_working_set_size per Mac, in hostfile order")
    p.add_argument("--headroom", type=float, default=0.25,
                   help="share of each working set kept free for KV cache and buffers")
    a = p.parse_args()
    assert len(a.working_set_gib) == a.nodes, "one working set per node"

    layers, shared = layer_sizes(a.model)
    total = sum(layers) + shared
    print(f"{a.model}: {len(layers)} layers, weights {total / GiB:.1f} GiB "
          f"(shared {shared / GiB:.2f} GiB)")

    print("\npipeline (--pipeline): every rank also keeps embeddings, final norm, head")
    for r, (s, e) in enumerate(pipeline_split(len(layers), a.nodes)):
        need = sum(layers[s:e]) + shared
        budget = a.working_set_gib[r] * (1 - a.headroom) * GiB
        flag = "OK" if need <= budget else "TOO BIG"
        print(f"  rank {r}: layers {s}-{e - 1}, {need / GiB:6.1f} GiB "
              f"of {budget / GiB:6.1f} GiB budget  {flag}")

    print("\ntensor (default): each layer is split evenly across ranks")
    for r in range(a.nodes):
        need = sum(layers) / a.nodes + shared  # upper bound: shared counted in full
        budget = a.working_set_gib[r] * (1 - a.headroom) * GiB
        flag = "OK" if need <= budget else "TOO BIG"
        print(f"  rank {r}: about {need / GiB:6.1f} GiB "
              f"of {budget / GiB:6.1f} GiB budget  {flag}")


if __name__ == "__main__":
    main()

Get each Mac's working set with mx.device_info()["max_recommended_working_set_size"], as shown in Keep a Mac MLX box stable. Output on our test box for a small MoE model and two example budgets:

$ uv run python size_shards.py --model mlx-community/Qwen3-30B-A3B-4bit --nodes 2 --working-set-gib 26.8 48
mlx-community/Qwen3-30B-A3B-4bit: 48 layers, weights 16.0 GiB (shared 0.33 GiB)

pipeline (--pipeline): every rank also keeps embeddings, final norm, head
  rank 0: layers 24-47,    8.2 GiB of   20.1 GiB budget  OK
  rank 1: layers 0-23,     8.2 GiB of   36.0 GiB budget  OK

tensor (default): each layer is split evenly across ranks
  rank 0: about    8.2 GiB of   20.1 GiB budget  OK
  rank 1: about    8.2 GiB of   36.0 GiB budget  OK

The 25% headroom is our choice, not an Apple number. Long contexts need more. Apple reports that Mistral Medium 3.5 (128B, FP8) used up to 158 GB on one Mac at a 65,536-token context, so plan for tens of gigabytes beyond the weights at long contexts. Measure on your own Macs.

Sources

• Apple Machine Learning Research: Distributed LLM inference with MLX and RDMA
• MLX docs: Distributed Communication
• mlx-lm repository (sharded_load, pipeline)