Skip to content

Only use MPI backend for world size > 1 - #4210

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
erwinzhang7:fix-3442-any-size1
Aug 14, 2026
Merged

Only use MPI backend for world size > 1#4210
zcbenz merged 1 commit into
ml-explore:mainfrom
erwinzhang7:fix-3442-any-size1

Conversation

@erwinzhang7

@erwinzhang7 erwinzhang7 commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

Related to #3442.

init(backend="any") tries backends in order and takes the first that initializes. Open MPI
initializes successfully for a program that was not started with mpirun, giving an
MPI_COMM_WORLD of size 1. That counted as success, so the chain stopped at mpi and never
reached the backends after it. A correctly configured JACCL job would then run single node and
report nothing.

mpi::init now returns nullptr when the world holds a single process, so any carries on to
the next backend. Asking for it directly with backend="mpi" and strict=true raises instead:

[mpi] The world has a single process. Launch with mpirun to initialize the mpi backend.

The size comes from mpi().world() rather than from a constructed MPIGroup, because building
one and discarding it runs ~MPIGroup(), which calls finalize_safe() when global_ is set,
and MPI_Init cannot be called again after that.

The second hunk is needed to make the first safe. register_group only caches under "any" for
a non-null group, so once mpi::init can decline, nothing is cached and the next init() call
re-runs the whole chain. init_safe() called MPI_Init unconditionally, and Open MPI treats a
second MPI_Init as fatal:

Open MPI has detected that this process has attempted to initialize
MPI (via MPI_INIT or MPI_INIT_THREAD) more than once.  This is erroneous.
*** MPI_ERRORS_ARE_FATAL

That is unreachable on main today, because a successful mpi group is cached and init_safe never
runs twice. init_safe is idempotent now.

Verified

Open MPI 5.0.9, macOS 26.6.

case result
libmpi present, no mpirun, init() mpi not selected
libmpi present, no mpirun, init(strict=true) [distributed] Couldn't initialize any backend
libmpi present, init(backend="mpi", strict=true) raises the message above
libmpi present, init() twice no crash
mpirun -n 2, init() size 2, ranks 0 and 1
mpirun -n 2, init(backend="mpi", strict=true) size 2, no raise
mpirun -n 2, all_sum [3.0, 3.0, 3.0, 3.0]

The first two rows match running with MLX_MPI_LIBNAME pointed at a missing file, which is the
intent: with Open MPI installed but no mpirun, MLX behaves as though it were not there.

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@erwinzhang7

Copy link
Copy Markdown
Contributor Author

Reworked: mpi::init now returns nullptr when the
world has a single process, and the warning is gone.

One thing falls out of it. Returning nullptr there made a latent crash reachable:
register_group only caches under "any" for a non-null group, so nothing is cached, the next
init() re-runs the chain, and init_safe() called MPI_Init a second time, which Open MPI
treats as fatal. init_safe() is idempotent now, which is the second hunk.

Checked with Open MPI 5.0.9: without mpirun, init() no longer selects mpi and behaves the
same as with libmpi absent; init(backend="mpi", strict=True) raises; under mpirun -n 2,
init() still gives size 2 and all_sum is correct.

@zcbenz zcbenz changed the title Warn when a size one group hides a configured jaccl backend Only use MPI backend for world size > 1 Aug 14, 2026
@zcbenz
zcbenz merged commit adf21de into ml-explore:main Aug 14, 2026
28 checks passed
erwinzhang7 added a commit to erwinzhang7/mlx that referenced this pull request Aug 16, 2026
…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants