[BUG][Metal] Deadlock in fence

[nastya236]

Overview of the bug:

There were issues reported about deadlocks while using MLX_METAL_FAST_SYNCH=1:

• [BUG] GPU locking using METAL_FAST_SYNCH=1 and the JACCL backend #3142
• [BUG] Distributed pipeline generation: Metal fence handoff deadlocks under MLX_METAL_FAST_SYNCH=1 (TCP ring; orphaned fence_wait kernel locks GPU until reboot) and hits the ~5s GPU watchdog when unset (kIOGPUCommandBufferCallbackErrorTimeout at ~7.3k tokens) #3830

The hypothesis was that stall happens because memory coherence is not guaranteed.

After some debugging, I noticed that the problem happens when we have a lot of independent reductions in one tape (More about it later).

CPU-GPU synchronization:

We do communication operations on cpu, compute (usually) on gpu. Gpu should mark inputs ready for cpu to do collective communication operation. With MLX_METAL_FAST_SYNCH=1 we can sync fast between cpu and gpu using counter in shared memory.
How operations were scheduled before this PR:

Main thread traverses the graph and constructs a tape, encodes operations in to gpu command buffer and cpu queue. To synchronise, we use one fence per stream, both fences have counters to track how many fence::update() were done -- how many inputs for consumer stream are ready. While encoding main thread tracks number of fence::update() scheduled by the producer stream.

Important part is that consumer stream is waiting on a counter that is equal to this scheduled number of updates, independently from the number of fence::updates() finished. To be concrete, suppose the graph of operations is:

x = M@a
y = N@b
all_reduce(x)
all_reduce(y) 

mx.eval(x, y)

In the above case we do independent reductions.
Then command buffer will look like: [matmul, matmul, fence::update(), fence::update(), fence::wait(), fence::wait(), ..]
Number of fences is 2 -- main threads counter is 2.
CPU queue would look like: [(while C<2) {}, all_reduce(), all reduce() ]

Basically, cpu won't start the execution of all reduce before gpu is finished preparing inputs for all all_reduce that are in the same tape. And this is exactly where stall is happening.

Debugging:

I was able to find a case when this bug is reproducible:

import mlx.core as mx
world = mx.distributed.init(strict=True, backend="jaccl")

num_grads = 32
dim = 4096
steps = 100
gpu_rounds = 8

def gpu_work(x, weight):
    for _ in range(gpu_rounds):
        x = x @ weight
    return x

x = mx.random.normal((dim, dim), dtype=mx.float32) + world.rank()
weight = mx.random.normal((dim, dim), dtype=mx.float32) / dim**0.5
mx.eval(x, weight)

for step in range(steps):
    grads = [gpu_work(x + i + step, weight) for i in range(num_grads)]
    reduced = [mx.distributed.all_sum(g, group=world) for g in grads]
    outputs = [gpu_work(g / world.size(), weight) for g in reduced]
    mx.eval(outputs)
    print(f"rank {world.rank()} step {step}: done", flush=True)

gpu work and back to back reduction.
I was debugging it on a cluster on 4 M5 ultras, connected with a ring of 3 cables per neighbour.
Steps to identify this stall with the above example:

1. ps -Ao pid,etime,%cpu,command | grep <script> — show the asymmetry between cpu utilisation

Stalled node: CPU utilisation is 100% -- 1 stalled system thread, idle jaccl ring threads
Healthy nodes: CPU utilisation is 300% -- 3 threads: system thread + 2 jaccl thread pool threads.

2. sample pid 3 -mayDie -f /tmp/out.txt — ttaches to a live process and captures the call stack of every thread, shows exact stalls on cpu threads

Stalled node: mlx::core::Fence::wait(mlx::core::Stream, mlx::core::array const&)::$_1 on system cpu thread, jaccl threads are idle -> we have not entered all reduce.
Healthy nodes: <ring_pass> -- healthy nodes are at all_reduce waiting for a stalled node

So based on the above, cpu system threads has not entered the all_reduce and waiting on the fence that was scheduled to wait on the input computed by gpu.

3. The last part is to identify the value for 3 counters: scheduled fences on a main thread, G -- gpu counter, number of fence::updates that GPU is done with, C -- cpu counter,

(here I asked Claude to help me and disassemble the code and identify registers where these counters are stored)
lldb -p <pid>
then we need to attache to stalled thread (for this we do thread list -> identify stalled one)

  (lldb) thread select <id>
  (lldb) p/u (unsigned)($x20+4) <- main's thread scheduled 
  (unsigned int) 12
  (lldb) p/u (unsigned)(void*)[(id)(void*)($x20+8) contents]
  (unsigned int) 8 <- gpu counter 
  (lldb)p/u *(unsigned*)(0x6000012ab340+0x1c)
  (unsigned int) 0 <- cpu counter

What above tells us: main thread scheduled 12 fence::updates. GPU is done with 8 of the inputs. CPU has 0 -- 0 all_reduces are done.

GPU is stuck in the middle of inputs preparation for some reason. It is not clear whether GPU updates 9–12 never executed or whether the CPU cannot see their writes, and the counters alone don’t distinguish those cases.

Sum up.

In general, this is incorrect. Every fence::wait should be waiting for a specific fence::update. This pull request fixes the scheduling and correct scheduling fixes the stall.

[zcbenz]

Very good job locating the deadlock!