MPI environment variables are not set

[fabiogeraci]

System Info
HPC ubuntu 22.04 2nodesx8H100

LSF as scheduler

[tool.poetry.dependencies]
python = "^3.10"

importlib-metadata = { version = "~=1.0", python = "<3.8" }
tensorboard = "^2.16.2"
sge-data-package = {version = "", source = "sgedata"}
torch = "2.2.1"
torchvision = "0.17.1"
torchaudio = "2.2.1"
transformers = "4.42.0"
datasets = "2.18."
accelerate = "0.28.0"
deepspeed = "0.13.4"
safetensors = "0.4.2"
mpi4py = "^4.0.0"

module load cuda-12.1.1
module load ISG/experimental/fg12/openmpi/5.0.4-cuda12.1-lsf
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

deepspeed \
    --hostfile=${HOSTFILE_PATH} \
    --launcher=OPENMPI \
    --launcher_args="-bind-to none -map-by slot --mca pml ob1 --oversubscribe --display-allocation --display-map" \
    --master_addr=${MASTER_ADDR} \
    --master_port=${_M_PORT} \
    --no_ssh_check \
    src/dna_mlm/runner.py

def setup_env_ranks() -> tp.Tuple[int, int, int]:

    # Map MPI environment variables to those expected by DeepSpeed/PyTorch
    if 'OMPI_COMM_WORLD_LOCAL_RANK' in os.environ:
        os.environ['LOCAL_RANK'] = os.environ['OMPI_COMM_WORLD_LOCAL_RANK']
        os.environ['RANK'] = os.environ['OMPI_COMM_WORLD_RANK']
        os.environ['WORLD_SIZE'] = os.environ['OMPI_COMM_WORLD_SIZE']
    else:
        raise EnvironmentError(
            "MPI environment variables are not set. "
            "Ensure you are running the script with an MPI-compatible launcher."
        )
 
 setup_env_ranks()

the function should set the env vars but instaed it raises the error

[fabiogeraci]

I found the error

deepspeed \
    --hostfile ${HOSTFILE_PATH} \
    --launcher "OPENMPI" \ #openmpi should have been between ""
    --launcher_args "-bind-to none -map-by slot --allow-run-as-root --mca pml ob1 --oversubscribe --display-allocation --display-map" \
    --master_addr ${MASTER_ADDR} \
    --master_port ${_M_PORT} \
    --no_ssh_check \
    src/runner.py

[fabiogeraci]

the real question is why I need to setup

    if 'OMPI_COMM_WORLD_LOCAL_RANK' in os.environ:
        os.environ['LOCAL_RANK'] = os.environ['OMPI_COMM_WORLD_LOCAL_RANK']
        os.environ['RANK'] = os.environ['OMPI_COMM_WORLD_RANK']
        os.environ['WORLD_SIZE'] = os.environ['OMPI_COMM_WORLD_SIZE']
    else:
        raise EnvironmentError(
            "MPI environment variables are not set. "
            "Ensure you are running the script with an MPI-compatible launcher."
        )

[loadams]

Hi @fabiogeraci - can you be more specific on the question here? Are you asking why you had to set OMPI_COMM_WORLD_LOCAL_RANK?

[loadams]

Hi @fabiogeraci - following up on this?

[fabiogeraci]

I was in AL, i will add more details tomorrow

[fabiogeraci]

@loadams I would think that if i pass --launcher "OPENMPI" deepspeed whould check and set OMPI_COMM_WORLD_LOCAL_RANK?

[loadams]

@fabiogeraci - that makes sense to me. Would you prefer to make the PR with these suggested changes and we can review/get feedback from others using the OpenMPI launcher there?

[fabiogeraci]

@loadams
Apologies, i have been unwell

according to the documentation, I should open an issue first

A description of the proposed feature.
A motivation of why it will be useful to DeepSpeed users.
A rough design of how you implement the feature inside DeepSpeed.
(Important) Results or planned experiments to demonstrate the effectiveness and correctness of the feature.

but i think this is already covered, could you provide guidance, please?

[loadams]

Hi @fabiogeraci - apologies, the documentation can be a bit confusing. Since we've already discussed it here, and the scope of the change isn't huge (the documents are more for major refactor/feature adds), I'd say you're fine to just go ahead and create the PR and we can move any discussion on specific implementation and any questions/reviews there.

[fabiogeraci]

@loadams I have prepared the PR, with tests (which I am not sure that they are needed, but just following guide lines). I have an issue the following tests

• test_openmpi_runner
• test_btl_nic_openmpi_runner
• test_btl_nic_two_dashes_openmpi_runner

they all fail similarly
FAILED tests/unit/launcher/test_multinode_runner.py::test_openmpi_runner - OSError: MPI environment variables are not set. Ensure you are running the script with an MPI-compatible launcher.

since I added to class OpenMPIRunner(MultiNodeRunner)

def _setup_mpi_environment(self):
        """Sets up MPI-related environment variables or raises an error if they're missing."""

        required_vars = [
            'OMPI_COMM_WORLD_LOCAL_RANK',
            'OMPI_COMM_WORLD_RANK',
            'OMPI_COMM_WORLD_SIZE'
        ]

        # Check if all these are present
        if not all(var in os.environ for var in required_vars):
            raise EnvironmentError(
                "MPI environment variables are not set. "
                "Ensure you are running the script with an MPI-compatible launcher."
            )

        # Now safe to read all
        os.environ['LOCAL_RANK'] = os.environ['OMPI_COMM_WORLD_LOCAL_RANK']
        os.environ['RANK'] = os.environ['OMPI_COMM_WORLD_RANK']
        os.environ['WORLD_SIZE'] = os.environ['OMPI_COMM_WORLD_SIZE']

I can fix those tests, for instance

def test_openmpi_runner(runner_info, monkeypatch):
    env, resource_pool, world_info, args = runner_info

    # Provide the 3 required MPI variables:
    monkeypatch.setenv('OMPI_COMM_WORLD_LOCAL_RANK', '0')
    monkeypatch.setenv('OMPI_COMM_WORLD_RANK', '0')
    monkeypatch.setenv('OMPI_COMM_WORLD_SIZE', '1')

    runner = mnrunner.OpenMPIRunner(args, world_info, resource_pool)
    cmd = runner.get_cmd(env, resource_pool)
    assert cmd[0] == 'mpirun'
    assert 'eth0' in cmd