mx.distributed ring backend: gradient-shaped CPU-stream all_sum deadlocks probabilistically (~3-4%/call) across a training loop; GPU-stream variant crashes deterministically at ~5s

Summary

Running a realistic LoRA-training-shaped gradient reduction (mlx.nn.utils.average_gradients)
repeatedly across a --backend ring cluster of ≥3 ranks fails in one of two ways depending
on which stream the collective runs on, and neither is a training-code bug — both reproduce
in a self-contained script with no model, dataset, or third-party library involved:

• GPU (default) stream + a real-world timing skew between ranks: the ~5s Metal command-buffer
  watchdog kills the waiting rank deterministically (kIOGPUCommandBufferCallbackErrorTimeout),
  tearing down the ring.
• CPU stream (communication_stream=mx.cpu, the documented workaround for the above):
  the watchdog no longer fires, but the ring's socket exchange deadlocks outright on a
  measured ~3-4% of calls, independent of skew, tree shape, chunking, or a pre-barrier —
  ranks sit blocked forever in recvfrom/sendto inside libmlx.dylib. Over a training run
  of realistic length (hundreds of gradient syncs), this failure rate makes the run's survival
  probability effectively zero.

Both are visible in this project's real training loop (mlx_vlm's sft_trainer.py, an
independent third-party package, not the source of the bug) but reproduce identically in the
attached minimal script with a synthetic model and synthetic gradients — no mlx_vlm, no
dataset, no external dependency of any kind.

Environment

• Hardware: 4x Apple M4 Mac mini, connected via Gigabit LAN
• OS: macOS 26.3 (build 25D125)
• mlx: 0.32.2 (pip show mlx)
• Backend: mlx.launch --backend ring, numeric IPs, one process per host

Repro 1 — deterministic GPU-stream watchdog kill

# gpu_stream_repro.py — run: mlx.launch --hosts <ip1>,<ip2> --backend ring -- python3 gpu_stream_repro.py
import time
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from mlx.nn.utils import average_gradients

group = mx.distributed.init()
rank, world_size = group.rank(), group.size()

model = nn.Linear(1024, 1024)
mx.eval(model.parameters())
optimizer = optim.Adam(learning_rate=1e-3)
loss_and_grad_fn = nn.value_and_grad(model, lambda m, x: m(x).sum())
x = mx.random.normal((64, 1024))

loss, grad = loss_and_grad_fn(model, x)
grad = average_gradients(grad)              # default (GPU) communication_stream
optimizer.update(model, grad)                # queued lazily on GPU

if rank == 1:
    time.sleep(7)                            # simulate a real per-rank timing skew

t0 = time.perf_counter()
mx.eval(model.state)                         # <-- rank 0 dies here at ~5.0s:
print(f"[{rank}] eval done in {time.perf_counter()-t0:.2f}s")

Result: rank 0's mx.eval(model.state) raises kIOGPUCommandBufferCallbackErrorTimeout
at 5.01s (measured, repeatable) whenever a peer is delayed past ~5s before the collective
resolves — regardless of how small the actual payload is (this uses a single 1024x1024
linear layer, ~4MB of gradients).

Repro 2 — probabilistic CPU-stream deadlock

The documented fix for Repro 1 is to force the collective onto the CPU stream and eval it
explicitly before any GPU-stream node depends on it (materialize-then-reduce-then-update).
This does eliminate the watchdog kill — but repeating it in a loop, with a gradient tree
shaped like a real LoRA fine-tune (~500 small fp32 arrays, 32KB-384KB each, ~87MB total —
the exact shape average_gradients's default 32MiB chunking produces from a real
nn.value_and_grad output), deadlocks intermittently:

# cpu_stream_loop_repro.py — run at >=3 ranks:
# mlx.launch --hosts <ip1>,<ip2>,<ip3>[,<ip4>] --backend ring -- python3 cpu_stream_loop_repro.py [delay_s] [iters]
import sys, time
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from mlx.nn.utils import average_gradients
from mlx.utils import tree_flatten

group = mx.distributed.init()
rank, world_size = group.rank(), group.size()
delay_s = float(sys.argv[1]) if len(sys.argv) > 1 else 7.0
n_iters = int(sys.argv[2]) if len(sys.argv) > 2 else 30

# Synthesize a gradient tree matching a real ~22M-param LoRA adapter's shape:
# 3 groups of 195/195/114 fp32 arrays that average_gradients's 32MiB
# _group_by_size chunking produces from a real training step's grad tree.
sizes = []
for count, total_mb in ((195, 33.82), (195, 33.55), (114, 19.92)):
    total = int(total_mb * 1e6)
    tail = 280_000
    base = (total - tail) // (count - 1)
    sizes.extend([base] * (count - 1) + [total - base * (count - 1)])

model = nn.Linear(64, 64)
mx.eval(model.parameters())
optimizer = optim.Adam(learning_rate=1e-3)
loss_and_grad_fn = nn.value_and_grad(model, lambda m, x: m(x).sum())
x = mx.random.normal((32, 64))

for it in range(1, n_iters + 1):
    _, g = loss_and_grad_fn(model, x)
    mx.eval(g)                                          # Phase 1: materialize local GPU work
    grads = [mx.random.normal((s // 4,)) for s in sizes]
    mx.eval(grads)

    if it == 1 and rank == 1 and delay_s > 0:
        time.sleep(delay_s)                              # optional: simulate arrival skew

    t0 = time.perf_counter()
    with mx.stream(mx.cpu):                               # Phase 2: CPU-stream-only reduction
        reduced = average_gradients(grads, communication_stream=mx.cpu)
        mx.eval(reduced)
    mx.synchronize(mx.cpu)
    print(f"[{rank}] iter {it}: reduction done in {time.perf_counter()-t0:.2f}s", flush=True)

    pflat = tree_flatten(model.parameters())              # Phase 3: GPU update, unrelated to the bug
    grads_tree, offset = {}, 0
    for k, p in pflat:
        grads_tree[k] = reduced[0][offset: offset + p.size].reshape(p.shape)
        offset += p.size
    optimizer.update(model, grads_tree)
    mx.eval(model.state)

print(f"[{rank}] all {n_iters} iterations completed")

Observed across our own runs at 4 ranks (master + 3 workers), same script/config, repeated:

Cumulative: 4 hangs across ~124 successful CPU-stream collectives at 4 ranks (~3-4%
per collective), independent of every variable we controlled for: gradient tree shape
(single tensor / 6 large arrays / ~500 small arrays), chunking (default 32MiB grouping vs
forcing one flat group), rank-arrival skew (7s vs none), and a rendezvous barrier
immediately before the reduction. At 2 ranks we have never observed a hang across any run.
When a hang occurs, sample <pid> on every host shows the ring's worker threads blocked in
__recvfrom/__sendto inside libmlx.dylib, with no CPU spin and no forward progress —
this is a genuine stuck-forever deadlock, not a slow collective (confirmed by waiting well
past the point any successful collective in the same run had ever taken).

At a ~3.5%/collective failure rate, a training run of realistic length (roughly 800+
gradient syncs for one epoch of a modest fine-tune) has a survival probability on the order
of 1e-13 — which is consistent with what we observed operationally: every real multi-hour
distributed training attempt eventually died at some point in the run, never at a fixed
iteration, while short smoke tests (a few dozen iterations) usually didn't hit it.

What we ruled out

• Not a tree-shape issue: single tensor, 6 large arrays, and ~500 small arrays (matching a
  real LoRA adapter's exact shape) all pass individually; only repeating the reduction in a
  loop surfaces the hang, and even then non-deterministically.
• Not the default 32MiB chunking specifically: forcing one flat 87MB group (single
  concatenate → single all_sum → single split, via average_gradients(..., all_reduce_size=2**30))
  hangs the same way.
• Not purely a rank-arrival-skew race: a scalar pre-barrier all_sum immediately before the
  big reduction visibly absorbs the arrival skew (0.004-0.872s barrier times observed) and
  the very next big reduction still deadlocks.
• Not skew-dependent at all: a fully skew-free loop (time.sleep removed entirely) also
  hung, at collective ~37 of an intended ~90.

Ask

This looks like a race condition inside the ring backend's socket-level exchange for
larger payloads (tens of MB) under mx.distributed.all_sum/average_gradients, independent
of which mx.Stream the call is issued on — the GPU-stream case just fails faster and more
visibly (watchdog kill) than the CPU-stream case (silent deadlock). We're happy to share the
full test harness (adds a few more probe modes on top of the two scripts above — chunking
overrides, pre-barrier toggle, etc.) if useful for reproducing on your end.

Start by running gpu_stream_repro.py and cpu_stream_loop_repro.py with the documented mlx.launch commands, then trace mx.distributed.all_sum and average_gradients through the ring backend. Reproduce both the GPU watchdog timeout and repeated CPU-stream hang across at least three ranks. Done means the supplied multi-rank loops complete without timeout or deadlock.