DeepSpeed cannot start without mpi4py on a machine with no launcher

[alanhuangyoo]

python train.py on a machine with one accelerator and no launcher ends before deepspeed.initialize returns:

ModuleNotFoundError: No module named 'mpi4py'

Reproduced on two machines that have nothing else in common — an Apple M5 (MPS, ZeRO 0/1/2/3) and an H20 (CUDA, ZeRO-1). Neither had mpi4py installed, which is the default state of a machine that was never going to run an MPI job.

Why

init_distributed fills in the distributed environment when a launcher did not:

required_env = ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
if auto_mpi_discovery and not all(map(lambda v: v in os.environ, required_env)):
    ...
    else:
        mpi_discovery(distributed_port=distributed_port, verbose=verbose)

and mpi_discovery opens with from mpi4py import MPI. It is the only route out of that branch, so a machine with no launcher and no MPI has to have mpi4py installed to get past it — and the error it gets instead says nothing about what to do.

That case is not exotic. It is DeepSpeed on a workstation or a laptop, which the MPS accelerator added in #8293 exists to serve, and it is also the first thing anyone hits trying DeepSpeed on a single GPU without deepspeed/torchrun.

comm/utils.py already has the answer to what that situation means:

def get_world_rank_from_launcher():
    rank = os.environ.get('RANK')
    if rank is None:
        rank = os.environ.get('OMPI_COMM_WORLD_RANK')
    # Make it a single process job and set rank to 0
    if rank is None:
        rank = 0

No launcher is a single-process job. That reading just never reached the environment the backend is initialized from.

The change

single_process_discovery sets RANK=0, WORLD_SIZE=1, LOCAL_RANK=0, MASTER_ADDR=127.0.0.1 and MASTER_PORT, with setdefault so a partially populated environment is left as the caller set it.

mpi_discovery still runs first and unchanged, and only an ImportError out of it is caught — so nothing changes for anyone who has mpi4py, including an unrecognized MPI launcher. If that ImportError fires while one of OMPI_COMM_WORLD_RANK, PMI_RANK, PMIX_RANK, MV2_COMM_WORLD_RANK or SLURM_PROCID is set, the run stops with a message naming mpi4py rather than silently becoming one rank of a many-rank job — which is the one outcome worse than today's crash.

Before and after

Apple M5, MPS, no launcher, no environment variables:

ZeRO stage	master	this PR
0	ModuleNotFoundError: No module named 'mpi4py'	trains, losses=[0.05809, 0.04466, 0.03354]
1	same	same losses
2	same	same losses
3	same	same losses

The losses are identical to the same script run with RANK/WORLD_SIZE/... set by hand, which is the path this PR does not touch.

H20, CUDA, no launcher:

master:   ModuleNotFoundError: No module named 'mpi4py'
this PR:  RESULT: trained on cuda:0 world 1

Tests

tests/unit/comm/test_single_process_discovery.py, 9 cases, CPU-only and no distributed context needed: an environment with no MPI variables is not an MPI job; each of the five launcher rank variables marks one; the fallback fills all five variables and leaves any the caller already set; and an MPI job whose mpi4py is missing raises with a message naming mpi4py and does not fill the environment.

On master the file cannot be imported, so it fails there for the trivial reason rather than the interesting one. The evidence that matters is the before/after table above, which is a real deepspeed.initialize on two accelerators.

tests/unit/comm/ passes on the branch. yapf clean; flake8 reports 51 pre-existing F824 global ... is unused warnings on comm.py, identical in count on master.

[alanhuangyoo]

Two things I checked afterwards that bear on the review.

mpi4py is an optional extra, not a core dependency. It appears in exactly one place in the requirements:

requirements/requirements-1bit-mpi.txt:1:mpi4py

requirements/requirements.txt does not list it. So a correctly installed DeepSpeed does not have it, and yet the no-launcher path cannot get past init_distributed without it. That inconsistency is the bug in one line.

Why CI never caught it. DistributedTest sets the environment itself:

# tests/unit/common.py
set_dist_env = True
...
if self.set_dist_env:
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = str(master_port)
    ...
    os.environ['RANK'] = str(local_rank)

so the all(v in os.environ for v in required_env) guard is satisfied and the MPI branch is skipped for essentially every test. The one test that sets set_dist_env = False, TestDistInitNoEnv, calls torch.distributed.init_process_group before deepspeed.init_distributed, so it returns early and does not reach the branch either — I ran it both ways to be sure:

master:   1 passed
this PR:  1 passed

It is a real test of a real thing, it just is not a test of this. The no-launcher-and-no-mpi4py combination has no coverage, which is what the new file adds.

tests/unit/comm/ on the branch: 49 passed, 5 skipped. There are 4 errors in that directory on this machine; I am running them against master as well to confirm they are pre-existing rather than mine, and will post the comparison.

(Editing this comment — my shell ate the inline code spans in the paragraph above when I first posted it.)

[alanhuangyoo]

The comparison I promised. Same machine, same command, only comm.py differing:

master:   40 passed, 5 skipped, 4 errors   in 324.92s
this PR:  49 passed, 5 skipped, 4 errors   in 304.67s

The same four errors on both sides, so they are pre-existing on this machine and not something this PR introduces. The 40 -> 49 is exactly the nine cases in the new file.

Also ran the new file on its own on both machines: 9 passed on the H20 (Linux, CUDA) and 9 passed on the Mac (macOS, MPS).

[ebarkhordar]

I ran this against 5ef94c45e100ff9f53da880caa5a45b2693803e6 in a clean container (python:3.12-slim, CPU torch, no mpi4py, DS_ACCELERATOR=cpu).

in_mpi_job() keys on rank variables, and SLURM_PROCID is one of them. If srun -n1 python train.py sets it, which is my reading of srun's environment rather than something I could run here, then a one-task step is a single process on one device and takes the raise instead of the path this PR adds:

# SLURM_PROCID=0, SLURM_NTASKS=1, no launcher vars, no mpi4py
in_mpi_job() -> True
RESULT: ImportError -> An MPI job is running but mpi4py is not installed, so the rank
and world size cannot be discovered from it. Install mpi4py, or set RANK, WORLD_SIZE, ...

# identical, minus SLURM_PROCID
in_mpi_job() -> False
ENV AFTER -> {'RANK': '0', 'WORLD_SIZE': '1', 'LOCAL_RANK': '0',
              'MASTER_ADDR': '127.0.0.1', 'MASTER_PORT': '29500'}

A rank variable says a launcher is present, not that the world is bigger than one. The size variables sitting next to them do say it: SLURM_NTASKS, OMPI_COMM_WORLD_SIZE, PMI_SIZE. Would keying the raise on one of those being above 1 be closer to what you mean here, so a one-task step still gets the single-process path?

[alanhuangyoo]

You are right, and the distinction you drew is the one I got wrong: a rank variable says a launcher is present, not that the world is bigger than one. srun -n1 is a launcher and a single process at the same time, and that is exactly the case this PR exists to serve, so keying the refusal on SLURM_PROCID sent it to the error instead of the fallback.

Changed to read the size variables the same launchers export alongside their rank ones — OMPI_COMM_WORLD_SIZE, PMI_SIZE, PMIX_SIZE, MV2_COMM_WORLD_SIZE, SLURM_NTASKS — and refuse only when one of them is above 1. Same machine as before, no mpi4py:

srun -n1      (SLURM_PROCID=0 SLURM_NTASKS=1)              trains, losses=[0.05809, 0.04466, 0.03354]
srun -n4      (SLURM_PROCID=0 SLURM_NTASKS=4)              ImportError naming mpi4py
mpirun -n2    (OMPI_COMM_WORLD_RANK=0 SIZE=2)              ImportError naming mpi4py
no launcher                                                trains, same losses

Four tests added for it: each size variable above 1 marks a multi-rank job, each one at 1 does not, a rank variable on its own does not, and an unparseable size does not. Plus the end-to-end one for your case — SLURM_PROCID=0, SLURM_NTASKS=1, mpi4py absent — asserting the fallback is reached rather than the raise. 17 passing in the file.

PMIX_SIZE I have kept on the list for symmetry with the rank set, though I have not confirmed PMIx exports it under that name and could not test it here. If you know it is spelled differently I will fix it; a wrong name there fails safe, since it only means that launcher is read from whichever of the others it also sets.

Thanks — you also read this in a clean container without mpi4py, which is the setup that makes the bug visible at all, and that is the second time on my PRs that the container is what caught the thing I could not see from my own machine.

[ebarkhordar]

I read 114ba85 rather than ran it, and keying on the size variables is right for the srun -n1 case.

On PMIX_SIZE: it is not an environment variable, so it is not a spelling to correct. I cloned openpmix (c1e1744) and PRRTE (bfab800) and enumerated every literal name either passes to a setenv variant. The only PMIx names that reach a launched process are PMIX_RANK and PMIX_NAMESPACE, set by the server fork setup in src/server/pmix_server.c:1617 and :1624, and nothing with SIZE, NPROC, NTASK or WORLD in the name is set anywhere in either tree:

$ grep -rhoiE '[a-z_]*setenv\(\s*&?"[A-Z0-9_]+"' --include='*.c' --include='*.h' . \
    | grep -oE '"[A-Z0-9_]+"' | sort | uniq -c | sort -rn | head -3
     17 "PMIX_RANK"
     16 "PMIX_NAMESPACE"
      9 "PMIX_SERVER_URI51"

$ grep -rniE '[a-z_]*setenv\(\s*&?"[A-Z0-9_]*(SIZE|NPROC|NTASK|WORLD)[A-Z0-9_]*"' \
    --include='*.c' --include='*.h' . | wc -l
0

The PMIX_SIZE a grep does find is the size_t data type, registered as PMIX_REGISTER_TYPE("PMIX_SIZE", PMIX_SIZE, ...) in src/mca/bfrops/*/bfrop_pmix*.c.

That matters beyond the dead entry, because PMIX_RANK was the one covering a job launched directly by PMIx (prterun, prun). srun --mpi=pmix still sets SLURM_NTASKS and Open MPI's mpirun still sets OMPI_COMM_WORLD_SIZE, so those stay covered. The case that changes is a multi-rank prterun job with no mpi4py: it used to raise, and now no size variable is set at all, so it reaches single_process_discovery and every rank comes up as world size 1.

I do not think the environment can settle that one, since prterun -n1 sets PMIX_RANK too and there is no size variable to separate the two. Dropping PMIX_SIZE and saying in the comment that a PMIx-direct launch is not detected would at least be accurate about the boundary.

[alanhuangyoo]

That is a better piece of work than the PR it is reviewing. I took the openpmix and PRRTE finding on your evidence rather than re-running the greps — the PMIX_REGISTER_TYPE("PMIX_SIZE", PMIX_SIZE, ...) explanation is the kind of thing that is obviously right once someone points at it, and I had put the name in for symmetry with the rank set without checking that it exists.

The part I want to push back on slightly is the conclusion. Dropping PMIX_SIZE and documenting the boundary is accurate, but it leaves the PMIx-direct case doing the thing this PR is meant to prevent: prterun -n4 with no mpi4py silently becoming four world-size-1 runs. You are right that the environment cannot separate prterun -n4 from prterun -n1. But the two failures are not symmetric — silently splitting a four-rank job is unrecoverable and looks like a slow run, while refusing a one-rank job costs an error message that names mpi4py and says what to set.

So 90820bf reads sizes when a launcher reports one, and refuses when a launcher is present and reports none:

no launcher                        trains
srun -n1      SLURM_NTASKS=1       trains
srun -n4      SLURM_NTASKS=4       ImportError naming the size
mpirun -n1    OMPI_..._SIZE=1      trains
mpirun -n2    OMPI_..._SIZE=2      ImportError naming the size
prterun       PMIX_RANK, no size   ImportError naming the missing size

PMIX_SIZE is gone; PMIX_RANK is back, but only in the rank list that answers "is a launcher present", never as a stand-in for the size. The message for that case says explicitly that a PMIx-direct launch reports no size, so someone hitting it on prterun -n1 knows why and what to do.

If you think refusing prterun -n1 is the wrong side of that trade I will switch it to your version — it is a judgement about which failure a user would rather have, and you have looked at this launcher landscape more carefully than I have.

One other thing your comment caused, indirectly. Building the three-way split, I shadowed init_distributed's own world_size parameter with a local of the same name and handed None to the backend, which broke the plain no-launcher case entirely. All 20 tests still passed, because every one of them stopped at the helpers or at a path where the shadowed value happened to be a valid int. I found it running the six-case matrix by hand. There is now a test that goes through init_distributed with a bare environment and fails on that version with world_size must be an integer. None.

[alanhuangyoo]

The red modal-torch-latest / DeepSpeedAI CI here is a GPU reservation timeout, not a test result. Both runs on this branch ended the same way:

__main__.SandboxStartTimeout: Sandbox did not start within 1800s, so no test ran.
This is a capacity problem rather than a test failure: the GPU reservation was
never satisfied.

collect tests and DCO are green on the same commit, and I re-pushed onto current master (a2e9052) partly to see whether it would clear — it did not. Not asking anyone to do anything about it; flagging it so the X does not read as a failing test.

The open question from earlier still stands, and it is the only thing I think needs a decision: whether refusing prterun -n1 is the right side of the trade, or whether a single-rank MPI launch should fall through to the single-process path like a bare environment does.