DeepSpeed support for device IDs

[SeanNaren]

What does this PR do?

Adds support for device IDs for DeepSpeed!

Fixes #9521
Fixes #9541

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

Merging #9847 (aa816da) into master (f16bfe9) will increase coverage by 0%.
The diff coverage is 100%.

@@          Coverage Diff           @@
##           master   #9847   +/-   ##
======================================
  Coverage      93%     93%           
======================================
  Files         178     178           
  Lines       15647   15648    +1     
======================================
+ Hits        14502   14503    +1     
  Misses       1145    1145           

[tchaton]

LGTM !

[awaelchli]

@SeanNaren does the current implementation support non-consecutive device ids? I tried running your branch with the following:

import os

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

from pytorch_lightning import LightningModule, Trainer


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)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        gpus=[1, 3],
        accelerator='deepspeed',
    )
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
    trainer.test(model, dataloaders=test_data)
if __name__ == "__main__":
    run()

However, I get the error:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:3! (when checking arugment for argument mat1 in method wrapper_addmm)

[SeanNaren]

@SeanNaren does the current implementation support non-consecutive device ids? I tried running your branch with the following:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:3! (when checking arugment for argument mat1 in method wrapper_addmm)

Thanks for checking this out @awaelchli!! Have you updated to the latest DeepSpeed? using the latest DeepSpeed seems to work fine!

[awaelchli]

🤦 yes, my bad