[DeepSpeed] `MPI world size 1 does not match torch world size 2` when launched with >1 GPU.
🐛 Bug
When trying to train a model with PyTorch Lightning and Deepspeed, the process fails with an assertion error:
AssertionError: MPI world size 1 does not match torch world size 2.
The number 2 in this case is replaced with the number of GPUs specified by the PyTorch Lightning Trainer.
Please reproduce using the BoringModel
Because Colab does not expose multiple GPUs, I cannot post a Colab link that reproduces this issue. The script to reproduce is simple:
from pytorch_lightning import Trainer
from pytorch_lightning.plugins import DeepSpeedPlugin, DeepSpeedPrecisionPlugin
from tests.helpers.boring_model import BoringModel

model = BoringModel()
trainer = Trainer(
        plugins=[DeepSpeedPlugin()],
        default_root_dir="/tmp",
        gpus=2,
        fast_dev_run=True,
        precision=16,
)
trainer.fit(model)
trainer.test(model)

Expected behavior
I expect the world size to be set to the number of GPUs specified so that DeepSpeed will run with my model on multiple GPUs.
Environment

PyTorch Version (e.g., 1.0): 1.9.0.dev20210305+cu111
PyTorch Lightning Version: 1.2.3
OS (e.g., Linux): Ubuntu 20.04
How you installed PyTorch (conda, pip, source): pip
Python version: 3.8
CUDA/cuDNN version: 11.1
GPU models and configuration: 4x A100-SXM4-40GB
Any other relevant information: Deepspeed version: 0.3.10

Additional context
I searched Google and this repository extensively to find some way to adjust the MPI world size, but this was to no avail.