[BUG] JACCL RDMA consecutive send/recv point-to-point with different shape produces wrong data or hangs.

[dev-tb5-tester]

Describe the bug
Consecutive send/recv point-to-point calls between two nodes with different shapes fail on the JACCL backend. Either I get wrong data sent/received or hang my system. It seems that adding an all_sum barrier before the send/recv fixes it. Ring backend always works fine with the same complex sequences of communications.

To Reproduce
I put in a repo the detailed setup instructions on my system, using conda and the mlx config setup. I added all the hostfiles generated by the setup and two test scripts (test_fail.py / test_pass.py):
https://github.com/dev-tb5-tester/rdma_jaccl

Expected behavior
Complex series of send/recv on jaccl backend are not currently working as expected. Prepending a barrier before each couple of point-to-point communication should not be required.

Desktop (please complete the following information):

• OS Version: MacOS 26.3
• Environment: M4 Max + M4 Pro, Thunderbolt 5, RDMA enabled.
• MLX version: 0.30.6

Additional context
Please review reproduction details and results on README

[angeloskath]

That's interesting. I think the bug has to do with synchronization of the initialization and not the send/recv. I did get occasional hangs when ran as is. Then putting a time.sleep(5) before the communication starts makes it always pass. I also never got wrong results 🤔

I 'll investigate further.

[dev-tb5-tester]

Thank you! Do I need to put a sleep only right after the first initialization group = mx.distributed.init() of my distributed processes or every time we retrieve the singleton? For example, in a distributed server, we would start the initialization with the backend specified by mlx-launch, and then the subsequent call to mx.distributed.init() would return that group - in other part of the application. I would expect the latter not to require either a sleep or a barrier to bring all the processes at the same timeline.

That's interesting. I think the bug has to do with synchronization of the initialization and not the send/recv. I did get occasional hangs when ran as is. Then putting a time.sleep(5) before the communication starts makes it always pass. I also never got wrong results 🤔

I 'll investigate further.

[angeloskath]

No only after the first time. To be fair that is still a bug and that 's not a fix... maybe a stop-gap, we 'll get a proper fix. But it would be helpful to check if that works for you too. You are also on heterogeneous systems which I haven't tried. It shouldn't be a problem but you never know...

[dev-tb5-tester]

No only after the first time. To be fair that is still a bug and that 's not a fix... maybe a stop-gap, we 'll get a proper fix. But it would be helpful to check if that works for you too. You are also on heterogeneous systems which I haven't tried. It shouldn't be a problem but you never know...

Thanks! I can confirm that your hunch is correct. I tested your idea of allowing the latency of initialization in heterogeneous hardware setup. Rather than setting a sleep which could be hardware dependent, I added a barrier similar to the one of test_pass.py. So I have a full suite of unit tests that does 14 sequence of point to point communication between controller and worker, on ints, float, strings, bytes and with different size payloads. They are all passing in JACCL (as it used in RING backend)!

But here is where things get really interesting. If I use a single element float array as the barrier, things go bad in exactly the same way - I can share the below trace:

[Rank 0] array of 5 integers sending  : [1, 2, 3, 4, 5], to 1 node(s)

[Rank 1] array of 1 integers receiving: [5], from Rank 0
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/Users/quantum/dev/mlx_mpi_server/scripts_mpi/test_mpi_send_recv.py", line 80, in <module>
    assert_int_list_eq(got_ints, EXP_INTS, ctx=f"[Rank {my_rank}] Slave_RecvInts")
  File "/Users/quantum/dev/mlx_mpi_server/scripts_mpi/test_mpi_send_recv.py", line 30, in assert_int_list_eq
    assert got == exp, f"{ctx}: expected {exp}, got {got}"
           ^^^^^^^^^^
AssertionError: [Rank 1] Slave_RecvInts: expected [1, 2, 3, 4, 5], got [5]
[WARN] Node with rank 1 exited with code 1 

As you see it is not only hanging but the array is sent all wrong, so this barrier does not work:
mx.eval(mx.distributed.all_sum(mx.array([1.0]), group=group, stream=mx.cpu))

Instead the barrier that works for me is a mx.ones(10). I have no idea why and maybe you know instead - if the single element or the integer vs float???
Anyway what I need is the the following, which is the same code of test_pass.py in my repo:
mx.eval(mx.distributed.all_sum(mx.ones(10), group=group, stream=mx.cpu))

So I totally think there is a delay initialization issue. it's not a big deal for me as long as I know how to work around it, though it would be nice if we could have feedback from mx.distributed.init() that the initialization is indeed finished and we can proceed...