Explore distributed inference and training with MLX
Scale MLX inference and fine-tuning across multiple Macs with RDMA over Thunderbolt 5, JACCL, mlx.launch, and MLX LM.
TL;DR
- MLX can distribute inference and fine-tuning across multiple Apple silicon Macs using RDMA over Thunderbolt 5 and Apple's open-source JACCL communication library.
- Use
mlx.distributed_configto generate a hostfile and configure a Thunderbolt RDMA cluster, then wrap existing MLX LM commands withmlx.launch. - MLX LM supports tensor parallelism for faster inference, pipeline parallelism for some large models, and data-parallel LoRA fine-tuning with scaled batch size.
- Distributed APIs are available at several levels: MLX LM Python helpers, lower-level MLX distributed primitives in Python/Swift/C++, and standalone JACCL C++ collectives.
Distributed stack on Apple silicon
Distributed MLX workloads use a layered stack: Thunderbolt 5 provides the physical interconnect, RDMA over Thunderbolt moves data directly between machine memories with low CPU and OS overhead, JACCL provides collective communication primitives, and MLX uses those primitives for distributed inference and training.
JACCL is an open-source collective communication library from Apple. It supports communication patterns needed for distributed ML, such as sending data between machines and combining results across a group, but can also be used for non-ML distributed workloads.
- RDMA over Thunderbolt 5 is supported starting in macOS 26.2.
- JACCL uses RDMA over Thunderbolt for low-latency communication across Apple silicon Macs.
- MLX and MLX LM build on JACCL to shard models, launch distributed jobs, and coordinate communication.
Set up a multi-Mac MLX cluster
The session demonstrates a four-node M3 Ultra cluster connected with Thunderbolt 5. Topology matters because communication cost has both latency and transfer-time components: small messages are latency-bound, while large messages are bandwidth-bound.
JACCL supports mesh and ring topologies. A full mesh gives direct one-hop communication between every pair of machines and is preferred for latency-sensitive patterns such as tensor parallelism. A ring uses fewer ports and cables and can use multiple cables per neighbor to improve bandwidth. With a mesh, JACCL can choose mesh or ring routing depending on message size and operation.
After connecting the machines, enable RDMA over Thunderbolt in System Settings on every Mac and reboot. mlx.launch starts distributed programs over SSH using a JSON hostfile, after which the nodes communicate over the Thunderbolt links.
- The hostfile contains one entry per node with SSH hostname, local-network IPs for JACCL coordination, and RDMA device names for Thunderbolt peer connections.
mlx.distributed_configcan generate the hostfile, check SSH reachability, probe Thunderbolt topology, and optionally configure Thunderbolt links for RDMA.- Environment variables can be embedded in the hostfile; the session uses
MLX_METAL_FAST_SYNCH=1for faster GPU-to-CPU synchronization in distributed jobs. - Use
--backend jacclfor mesh and--backend jaccl-ringfor ring.
Generate a mesh hostfile for four Macs
Creates a cluster configuration for mlx.launch, embeds MLX_METAL_FAST_SYNCH=1, and uses the JACCL mesh backend.
mlx.distributed_config \
--hosts m3-ultra-0,m3-ultra-1,m3-ultra-2,m3-ultra-3 \
--output "m3-ultra-jaccl.json" \
--env MLX_METAL_FAST_SYNCH=1 \
--auto-setup \
--backend jacclHostfile shape
Each node declares the SSH hostname, coordination IP, and RDMA interfaces for its Thunderbolt peer links.
[
{
"ssh": "m3-ultra-0",
"ips": ["192.168.1.10"],
"rdma": [null, "rdma_en5", "rdma_en4", "rdma_en3"]
}
]Run distributed LLM inference with MLX LM
The simplest migration path is to keep the existing MLX LM command and wrap it with mlx.launch. The launcher connects to each host from the hostfile, starts the executable on each node, and MLX LM shards the model and coordinates distributed inference.
The session compares Qwen 3.6 27B on one M3 Ultra versus four M3 Ultras and reports nearly 3x token generation rate on the cluster. It also demonstrates running Kimi 2.6, a one-trillion-parameter model whose 8-bit weights alone require about one terabyte of memory, across four Macs.
- All required libraries, including MLX, must be installed on every Mac.
- The executable path passed after
--must be valid on each remote node. - Tensor parallelism is the default MLX LM sharding strategy.
- Pipeline parallelism can be requested with
--pipeline, but not all models support it.
Wrap a single-device chat command with mlx.launch
The distributed command is the same MLX LM chat invocation launched across the cluster.
# Single-device LLM inference
mlx_lm.chat --model "Qwen/Qwen3.6-27B" --max-tokens 2048
# Distributed LLM inference across the cluster
mlx.launch --hostfile "m3-ultra-jaccl.json" -- \
/remote/path/to/mlx_lm.chat --model "Qwen/Qwen3.6-27B" --max-tokens 2048Select tensor or pipeline parallelism
Tensor parallelism splits layers by width; pipeline parallelism splits the model by depth where supported.
# Tensor parallelism (default)
mlx.launch --hostfile "m3-ultra-jaccl.json" -- \
/remote/path/to/mlx_lm.chat --model "moonshotai/Kimi-K2.6" \
--max-tokens 2048
# Pipeline parallelism
mlx.launch --hostfile "m3-ultra-jaccl.json" -- \
/remote/path/to/mlx_lm.chat --model "moonshotai/Kimi-K2.6" \
--max-tokens 2048 \
--pipelineParallelism strategies and distributed fine-tuning
MLX LM uses model parallelism for large-model inference. Pipeline parallelism assigns groups of layers to different machines, reducing communication to boundary activations but not speeding up per-token execution because tokens still pass through layer groups sequentially. Tensor parallelism splits each layer by width so machines process the same token concurrently, improving inference speed at the cost of communication at every layer and token.
For fine-tuning, MLX LM uses data parallelism: each Mac holds a model replica, receives different batches, computes gradients locally, and averages gradients across machines. The session demonstrates LoRA fine-tuning Qwen 3.5 9B and scales --batch-size by the number of devices so each machine processes the same number of samples per step.
- Tensor parallelism benefits from low-latency mesh communication because every layer can require cross-node communication.
- Pipeline parallelism has simpler communication but does not inherently accelerate single-token inference.
- Data-parallel training can process data up to roughly N times faster with N machines, depending on workload and overhead.
- The demonstrated four-node fine-tuning run processes about 600 tokens/s versus about 180 tokens/s on one M3 Ultra.
Run distributed LoRA fine-tuning
The distributed LoRA command uses mlx.launch; MLX LM handles data sharding and gradient averaging.
# Single-device fine-tuning
mlx_lm.lora --model "Qwen/Qwen3.5-9B" \
--data "mlx-community/wikisql" \
--train --batch-size 4
# Distributed fine-tuning: scale batch size by number of devices
mlx.launch --hostfile "hostfile.json" -- \
/remote/path/to/mlx_lm.lora --model "Qwen/Qwen3.5-9B" \
--data "mlx-community/wikisql" \
--train --batch-size 16Python, Swift, C++, and standalone JACCL APIs
Beyond the CLI, MLX exposes distributed functionality through Python, Swift, and C++ APIs. At the MLX LM level, Python code can initialize a distributed group, choose tensor or pipeline groups, load a sharded model, and generate as if using a single-device model while MLX LM handles communication.
Lower-level MLX APIs expose layer sharding and collective operations such as all-reduce. JACCL can also be used directly from C++ without MLX for distributed applications that need collective communication but are not necessarily machine-learning workloads.
- Use
mx.distributed.init(strict=True, backend="jaccl")to initialize a Python distributed group with the JACCL backend. - Use
sharded_loadfrom MLX LM to load a distributed model for inference. - Use MLX distributed collectives such as
all_sumfrom Python, Swift, or C++. - Use standalone JACCL C++ APIs when an application needs collective communication without MLX.
Distributed inference with MLX LM Python API
Initializes distributed communication, loads a tensor-sharded model, and prints generated text from rank 0.
import mlx.core as mx
from mlx_lm import stream_generate
from mlx_lm.utils import sharded_load
group = mx.distributed.init(strict=True, backend="jaccl")
tensor_group, pipeline_group = group, None
model, tokenizer = sharded_load("moonshotai/Kimi-K2.6", pipeline_group, tensor_group)
for response in stream_generate(model, tokenizer, prompt, max_tokens=1024):
if group.rank() == 0:
print(response.text, end="", flush=True)Shard a layer and use a collective
Shows lower-level MLX control over tensor-parallel layer sharding and an all-reduce sum.
import mlx.core as mx
import mlx.nn as nn
group = mx.distributed.init(strict=True, backend="jaccl")
layer = nn.Linear(1024, 1024)
sharded_layer = nn.layers.distributed.shard_linear(
layer, strategy="all-to-sharded", group=group
)
data = mx.random.normal((1, 1, 1024))
output = sharded_layer(data)
mx.eval(output)
values = mx.full((4,), float(group.rank()), dtype=mx.float32)
result = mx.distributed.all_sum(values, group=group)
mx.eval(result)Resources
- MLX Swift LM on GitHub
- MLX Swift Examples
- MLX Examples
- MLX Swift
- MLX LM - Python API
- MLX Explore - Python API
- MLX Framework
- MLX
Related Sessions
- Explore numerical computing in Swift with MLX
- Run local agentic AI on the Mac using MLX
- Explore large language models on Apple silicon with MLX
- Get started with MLX for Apple silicon
Run local agentic AI on the Mac using MLX
Build fully local agentic AI workflows on Mac with MLX-LM Server, OpenAI-compatible agents, continuous batching, and distributed inference.
What’s new in image understanding
Use Vision tap-to-segment, Foundation Models image inputs and image tool calling, and Vision on watchOS for richer app image understanding.