DeepSpeed requires the LOCAL_RANK environment variable

[nickKyr]

Hello,

I am trying to use DeepSpeed in the pl.Trainer of PyTorch Lightning and I get the following error:
AssertionError: DeepSpeed requires the LOCAL_RANK environment variable, it is set by the deepspeed launcher, deepspeed.init_distributed, or the torch.distributed launcher. If using a different launcher please ensure LOCAL_RANK is set prior to initializing deepspeed.

The trainer is:

trainer = pl.Trainer(
           gpus=2, accelerator='ddp',
           callbacks=[checkpoint, callback1],
           max_epochs=epochs,
           logger=logger1,
           plugins=[DeepSpeedPlugin()], precision=16
           )

and the corresponding import at the beginning of the file:
from pytorch_lightning.plugins import DeepSpeedPlugin
Is the above problem related to Distributed Initialization and therefore, deepspeed.init_distributed(init_method="Torch distributed") must be added prior to the Trainer?

Any other suggestion is more than welcome.

Thanks,
Nikos

[tjruwase]

@nickKyr, is this a general problem with DeepSpeed+Lightning or specific to your scenario?
@SeanNaren for awareness.

[nickKyr]

@tjruwase that worked! Thank you both @tjruwase and @SeanNaren for the combined effort and your ideas!
I had read somewhere that DeepSpeed can run in 32 mode, but I was getting an error that precision must be set to 16 (with older Lightning version). It wasn't until the most recent PyTorch Lightning v1.4 that this was fixed based on the link you sent and I hadn't tried this with v1.4...

To sum up, Lightning v1.4 is required to integrate DeepSpeed. The only changes needed are the addition of the corresponding arguments in the trainer call. On a similar note, I noticed that the order of setup and train_dataloader methods is reversed when integrating DeepSpeed in Lightning.

Cheers,
Nikos

[tjruwase]

@nickKyr, thanks for sharing this update. I am so glad to hear that Lightning+DeepSpeed is now working for you. That is awesome.

It does seem like some work is needed to make mixed-precision training a smoother process. But perhaps, mixed-precision is not an urgent need for you at this time. But whenever you do plan to try mixed-precision, do let us know and we can solve the issues together :).

[nickKyr]

Hi, mixed-precision seems to work in Lightning v1.4 as well! I just added precision=16 in the trainer and modified the inputs and targets, which are actually float, in training_step and validation_step:

    def training_step(self, batch, batch_idx):
        inputs, targets = batch
        inputs, targets = inputs.half(), targets.half()
        outputs = self.forward(inputs)
        ...

[tjruwase]

Thanks for sharing that update. That is absolutely awesome to hear. Big kudos to @SeanNaren and co for enabling a seamless mixed-precision training experience in Lightning. That is no trivial achievement.

@nickKyr, do you think we can now close this question as answered?

[nickKyr]

Yes, the issue has been resolved. Thanks for the new capabilities that you have added!

[tjruwase]

LOCAL_RANK (and mixed-precision) issue resolved by upgrading to Lightning v1.4.

[shoang22]

Hello,

I'm running into a similar issue. I'm trying to run DeepSpeed (0.3.12) + Lightning (1.2.3) in the following toy example train_boring.py using a SLURM cluster:

import os
import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.plugins import DeepSpeedPlugin


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    test_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    ngpus = int(os.environ["SLURM_NTASKS_PER_NODE"])
    nnodes = int(os.environ["SLURM_NNODES"])

    ds_config = {
        "zero_optimization": {
            "stage": 2,
            "cpu_offload": False,
            "contiguous_gradients": True,
            "overlap_comm": True,
            "allgather_bucket_size": 5e8,
            "reduce_bucket_size": 5e8
        },
        "fp16": {
            "enabled": True,
            "loss_scale": 0.0,
            "initial_scale_power": 32,
            "loss_scale_window": 1000,
            "hysteresis": 2,
            "min_loss_scale": 1000
        }
    }

    plugins = [DeepSpeedPlugin(ds_config)]

    print("N GPUS: {}".format(ngpus))
    print("N NODES: {}".format(nnodes))

    model = BoringModel()
    trainer = Trainer(
        accelerator="ddp",
        gpus=ngpus,
        num_nodes=nnodes,
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        limit_test_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        plugins=plugins,
        precision=16
    )
    trainer.fit(model, train_data)
    trainer.test(model, test_data)


if __name__ == "__main__":
    run()

The following script is used to launch the SLURM job:

#!/bin/bash
#SBATCH --job-name=boring
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH -p gpu
#SBATCH --gres=gpu:4
#SBATCH --mem=64gb
#SBATCH --time=0:10:00
#SBATCH --output=slurm_out/output.%A_%a.txt

module load miniconda3
module load cuda/11.8.0
nvidia-smi

export MASTER_PORT=$(expr 10000 + $(echo -n $SLURM_JOBID | tail -c 4))
master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_ADDR=$master_addr
echo "MASTER_ADDR="$MASTER_ADDR    

srun python train_boring.py

executing this results in the error below:

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla V100-SXM2-32GB           On  | 00000000:1A:00.0 Off |                    0 |
| N/A   30C    P0              43W / 300W |      0MiB / 32768MiB |      0%   E. Process |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   1  Tesla V100-SXM2-32GB           On  | 00000000:1C:00.0 Off |                    0 |
| N/A   30C    P0              43W / 300W |      0MiB / 32768MiB |      0%   E. Process |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   2  Tesla V100-SXM2-32GB           On  | 00000000:1D:00.0 Off |                    0 |
| N/A   29C    P0              41W / 300W |      0MiB / 32768MiB |      0%   E. Process |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   3  Tesla V100-SXM2-32GB           On  | 00000000:1E:00.0 Off |                    0 |
| N/A   32C    P0              43W / 300W |      0MiB / 32768MiB |      0%   E. Process |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|  No running processes found                                                           |
+---------------------------------------------------------------------------------------+
MASTER_ADDR=mg090
GpuFreq=control_disabled
INFO:    underlay of /etc/localtime required more than 50 (80) bind mounts
INFO:    underlay of /usr/bin/nvidia-smi required more than 50 (251) bind mounts
INFO:    underlay of /etc/localtime required more than 50 (80) bind mounts
INFO:    underlay of /etc/localtime required more than 50 (80) bind mounts
INFO:    underlay of /etc/localtime required more than 50 (80) bind mounts
INFO:    underlay of /usr/bin/nvidia-smi required more than 50 (251) bind mounts
INFO:    underlay of /usr/bin/nvidia-smi required more than 50 (251) bind mounts
INFO:    underlay of /usr/bin/nvidia-smi required more than 50 (251) bind mounts
GPU available: True, used: True
TPU available: None, using: 0 TPU cores
Multi-processing is handled by Slurm.
GPU available: True, used: True
TPU available: None, using: 0 TPU cores
Multi-processing is handled by Slurm.
GPU available: True, used: True
TPU available: None, using: 0 TPU cores
Multi-processing is handled by Slurm.
GPU available: True, used: True
TPU available: None, using: 0 TPU cores
Multi-processing is handled by Slurm.
/opt/conda/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:50: UserWarning: you defined a validation_step but have no val_dataloader. Skipping validation loop
  warnings.warn(*args, **kwargs)
initializing ddp: GLOBAL_RANK: 3, MEMBER: 4/4
/opt/conda/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:50: UserWarning: you defined a validation_step but have no val_dataloader. Skipping validation loop
  warnings.warn(*args, **kwargs)
initializing ddp: GLOBAL_RANK: 1, MEMBER: 2/4
/opt/conda/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:50: UserWarning: you defined a validation_step but have no val_dataloader. Skipping validation loop
  warnings.warn(*args, **kwargs)
initializing ddp: GLOBAL_RANK: 2, MEMBER: 3/4
/opt/conda/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:50: UserWarning: you defined a validation_step but have no val_dataloader. Skipping validation loop
  warnings.warn(*args, **kwargs)
initializing ddp: GLOBAL_RANK: 0, MEMBER: 1/4
You have not specified an optimizer or scheduler within the DeepSpeed config.Using `configure_optimizers` to define optimizer and scheduler.
INFO:lightning:You have not specified an optimizer or scheduler within the DeepSpeed config.Using `configure_optimizers` to define optimizer and scheduler.
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/ui/abv/hoangsx/Chemformer/temp/train_boring.py", line 100, in <module>
    run()
  File "/ui/abv/hoangsx/Chemformer/temp/train_boring.py", line 95, in run
    trainer.fit(model, train_data)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 511, in fit
    self.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 540, in pre_dispatch
    self.accelerator.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 84, in pre_dispatch
    self.training_type_plugin.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 174, in pre_dispatch
    self.init_deepspeed()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 193, in init_deepspeed
    self._initialize_deepspeed_train(model)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 218, in _initialize_deepspeed_train
    model, optimizer, _, lr_scheduler = deepspeed.initialize(
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/__init__.py", line 112, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 145, in __init__
    self._configure_with_arguments(args, mpu)
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 499, in _configure_with_arguments
    args.local_rank = int(os.environ['LOCAL_RANK'])
  File "/opt/conda/lib/python3.8/os.py", line 675, in __getitem__
    raise KeyError(key) from None
KeyError: 'LOCAL_RANK'
N GPUS: 4
N NODES: 1
You have not specified an optimizer or scheduler within the DeepSpeed config.Using `configure_optimizers` to define optimizer and scheduler.
INFO:lightning:You have not specified an optimizer or scheduler within the DeepSpeed config.Using `configure_optimizers` to define optimizer and scheduler.
You have not specified an optimizer or scheduler within the DeepSpeed config.Using `configure_optimizers` to define optimizer and scheduler.
INFO:lightning:You have not specified an optimizer or scheduler within the DeepSpeed config.Using `configure_optimizers` to define optimizer and scheduler.
You have not specified an optimizer or scheduler within the DeepSpeed config.Using `configure_optimizers` to define optimizer and scheduler.
INFO:lightning:You have not specified an optimizer or scheduler within the DeepSpeed config.Using `configure_optimizers` to define optimizer and scheduler.
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/ui/abv/hoangsx/Chemformer/temp/train_boring.py", line 100, in <module>
    run()
  File "/ui/abv/hoangsx/Chemformer/temp/train_boring.py", line 95, in run
    trainer.fit(model, train_data)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 511, in fit
    self.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 540, in pre_dispatch
    self.accelerator.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 84, in pre_dispatch
    self.training_type_plugin.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 174, in pre_dispatch
    self.init_deepspeed()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 193, in init_deepspeed
    self._initialize_deepspeed_train(model)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 218, in _initialize_deepspeed_train
    model, optimizer, _, lr_scheduler = deepspeed.initialize(
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/__init__.py", line 112, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 145, in __init__
    self._configure_with_arguments(args, mpu)
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 499, in _configure_with_arguments
    args.local_rank = int(os.environ['LOCAL_RANK'])
  File "/opt/conda/lib/python3.8/os.py", line 675, in __getitem__
    raise KeyError(key) from None
KeyError: 'LOCAL_RANK'
N GPUS: 4
N NODES: 1
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/ui/abv/hoangsx/Chemformer/temp/train_boring.py", line 100, in <module>
    run()
  File "/ui/abv/hoangsx/Chemformer/temp/train_boring.py", line 95, in run
    trainer.fit(model, train_data)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 511, in fit
    self.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 540, in pre_dispatch
    self.accelerator.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 84, in pre_dispatch
    self.training_type_plugin.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 174, in pre_dispatch
    self.init_deepspeed()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 193, in init_deepspeed
    self._initialize_deepspeed_train(model)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 218, in _initialize_deepspeed_train
    model, optimizer, _, lr_scheduler = deepspeed.initialize(
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/__init__.py", line 112, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 145, in __init__
    self._configure_with_arguments(args, mpu)
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 499, in _configure_with_arguments
    args.local_rank = int(os.environ['LOCAL_RANK'])
  File "/opt/conda/lib/python3.8/os.py", line 675, in __getitem__
    raise KeyError(key) from None
KeyError: 'LOCAL_RANK'
N GPUS: 4
N NODES: 1
N GPUS: 4
N NODES: 1
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/ui/abv/hoangsx/Chemformer/temp/train_boring.py", line 100, in <module>
    run()
  File "/ui/abv/hoangsx/Chemformer/temp/train_boring.py", line 95, in run
    trainer.fit(model, train_data)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 511, in fit
    self.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 540, in pre_dispatch
    self.accelerator.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 84, in pre_dispatch
    self.training_type_plugin.pre_dispatch()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 174, in pre_dispatch
    self.init_deepspeed()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 193, in init_deepspeed
    self._initialize_deepspeed_train(model)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py", line 218, in _initialize_deepspeed_train
    model, optimizer, _, lr_scheduler = deepspeed.initialize(
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/__init__.py", line 112, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 145, in __init__
    self._configure_with_arguments(args, mpu)
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 499, in _configure_with_arguments
    args.local_rank = int(os.environ['LOCAL_RANK'])
  File "/opt/conda/lib/python3.8/os.py", line 675, in __getitem__
    raise KeyError(key) from None
KeyError: 'LOCAL_RANK'

Attempting to upgrade to Pytorch Lightning 1.4.0 resulted in the same error. Is there something I'm missing?