Only use MPI backend for world size > 1 - #4210
Conversation
angeloskath
left a comment
There was a problem hiding this comment.
I don't think that this is the correct fix. I would instead only return successful MPI initialization if the world size is larger than 1.
0bd40fe to
920c22d
Compare
|
Reworked: One thing falls out of it. Returning Checked with Open MPI 5.0.9: without |
…uses launch_jaccl checked the shape of the rdma matrix, its dimensions and its null diagonal, but never whether the pairs the backend is about to communicate over are actually in it. A hostfile whose rank order places unconnected nodes beside each other passes both checks. What follows is the worst failure available. The group does not fail to form and it does not run slowly: every rank ends up alone in a group of size one, reports success, and the job computes N separate wrong answers at full speed with nothing on stderr. Verified on four M4 Pro minis wired as a full mesh, where permuting two hosts in a working hostfile, rows and columns together so the file stays internally consistent, produced four independent groups of one and four 'membership OK' lines. This is the same class as ml-explore#4210, where a group of size 1 counted as a successful init. missing_jaccl_links reports the pairs that are used but absent: for a ring each rank's two neighbours, for a mesh every pair. With two ranks the previous and next neighbour are the same node, which is harmless because the pair is simply examined twice. The error names the hosts and says that for jaccl-ring the host order is the ring order, since that is the mistake that produces it. Checked against the live mesh: the misordered hostfile is now rejected, and valid mesh and valid ring hostfiles still run at 20.87 and 16.33 GB/s.
Related to #3442.
init(backend="any")tries backends in order and takes the first that initializes. Open MPIinitializes successfully for a program that was not started with
mpirun, giving anMPI_COMM_WORLDof size 1. That counted as success, so the chain stopped at mpi and neverreached the backends after it. A correctly configured JACCL job would then run single node and
report nothing.
mpi::initnow returnsnullptrwhen the world holds a single process, soanycarries on tothe next backend. Asking for it directly with
backend="mpi"andstrict=trueraises instead:The size comes from
mpi().world()rather than from a constructedMPIGroup, because buildingone and discarding it runs
~MPIGroup(), which callsfinalize_safe()whenglobal_is set,and
MPI_Initcannot be called again after that.The second hunk is needed to make the first safe.
register_grouponly caches under"any"fora non-null group, so once
mpi::initcan decline, nothing is cached and the nextinit()callre-runs the whole chain.
init_safe()calledMPI_Initunconditionally, and Open MPI treats asecond
MPI_Initas fatal:That is unreachable on main today, because a successful mpi group is cached and
init_safeneverruns twice.
init_safeis idempotent now.Verified
Open MPI 5.0.9, macOS 26.6.
init()init(strict=true)[distributed] Couldn't initialize any backendinit(backend="mpi", strict=true)init()twicempirun -n 2,init()mpirun -n 2,init(backend="mpi", strict=true)mpirun -n 2,all_sum[3.0, 3.0, 3.0, 3.0]The first two rows match running with
MLX_MPI_LIBNAMEpointed at a missing file, which is theintent: with Open MPI installed but no
mpirun, MLX behaves as though it were not there.