Implement Group.split() for the JACCL mesh backend - #4282
Open
erwinzhang7 wants to merge 1 commit into
Open
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Addresses #3205 for the mesh case.
Only MPI implemented split, by handing the problem to
MPI_Comm_split. Ring andJACCL both threw, so there was no in-tree example of a backend building the child
itself.
A
MeshGroupneeds three things: the caller's rank inside the child, the RDMAdevices 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:
SideChannelis 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
splitall_gathers(color, key)over the parent, selects the members sharing acolor, orders them by key and then by parent rank as MPI does, and builds the child
from the parent's own device row.
MeshGroupnow keeps that row, which it previouslytook 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_ptrto the parent, so a child keepsits 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:
membership and isolation both hold. Rank i contributes
2**i, so the sums nameexactly who took part and a child leaking into the parent would return 15
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, currentconfigure_jaccl_ringwritesnullfor every pair that is nota 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.