Skip to content

Implement Group.split() for the ring backend - #4218

Open
erwinzhang7 wants to merge 1 commit into
ml-explore:mainfrom
erwinzhang7:ring-split
Open

Implement Group.split() for the ring backend#4218
erwinzhang7 wants to merge 1 commit into
ml-explore:mainfrom
erwinzhang7:ring-split

Conversation

@erwinzhang7

Copy link
Copy Markdown
Contributor

Part of #3205. RingGroup::split threw, so mx.distributed.init().split(...) was unavailable
on ring. MPI and NCCL both implement it, so this brings ring to the same contract:

  • ranks sharing a color end up in the same group,
  • new ranks are ordered by ascending key, with the parent rank breaking ties,
  • key < 0 (the default) means keep the parent's order.

Two things made this smaller than expected.

TCPSocket::listen already sets SO_REUSEADDR and SO_REUSEPORT, and the parent's listening
sockets are closed once accept() returns, so a subgroup rebinds the same hostfile addresses.
No port negotiation is needed.

RingGroup is constructed from the full address table, so once every rank knows every color it
can work out the subgroups by itself. The constructor was dropping that table, it is kept now.
The colors and keys rotate around the existing ring, so the connection side channel is not
involved.

RingGroup also returns early at size one instead of trying to connect to itself, which is what
a color nobody else shares produces. Collectives over a group that size are resolved in
ops.cpp before they reach a backend.

test_groups in ring_test_distributed.py asserted that split raises. It now covers the
behaviour instead.

Verified

Ring backend on localhost at 2, 4 and 8 ranks. Beyond rank and size, every case checks an
all_sum inside the subgroup to confirm the collective is carried by the subgroup and not the
parent.

2 ranks 4 ranks 8 ranks
even/odd split pass pass pass
uneven blocks (rank // 3) pass pass pass
nested split of a subgroup pass pass pass
identity split reproduces the world pass pass pass
singleton subgroups pass pass pass
key-driven reordering pass pass pass
ring_test_distributed.py 13/13 13/13 13/13

@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants