Repository: Lightning-AI/pytorch-lightning
Pull request: Only import PostLocalSGD related modules when it's needed
## What does this PR do?

<!--
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)
-->

Partially Fixes #10348 

### Does your PR introduce any breaking changes? If yes, please list them.

DistributedOptimizer, PostLocalSGDOptimizer, ZeroRedundancyOptimizer and torch.distributed.algorithms modules only needed when `isinstance(self._ddp_comm_state, post_localSGD.PostLocalSGDState)`, move the import into _reinit_optimizers_with_post_localSGD() where these modules are used.

## 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/PyTorchLightning/pytorch-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?
- [x] 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)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [ ] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- 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 [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] 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:
Module 'torch.distributed' has no attribute 'ProcessGroup' when importing PyTorch Lightning
## 🐛 Bug

When importing PyTorch Lightning, it throws an `AttributeError: module 'torch.distributed' has no attribute 'ProcessGroup'`. I guess it comes from the fact that I am on macOS (M1) and PyTorch does not provide `torch.distributed` with its pre-built package. Indeed, `torch.distributed.is_available()` is `False`. 

### To Reproduce

```
import pytorch_lightning
```

### Environment

- PyTorch Lightning Version: 1.5.0
- PyTorch Version: 1.10
- Python version: 3.9
- OS: macOS
- How you installed PyTorch: `conda`
- Any other relevant information:

Hints from issue discussion:
Hello @Riccorl 
Thanks for reporting. Can you please show us the full error so we can check in which module it occurs?

I'm not sure why you get `torch.distributed.is_available() = False` on MacOS, it should be True. It is for me.
> Hello @Riccorl Thanks for reporting. Can you please show us the full error so we can check in which module it occurs?

This is the stack trace

```
Traceback (most recent call last):
  File "/Users/ric/Documents/PhD/Projects/invero-xl/invero_xl/train.py", line 7, in <module>
    import pytorch_lightning as pl
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/__init__.py", line 20, in <module>
    from pytorch_lightning.callbacks import Callback  # noqa: E402
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/callbacks/__init__.py", line 26, in <module>
    from pytorch_lightning.callbacks.pruning import ModelPruning
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/callbacks/pruning.py", line 31, in <module>
    from pytorch_lightning.core.lightning import LightningModule
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/core/__init__.py", line 16, in <module>
    from pytorch_lightning.core.lightning import LightningModule
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/core/lightning.py", line 39, in <module>
    from pytorch_lightning.trainer.connectors.logger_connector.fx_validator import _FxValidator
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/trainer/__init__.py", line 16, in <module>
    from pytorch_lightning.trainer.trainer import Trainer
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 30, in <module>
    from pytorch_lightning.accelerators import Accelerator, IPUAccelerator
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/accelerators/__init__.py", line 13, in <module>
    from pytorch_lightning.accelerators.accelerator import Accelerator  # noqa: F401
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/accelerators/accelerator.py", line 26, in <module>
    from pytorch_lightning.plugins.precision import ApexMixedPrecisionPlugin, NativeMixedPrecisionPlugin, PrecisionPlugin
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/plugins/__init__.py", line 8, in <module>
    from pytorch_lightning.plugins.plugins_registry import (  # noqa: F401
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/plugins/plugins_registry.py", line 20, in <module>
    from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/plugins/training_type/__init__.py", line 1, in <module>
    from pytorch_lightning.plugins.training_type.ddp import DDPPlugin  # noqa: F401
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 68, in <module>
    from torch.distributed.optim import DistributedOptimizer
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/torch/distributed/optim/__init__.py", line 37, in <module>
    from .post_localSGD_optimizer import PostLocalSGDOptimizer
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/torch/distributed/optim/post_localSGD_optimizer.py", line 2, in <module>
    import torch.distributed.algorithms.model_averaging.averagers as averagers
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/torch/distributed/algorithms/model_averaging/averagers.py", line 5, in <module>
    import torch.distributed.algorithms.model_averaging.utils as utils
  File "/Users/ric/mambaforge/envs/srl/lib/python3.9/site-packages/torch/distributed/algorithms/model_averaging/utils.py", line 10, in <module>
    params: Iterator[torch.nn.Parameter], process_group: dist.ProcessGroup
AttributeError: module 'torch.distributed' has no attribute 'ProcessGroup'
```


> I'm not sure why you get `torch.distributed.is_available() = False` on MacOS, it should be True. It is for me.

I installed PyTorch like this:

```
conda install pytorch -c pytorch
```

But I guess that the problem is the ARM build (I'm on an M1 cpu).
We can fix this easily as the error comes from a typing annotation, but we'll also have to add a M1 CI job when it becomes available.
Current init_dist_connection() will do nothing if torch.distributed.is_avalible = False. To wrap model with `DistributedDataParallel()`, something like `torch.distributed.init_process_group(
     backend='nccl', world_size=N, init_method='...'
      )`
Is required right?
Should we raise exceptions at set_distributed() in ddp if torch.distributed.is_avalible = False?

Commit message: Only import PostLocalSGD related modules when it's needed

Patch to review:
diff --git a/pytorch_lightning/plugins/training_type/ddp.py b/pytorch_lightning/plugins/training_type/ddp.py
index ea4820f61ec7c..229cc3e3ef1f3 100644
--- a/pytorch_lightning/plugins/training_type/ddp.py
+++ b/pytorch_lightning/plugins/training_type/ddp.py
@@ -63,11 +63,6 @@
 from pytorch_lightning.utilities.seed import reset_seed
 from pytorch_lightning.utilities.types import STEP_OUTPUT
 
-if _TORCH_GREATER_EQUAL_1_10:
-    if not _IS_WINDOWS:
-        from torch.distributed.optim import DistributedOptimizer
-    from torch.distributed.optim import PostLocalSGDOptimizer, ZeroRedundancyOptimizer
-
 if _FAIRSCALE_AVAILABLE:
     from fairscale.optim import OSS
 if _HYDRA_AVAILABLE:
@@ -75,9 +70,7 @@
     from hydra.utils import get_original_cwd, to_absolute_path
 if _TORCH_GREATER_EQUAL_1_8:
     from pytorch_lightning.utilities.distributed import register_ddp_comm_hook
-if _TORCH_GREATER_EQUAL_1_10:
-    import torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook as post_localSGD
-    import torch.distributed.algorithms.model_averaging.averagers as averagers
+
 
 log = logging.getLogger(__name__)
 
@@ -323,7 +316,8 @@ def _register_ddp_hooks(self) -> None:
                 ddp_comm_hook=self._ddp_comm_hook,
                 ddp_comm_wrapper=self._ddp_comm_wrapper,
             )
-
+            if _TORCH_GREATER_EQUAL_1_10:
+                import torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook as post_localSGD
             if (
                 _TORCH_GREATER_EQUAL_1_10
                 and isinstance(self._ddp_comm_state, post_localSGD.PostLocalSGDState)
@@ -337,6 +331,12 @@ def _reinit_optimizers_with_post_localSGD(self, warmup_steps: int):
             raise ValueError(
                 "Post-localSGD algorithm is used, but model averaging period is not provided to DDP plugin."
             )
+        if _TORCH_GREATER_EQUAL_1_10:
+            if not _IS_WINDOWS:
+                from torch.distributed.optim import DistributedOptimizer
+            import torch.distributed.algorithms.model_averaging.averagers as averagers
+            from torch.distributed.optim import PostLocalSGDOptimizer, ZeroRedundancyOptimizer
+
         averager = averagers.PeriodicModelAverager(period=self._model_averaging_period, warmup_steps=warmup_steps)
         for x, optimizer in enumerate(optimizers):
             if isinstance(optimizer, LightningOptimizer):


Code review comments:
diff_hunk: @@ -323,7 +316,8 @@ def _register_ddp_hooks(self) -> None:
                 ddp_comm_hook=self._ddp_comm_hook,
                 ddp_comm_wrapper=self._ddp_comm_wrapper,
             )
-
+            if _TORCH_GREATER_EQUAL_1_10:
+                import torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook as post_localSGD
             if (
                 _TORCH_GREATER_EQUAL_1_10
path: pytorch_lightning/plugins/training_type/ddp.py
text: @user1:
Merge this with the if above you added