Skip to content

Implement Group.split() for the JACCL mesh backend - #4282

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

Implement Group.split() for the JACCL mesh backend#4282
erwinzhang7 wants to merge 1 commit into
ml-explore:mainfrom
erwinzhang7:jaccl-group-split

Conversation

@erwinzhang7

Copy link
Copy Markdown
Contributor

Addresses #3205 for the mesh case.

Only MPI implemented split, by handing the problem to MPI_Comm_split. Ring and
JACCL both threw, so there was no in-tree example of a backend building the child
itself.

A MeshGroup needs three things: the caller's rank inside the child, the RDMA
devices reaching the other members, and a side channel. The first two are arithmetic
once membership is known. The third looked like the hard part and is not:
SideChannel is defined by an all-gather function, so the parent can supply it,
which is the same shape as MPI bootstrapping a child communicator from its parent.
No new coordinator and no second discovery pass.

So split all_gathers (color, key) over the parent, selects the members sharing a
color, orders them by key and then by parent rank as MPI does, and builds the child
from the parent's own device row. MeshGroup now keeps that row, which it previously
took and discarded.

Children of different sizes stay in step because the child's all-gather is a
parent-wide collective: every rank participates in every round regardless of which
child it joined, and the length negotiation inside the parent's container all_gather
covers the whole parent at once, so two children can never disagree about how many
bytes are in flight. The closure holds a shared_ptr to the parent, so a child keeps
its parent alive.

A child can only contain directly connected members, and the hostfile is the only
description of connectivity the runtime sees, so a missing pair raises naming both
ranks rather than failing later inside the transport.

Ring reports why instead of the generic message: a subset of a ring is only a ring
when its members are contiguous in ring order, and there is no route around the gap
when they are not. So this is mesh only, deliberately.

Verified on hardware

Four M4 Pro minis wired as a full mesh, RDMA over Thunderbolt, MLX built from this
branch and installed on all four:

  • splitting by parity gives children summing 5 and 10 where the parent sums 15, so
    membership and isolation both hold. Rank i contributes 2**i, so the sums name
    exactly who took part and a child leaking into the parent would return 15
  • the parent still sums 15 afterwards
  • a child of a child works
  • an uneven 3/1 split works, including the child of size one
  • a key reverses the order within a child without losing members
  • a two rank child reaches 8.62 GB/s, the same as a natively configured two rank
    group, so the child uses its own RDMA connections rather than relaying through the
    parent

Depends on #4269 for the mesh case

A child can only be built from pairs the hostfile recorded. On a mesh configured as
jaccl-ring, current configure_jaccl_ring writes null for every pair that is not
a ring neighbour, so a third of the connectivity is discarded before split is ever
called. Measured on this hardware, same command, only the config code differing:
main records 8 of 12 directed pairs, #4269 records 12 of 12.

#4269 is byte identical to main on a genuine ring, so it adds where there is
something to add and changes nothing otherwise.

Only MPI implemented split, by handing the problem to MPI_Comm_split. Ring and
JACCL both threw, so there was no in-tree example for a backend that has to
build the child itself.

A MeshGroup needs three things: the caller's rank inside the child, the RDMA
devices reaching the other members, and a side channel. The first two are
arithmetic once the membership is known. The third looked like the hard part
and is not: SideChannel is defined by an all-gather function, so the parent can
supply it, which is the same shape as MPI bootstrapping a child communicator
from its parent. No new coordinator, no second discovery pass.

So split all_gathers (color, key) over the parent, selects the members sharing
a color, orders them by key and then by parent rank as MPI does, and builds the
child from the parent's own device row. MeshGroup now keeps that row, which it
previously took and discarded.

Children of different sizes stay in step because the child's all-gather is a
parent-wide collective: every rank participates in every round regardless of
which child it joined, and the length negotiation inside the parent's container
all_gather covers the whole parent at once, so two children can never disagree
about how many bytes are in flight. The closure holds a shared_ptr to the
parent, so a child keeps its parent alive.

A child can only contain directly connected members, and the hostfile is the
only description of connectivity the runtime sees, so a missing pair raises
naming both ranks rather than failing later inside the transport.

Ring reports why instead of the generic message: a subset of a ring is only a
ring when its members are contiguous in ring order, and there is no route
around the gap when they are not.

Verified on four M4 Pro minis wired as a full mesh, MLX built from this branch
and installed on all four. Splitting by parity gives children summing 5 and 10
where the parent sums 15, so membership and isolation both hold; the parent
still sums 15 afterwards; a child of a child works; an uneven 3/1 split works
including the child of size one; a key reverses the order without losing
members. A two rank child reaches 8.62 GB/s, the same as a natively configured
two rank group, so the child uses its own RDMA connections rather than relaying
through the parent.
@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 16, 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