ml-explore / ml-explore/mlx

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

Open
#3,149 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug distributed
Dominant language
C++
Stars
28.5k
Forks
2.3k
Avg merge
3d 8h
Merged PRs (30d)
62

Description

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

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the external reproduction repository's README and run test_fail.py and test_pass.py to confirm the JACCL behavior. Then trace the JACCL point-to-point send/recv implementation and its handling of consecutive calls with different shapes. Done means the reproduced sequences transfer correct data and do not hang without requiring an all_sum barrier.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
distributed-systems, networking
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.