Repository: Lightning-AI/pytorch-lightning
Pull request: Error handling for `accelerator="mps"` and `ddp` strategy pairing
## What does this PR do?
Fixes #16148 

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] 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)
- [x] 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 minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/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 🙃


Problem statement:
Error handling for `accelerator="mps"` and `strategy="ddp"`
### Bug description

The MPS backend does not support `torch.distributed`. We should fail early in the AcceleratorConnector and produce a user friendly error message.

### How to reproduce the bug

```python
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 configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    model = BoringModel()
    trainer = Trainer(
        accelerator="mps",
        devices=1,
        strategy="ddp"
    )
    trainer.fit(model, train_dataloaders=train_data)


if __name__ == "__main__":
    run()
```


### Error messages and logs

```
Traceback (most recent call last):
  File "/Users/adrian/repositories/lightning/examples/pl_bug_report/bug_report_model.py", line 49, in <module>
    run()
  File "/Users/adrian/repositories/lightning/examples/pl_bug_report/bug_report_model.py", line 45, in run
    trainer.fit(model, train_dataloaders=train_data)
  File "/Users/adrian/repositories/lightning/src/pytorch_lightning/trainer/trainer.py", line 598, in fit
    call._call_and_handle_interrupt(
  File "/Users/adrian/repositories/lightning/src/pytorch_lightning/trainer/call.py", line 36, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/Users/adrian/repositories/lightning/src/pytorch_lightning/strategies/launchers/subprocess_script.py", line 88, in launch
    return function(*args, **kwargs)
  File "/Users/adrian/repositories/lightning/src/pytorch_lightning/trainer/trainer.py", line 640, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "/Users/adrian/repositories/lightning/src/pytorch_lightning/trainer/trainer.py", line 1074, in _run
    self.strategy.setup(self)
  File "/Users/adrian/repositories/lightning/src/pytorch_lightning/strategies/ddp.py", line 159, in setup
    self._share_information_to_prevent_deadlock()
  File "/Users/adrian/repositories/lightning/src/pytorch_lightning/strategies/ddp.py", line 405, in _share_information_to_prevent_deadlock
    self._share_pids()
  File "/Users/adrian/repositories/lightning/src/pytorch_lightning/strategies/ddp.py", line 423, in _share_pids
    pids = self.all_gather(torch.tensor(os.getpid(), device=self.root_device))
  File "/Users/adrian/repositories/lightning/src/pytorch_lightning/strategies/parallel.py", line 90, in all_gather
    return _all_gather_ddp_if_available(tensor, group=group, sync_grads=sync_grads)
  File "/Users/adrian/repositories/lightning/src/lightning_lite/utilities/distributed.py", line 202, in _all_gather_ddp_if_available
    return _AllGather.apply(tensor, group)
  File "/Users/adrian/repositories/lightning/src/lightning_lite/utilities/distributed.py", line 170, in forward
    torch.distributed.all_gather(gathered_tensor, tensor, group=group)
  File "/Users/adrian/miniconda3/envs/lightning/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 2277, in all_gather
    work = group.allgather([tensor_list], [tensor])
RuntimeError: ProcessGroupGloo::allgather: unsupported device type mps
```

This applies for the entire DDPStrategy family (ddp, ddp_fork, ddp_spawn, etc.)

### Environment

<details>
  <summary>Current environment</summary>

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): Trainer, LightningLite
#- PyTorch Lightning Version (e.g., 1.5.0): 1.9.0dev
#- Lightning App Version (e.g., 0.5.2): x
#- PyTorch Version (e.g., 1.10): 1.13
#- Python version (e.g., 3.9): 3.10
#- OS (e.g., Linux): MacOS
#- CUDA/cuDNN version: x
#- GPU models and configuration: x
#- How you installed Lightning(`conda`, `pip`, source): pip 
#- Running environment of LightningApp (e.g. local, cloud): x
```

</details>


### More info

_No response_

cc @borda @justusschock

Hints from issue discussion:
I'd like to take this up!
Perfect! Ping me here or on the PR if you have any questions

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/src/pytorch_lightning/trainer/connectors/accelerator_connector.py b/src/pytorch_lightning/trainer/connectors/accelerator_connector.py
index 03f7d53732a6e..264090a84d7bd 100644
--- a/src/pytorch_lightning/trainer/connectors/accelerator_connector.py
+++ b/src/pytorch_lightning/trainer/connectors/accelerator_connector.py
@@ -283,6 +283,24 @@ def _check_config_and_set_final_flags(
                 f" Available names are: {', '.join(self._accelerator_types)}."
             )
 
+        # mps accelerator incompatible with ddp-family
+        ddp_classes = (
+            DDPStrategy,
+            DDPSpawnShardedStrategy,
+            DDPShardedStrategy,
+            DDPFullyShardedNativeStrategy,
+            DDPFullyShardedStrategy,
+            DDPSpawnStrategy,
+        )
+        is_ddp_str = isinstance(strategy, str) and "ddp" in strategy
+        if (
+            accelerator is not None
+            and strategy is not None
+            and (isinstance(accelerator, str) and accelerator == "mps")
+            and (is_ddp_str or isinstance(strategy, ddp_classes))
+        ):
+            raise MisconfigurationException("With `accelerator=mps`, strategies from DDP Family are not compatible")
+
         self._accelerator_flag = accelerator
 
         if precision is not None:
diff --git a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py
index b33400d2aa227..f0a5658a4c479 100644
--- a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py
+++ b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py
@@ -369,6 +369,33 @@ def test_exception_invalid_strategy():
         Trainer(strategy="tpu_spawn")
 
 
+@pytest.mark.parametrize(
+    ["strategy", "strategy_class"],
+    (
+        ("ddp_spawn", DDPSpawnStrategy),
+        ("ddp_spawn_find_unused_parameters_false", DDPSpawnStrategy),
+        ("ddp", DDPStrategy),
+        ("ddp_find_unused_parameters_false", DDPStrategy),
+        ("dp", DataParallelStrategy),
+        ("ddp_sharded", DDPShardedStrategy),
+        ("ddp_sharded_spawn", DDPSpawnShardedStrategy),
+    ),
+)
+@RunIf(mps=True)
+def test_exception_invalid_ddp_strategy_with_mps(strategy, strategy_class):
+    # check with str
+    with pytest.raises(
+        MisconfigurationException, match=r"With `accelerator=mps`, strategies from DDP Family are not compatible"
+    ):
+        Trainer(accelerator="mps", strategy=strategy)
+
+    # check with instance
+    with pytest.raises(
+        MisconfigurationException, match=r"With `accelerator=mps`, strategies from DDP Family are not compatible"
+    ):
+        Trainer(accelerator="mps", strategy=strategy_class())
+
+
 @pytest.mark.parametrize(
     ["strategy", "strategy_class"],
     [


Code review comments:
diff_hunk: @@ -283,6 +283,24 @@ def _check_config_and_set_final_flags(
                 f" Available names are: {', '.join(self._accelerator_types)}."
             )
 
+        # mps accelerator incompatible with ddp-family
+        ddp_classes = (
+            DDPStrategy,
+            DDPSpawnShardedStrategy,
+            DDPShardedStrategy,
+            DDPFullyShardedNativeStrategy,
+            DDPFullyShardedStrategy,
+            DDPSpawnStrategy,
+        )
+        is_ddp_str = isinstance(strategy, str) and "ddp" in strategy
+        if (
+            accelerator is not None
+            and strategy is not None
+            and (isinstance(accelerator, str) and accelerator == "mps")
+            and (is_ddp_str or isinstance(strategy, ddp_classes))
path: src/pytorch_lightning/trainer/connectors/accelerator_connector.py
text: @user1:
```suggestion
            and (is_ddp_str or isinstance(strategy, ParallelStrategy))
```
this should cover all kinds of parallelism

diff_hunk: @@ -283,6 +283,24 @@ def _check_config_and_set_final_flags(
                 f" Available names are: {', '.join(self._accelerator_types)}."
             )
 
+        # mps accelerator incompatible with ddp-family
+        ddp_classes = (
+            DDPStrategy,
+            DDPSpawnShardedStrategy,
+            DDPShardedStrategy,
+            DDPFullyShardedNativeStrategy,
+            DDPFullyShardedStrategy,
+            DDPSpawnStrategy,
+        )
path: src/pytorch_lightning/trainer/connectors/accelerator_connector.py
text: @user1:
```suggestion
```