Test the link before loading weights

Why test first

In the r/LocalLLM thread, a 4 × M3 Ultra cluster failed the same way for the full model launch and for a trivial 10-element all_sum: [jaccl] Couldn't connect (error: 60), before any weights loaded. The same report is in mlx #3755. If the link is broken, you want to know in seconds, not after a long weight load. The MLX docs also say to test locally first.

The test script

It checks three things: the world size is what you expect, all_sum gives the right answer, and how long all_sum takes for small, medium, and large messages.

import os
import sys
import time

import mlx.core as mx

EXPECTED = int(os.environ.get("EXPECTED_WORLD_SIZE", "0"))
BACKEND = os.environ.get("CHECK_BACKEND", "any")

world = mx.distributed.init(backend=BACKEND, strict=True)
rank, size = world.rank(), world.size()

def say(msg):
    if rank == 0:
        print(msg, flush=True)

if EXPECTED and size != EXPECTED:
    print(f"rank {rank}: world size is {size}, expected {EXPECTED}", flush=True)
    sys.exit(2)

# 1. Correctness: every rank adds its rank + 1. The sum must be 1 + 2 + ... + size.
x = mx.full((10,), rank + 1, dtype=mx.float32)
total = mx.distributed.all_sum(x)
mx.eval(total)
want = size * (size + 1) / 2
if not mx.all(total == want).item():
    print(f"rank {rank}: all_sum gave {total.tolist()[:3]}, expected {want}", flush=True)
    sys.exit(3)
say(f"ok: {size} ranks, all_sum correct ({want:.0f})")

# 2. Speed: time all_sum for a few sizes. This uses the same link the model will use.
for n_bytes in (4 * 1024, 1024 * 1024, 64 * 1024 * 1024):
    a = mx.ones((n_bytes // 4,), dtype=mx.float32)
    mx.eval(mx.distributed.all_sum(a))  # warm up
    runs = 10
    start = time.perf_counter()
    for _ in range(runs):
        mx.eval(mx.distributed.all_sum(a))
    ms = (time.perf_counter() - start) / runs * 1000
    say(f"all_sum {n_bytes / 1024:>8.0f} KiB: {ms:8.3f} ms per call")

Run it on one machine first

The ring backend over localhost needs no cables. This checks your Python setup and the script itself:

cd ~/mlx-cluster
EXPECTED_WORLD_SIZE=2 CHECK_BACKEND=ring \
  uv run mlx.launch -n 2 --backend ring --python "$PWD/.venv/bin/python" check_link.py

Output on our Linux test box (CPU only, mlx 0.32.4). Your Mac numbers will be different:

ok: 2 ranks, all_sum correct (3)
all_sum        4 KiB:    0.107 ms per call
all_sum     1024 KiB:    0.621 ms per call
all_sum    65536 KiB:   46.188 ms per call

Pass --python with the full path to the venv Python. Without it, mlx.launch runs the command as given, and on our box -- check_link.py failed with exec: check_link.py: not found.

Run it on the cluster

Environment variables from your shell do not reach other Macs. Pass them with --env:

cd ~/mlx-cluster
uv run mlx.launch --verbose --hostfile jaccl.json --backend jaccl \
  --env EXPECTED_WORLD_SIZE=4 --env CHECK_BACKEND=jaccl \
  --python ~/mlx-cluster/.venv/bin/python check_link.py

We do not have Macs on our test box, so we have no JACCL numbers to show you. The docs say JACCL latency is about an order of magnitude lower than the ring backend. Save your first good output as a baseline for each cable layout, and run the test again after every reboot, cable change, or macOS update.

Read the output, not the exit code

On our test box, when both ranks exited with code 2, mlx.launch printed warnings and still returned exit code 0:

[WARN] Node with rank 0 exited with code 2
[WARN] Node with rank 1 exited with code 2
rank 0: world size is 2, expected 3
rank 1: world size is 2, expected 3
$ echo $?
0

So a script that runs this test must look for the ok: line. Lesson 04 does that.

Size the test like the real job

• The 4 KiB time is mostly latency. Tensor parallel sends many messages this small for every token, so this number matters most for decode speed.
• The 64 MiB time shows bandwidth. Pipeline and long-prompt prefill depend more on this.
• If small messages are slow but large ones are fine, tensor parallel will be slow. Try pipeline.

Sources

• MLX docs: Distributed Communication (Getting Started, Tips and Tricks)
• MLX docs: mlx.core.distributed.all_sum
• MLX source: mlx.launch