Skip to content

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

New issue
New issue
Open
Open
[BUG] JACCL RDMA consecutive send/recv point-to-point with different shape produces wrong data or hangs.#3149
Labels
bugdistributed

Description

@dev-tb5-tester
dev-tb5-tester
opened on Feb 19, 2026

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

Activity

  1. angeloskath commented on Feb 19, 2026

    @angeloskath
    angeloskath
    on Feb 19, 2026
    Member

    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.

  2. dev-tb5-tester commented on Feb 19, 2026

    @dev-tb5-tester
    dev-tb5-tester
    on Feb 19, 2026
    Author

    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.

  3. angeloskath commented on Feb 19, 2026

    @angeloskath
    angeloskath
    on Feb 19, 2026
    Member

    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...

  4. dev-tb5-tester commented on Feb 20, 2026

    @dev-tb5-tester
    dev-tb5-tester
    on Feb 20, 2026
    Author

    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...

  5. zcbenz
    added
    distributed
    on Mar 19, 2026

  6. erwinzhang7 commented on Aug 12, 2026

    @erwinzhang7
    erwinzhang7
    on Aug 12, 2026
    Contributor

    Tried to reproduce this on current main (596dc79, 0.32.1) and could not across 165 runs.
    Posting since the issue is still open and may be worth a retest.

    Setup: test_fail.py from your repo, unmodified except that init() names the backend, which
    your hostfile does for you. Thunderbolt 5, RDMA enabled, IPv4 on the TB ports (#3467).

    2x M4 Pro mini, simultaneous and staggered start        80 runs, 0 failures
2x M4 Pro mini, 12 induced-skew configs                 36 runs, 0 failures
M5 Max <-> M4 Pro                                       25 runs, 0 failures
M5 Max <-> M4 Pro, 8 induced-skew configs               24 runs, 0 failures

    Since @angeloskath's read was an init-synchronisation race, I tried to force it rather than
    just repeat the test: delays from 0.1s to 1s applied to each rank in turn, and nice -n 19 on
    one rank so its initialisation runs slower rather than merely later. Nothing failed.

    He also flagged heterogeneous systems as untested. The M5 Max <-> M4 Pro pair covers that
    (different generation, core counts and memory bandwidth) and it is clean there too.

    This was filed against 0.30.6 and there has been a lot of distributed work since, so the most
    likely explanation is that it was fixed somewhere in between. If you still have the cluster,
    I'd recommend retesting on current main.

  7. zcbenz
    added
    bug
    on Sep 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugdistributed

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions