diff --git a/deepspeed/comm/comm.py b/deepspeed/comm/comm.py index 635f64fef8b8..592a239d754d 100755 --- a/deepspeed/comm/comm.py +++ b/deepspeed/comm/comm.py @@ -843,7 +843,24 @@ def init_distributed(dist_backend: Optional[str] = None, elif in_aws_sm(): patch_aws_sm_env_for_torch_nccl_backend(verbose=verbose) else: - mpi_discovery(distributed_port=distributed_port, verbose=verbose) + try: + mpi_discovery(distributed_port=distributed_port, verbose=verbose) + except ImportError as err: + launcher_world_size = mpi_world_size_from_env() + if launcher_world_size is not None and launcher_world_size > 1: + raise ImportError( + f"A launcher reports a world size of {launcher_world_size} but mpi4py is not " + "installed, so " + "the rank cannot be discovered from it. Install mpi4py, or set RANK, WORLD_SIZE, " + "LOCAL_RANK, MASTER_ADDR and MASTER_PORT yourself.") from err + if launcher_world_size is None and launched_by_mpi(): + raise ImportError( + "A launcher started this process but does not report a world size in the " + "environment - PMIx launched directly, prterun or prun, sets a rank and no size - " + "so whether this is one rank of several cannot be determined without mpi4py. " + "Install mpi4py, or set RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR and MASTER_PORT " + "yourself.") from err + single_process_discovery(distributed_port=distributed_port, verbose=verbose) if cdb is not None and cdb.is_initialized(): if int(os.getenv('RANK', '0')) == 0: @@ -858,6 +875,60 @@ def init_distributed(dist_backend: Optional[str] = None, cdb = TorchBackend(dist_backend, timeout, init_method, rank, world_size) +# World sizes the launchers export. The size rather than the rank, because a rank only says a +# launcher is present: `srun -n1` sets SLURM_PROCID for a single-task step, which is one process +# on one device and wants the fallback below rather than an error about mpi4py. +MPI_WORLD_SIZE_ENV_VARS = ("OMPI_COMM_WORLD_SIZE", "PMI_SIZE", "MV2_COMM_WORLD_SIZE", "SLURM_NTASKS") + +# Ranks the launchers export, used only to tell "no launcher" from "a launcher that reports no +# size". PMIx launched directly - prterun, prun - is the case that needs it: it sets PMIX_RANK +# and PMIX_NAMESPACE and no size at all, so the world cannot be read from the environment. +MPI_RANK_ENV_VARS = ("OMPI_COMM_WORLD_RANK", "PMI_RANK", "PMIX_RANK", "MV2_COMM_WORLD_RANK", "SLURM_PROCID") + + +def mpi_world_size_from_env(): + """The launcher's world size, or None when no launcher reports one.""" + for var in MPI_WORLD_SIZE_ENV_VARS: + try: + return int(os.environ[var]) + except (KeyError, ValueError): + continue + return None + + +def launched_by_mpi(): + """Whether a launcher started this process, whatever world size it reports.""" + return any(var in os.environ for var in MPI_RANK_ENV_VARS) + + +def single_process_discovery(distributed_port=TORCH_DISTRIBUTED_DEFAULT_PORT, verbose=True): + """Fill in the distributed environment for one process on one device. + + Reached when no launcher set the variables and no MPI job is running - `python train.py` on + a single-accelerator machine, which is how DeepSpeed is used on a laptop. `mpi_discovery` + cannot serve that case: it imports mpi4py, so without that package the run ends at + `ModuleNotFoundError: No module named 'mpi4py'` before `deepspeed.initialize` returns. + + `comm/utils.py` already reads a missing launcher as rank 0 of a world of 1; this puts the + same reading in the environment the backend is initialized from. + """ + defaults = { + "RANK": "0", + "WORLD_SIZE": "1", + "LOCAL_RANK": "0", + "MASTER_ADDR": "127.0.0.1", + "MASTER_PORT": str(distributed_port), + } + for name, value in defaults.items(): + os.environ.setdefault(name, value) + + if verbose: + utils.logger.info("No launcher and no MPI job detected; running as a single process with world_rank={}, " + "local_rank={}, world_size={}, master_addr={}, master_port={}".format( + os.environ["RANK"], os.environ["LOCAL_RANK"], os.environ["WORLD_SIZE"], + os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"])) + + def mpi_discovery(distributed_port=TORCH_DISTRIBUTED_DEFAULT_PORT, verbose=True): ''' Discovery MPI environment via mpi4py and map to relevant dist state diff --git a/tests/unit/comm/test_single_process_discovery.py b/tests/unit/comm/test_single_process_discovery.py new file mode 100644 index 000000000000..5064dfb68c2d --- /dev/null +++ b/tests/unit/comm/test_single_process_discovery.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""`deepspeed.initialize` on a machine with no launcher must not require mpi4py. + +`init_distributed` fills in the distributed environment when a launcher did not, and its only +route for that was `mpi_discovery`, which imports mpi4py. Running `python train.py` on a single +accelerator - no launcher, no MPI, no mpi4py - therefore ended at `ModuleNotFoundError: No +module named 'mpi4py'` before `deepspeed.initialize` returned. +""" + +import os + +import pytest + +from deepspeed.comm.comm import (MPI_RANK_ENV_VARS, MPI_WORLD_SIZE_ENV_VARS, launched_by_mpi, mpi_world_size_from_env, + single_process_discovery) + +LAUNCHER_ENV = ("RANK", "WORLD_SIZE", "LOCAL_RANK", "MASTER_ADDR", "MASTER_PORT") + + +@pytest.fixture +def clean_env(monkeypatch): + for name in LAUNCHER_ENV + MPI_WORLD_SIZE_ENV_VARS + MPI_RANK_ENV_VARS + ("PMIX_NAMESPACE", ): + monkeypatch.delenv(name, raising=False) + + +def test_a_bare_environment_reports_no_launcher_and_no_size(clean_env): + assert mpi_world_size_from_env() is None + assert launched_by_mpi() is False + + +@pytest.mark.parametrize("var", MPI_WORLD_SIZE_ENV_VARS) +def test_each_launcher_size_variable_is_read(clean_env, monkeypatch, var): + """OpenMPI, MPICH/Intel MPI, MVAPICH and srun each export a different one.""" + monkeypatch.setenv(var, "4") + + assert mpi_world_size_from_env() == 4 + + +@pytest.mark.parametrize("var", MPI_WORLD_SIZE_ENV_VARS) +def test_a_launcher_reporting_one_task_reports_one(clean_env, monkeypatch, var): + """`srun -n1` is a launcher and a single process at once; it wants the fallback, not an error.""" + monkeypatch.setenv(var, "1") + + assert mpi_world_size_from_env() == 1 + + +@pytest.mark.parametrize("var", MPI_RANK_ENV_VARS) +def test_a_rank_variable_marks_a_launcher_but_gives_no_size(clean_env, monkeypatch, var): + """A rank says a launcher is present, not how big the world is.""" + monkeypatch.setenv(var, "0") + + assert launched_by_mpi() is True + assert mpi_world_size_from_env() is None + + +def test_an_unparseable_size_falls_through_to_the_next_variable(clean_env, monkeypatch): + monkeypatch.setenv("SLURM_NTASKS", "") + + assert mpi_world_size_from_env() is None + + +def test_single_process_discovery_fills_the_environment(clean_env): + single_process_discovery(distributed_port=29501, verbose=False) + + assert os.environ["RANK"] == "0" + assert os.environ["LOCAL_RANK"] == "0" + assert os.environ["WORLD_SIZE"] == "1" + assert os.environ["MASTER_ADDR"] == "127.0.0.1" + assert os.environ["MASTER_PORT"] == "29501" + + +def test_single_process_discovery_leaves_what_the_caller_set(clean_env, monkeypatch): + """A partially set environment is the caller's, not something to overwrite.""" + monkeypatch.setenv("MASTER_PORT", "12345") + monkeypatch.setenv("MASTER_ADDR", "10.0.0.7") + + single_process_discovery(distributed_port=29501, verbose=False) + + assert os.environ["MASTER_PORT"] == "12345" + assert os.environ["MASTER_ADDR"] == "10.0.0.7" + assert os.environ["WORLD_SIZE"] == "1" + + +def test_a_multi_rank_job_without_mpi4py_says_so(clean_env, monkeypatch): + """Falling back to a single process there would silently run one rank of a many-rank job.""" + import deepspeed.comm.comm as comm + + def no_mpi4py(*args, **kwargs): + raise ImportError("No module named 'mpi4py'") + + monkeypatch.setattr(comm, "mpi_discovery", no_mpi4py) + monkeypatch.setenv("OMPI_COMM_WORLD_SIZE", "4") + + with pytest.raises(ImportError, match="mpi4py"): + comm.init_distributed(dist_backend="gloo", auto_mpi_discovery=True, dist_init_required=True) + + assert os.environ.get("WORLD_SIZE") != "1", "the environment must not be filled in for a multi-rank job" + + +def test_a_single_task_slurm_step_without_mpi4py_falls_back(clean_env, monkeypatch): + """`srun -n1 python train.py` with no mpi4py: the case ebarkhordar raised on the PR.""" + import deepspeed.comm.comm as comm + + def no_mpi4py(*args, **kwargs): + raise ImportError("No module named 'mpi4py'") + + monkeypatch.setattr(comm, "mpi_discovery", no_mpi4py) + monkeypatch.setenv("SLURM_PROCID", "0") + monkeypatch.setenv("SLURM_NTASKS", "1") + + reached = {} + real = comm.single_process_discovery + + def wrapped(*args, **kwargs): + reached["yes"] = True + return real(*args, **kwargs) + + monkeypatch.setattr(comm, "single_process_discovery", wrapped) + + comm.init_distributed(dist_backend="gloo", auto_mpi_discovery=True, dist_init_required=True) + + assert reached.get("yes"), "a one-task step took the mpi4py error instead of the fallback" + assert os.environ["WORLD_SIZE"] == "1" + + +def test_a_launcher_that_reports_no_size_is_refused(clean_env, monkeypatch): + """PMIx launched directly sets PMIX_RANK and no size at all. + + `prterun -n4` and `prterun -n1` are indistinguishable from the environment, so falling back + would turn the four-rank case into four separate world-size-1 runs. Refusing costs the + one-rank case an error naming mpi4py, which is the recoverable half of that trade. + """ + import deepspeed.comm.comm as comm + + def no_mpi4py(*args, **kwargs): + raise ImportError("No module named 'mpi4py'") + + monkeypatch.setattr(comm, "mpi_discovery", no_mpi4py) + monkeypatch.setenv("PMIX_RANK", "0") + monkeypatch.setenv("PMIX_NAMESPACE", "prterun-host-1234@1") + + with pytest.raises(ImportError, match="does not report a world size"): + comm.init_distributed(dist_backend="gloo", auto_mpi_discovery=True, dist_init_required=True) + + assert "WORLD_SIZE" not in os.environ + + +def test_a_bare_environment_initializes_end_to_end(clean_env, monkeypatch): + """`python train.py` with nothing set, all the way through init_distributed. + + The helper tests above pass on a version of this that shadows `init_distributed`'s own + `world_size` parameter and hands `None` to the backend, because they never reach the + backend. This one does. + """ + import deepspeed.comm.comm as comm + + def no_mpi4py(*args, **kwargs): + raise ImportError("No module named 'mpi4py'") + + monkeypatch.setattr(comm, "mpi_discovery", no_mpi4py) + + comm.init_distributed(dist_backend="gloo", auto_mpi_discovery=True, dist_init_required=True) + + assert os.environ["WORLD_SIZE"] == "1" + assert os.environ["RANK"] == "0"