WWDC.ai

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.

Watch on Apple Developer

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_config to generate a hostfile and configure a Thunderbolt RDMA cluster, then wrap existing MLX LM commands with mlx.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_config can 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=1 for faster GPU-to-CPU synchronization in distributed jobs.
  • Use --backend jaccl for mesh and --backend jaccl-ring for 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 jaccl

Hostfile 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 2048

Select 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 \
  --pipeline

Parallelism 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 16

Python, 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_load from MLX LM to load a distributed model for inference.
  • Use MLX distributed collectives such as all_sum from 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

Unofficial, not associated with Apple. BySuperwall

On this page

Ask AI