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

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

Author: pre-commit-ci[bot]

pytorch_lightning/strategies/__init__.py

@@ -19,8 +19,8 @@
 from pytorch_lightning.strategies.ddp_spawn import DDPSpawnStrategy  # noqa: F401
 from pytorch_lightning.strategies.deepspeed import DeepSpeedStrategy  # noqa: F401
 from pytorch_lightning.strategies.dp import DataParallelStrategy  # noqa: F401
-from pytorch_lightning.strategies.fully_sharded_native import DDPFullyShardedNativeStrategy  # noqa: F401
 from pytorch_lightning.strategies.fully_sharded import DDPFullyShardedStrategy  # noqa: F401
+from pytorch_lightning.strategies.fully_sharded_native import DDPFullyShardedNativeStrategy  # noqa: F401
 from pytorch_lightning.strategies.horovod import HorovodStrategy  # noqa: F401
 from pytorch_lightning.strategies.ipu import IPUStrategy  # noqa: F401
 from pytorch_lightning.strategies.parallel import ParallelStrategy  # noqa: F401

pytorch_lightning/strategies/fully_sharded_native.py

@@ -14,39 +14,31 @@
 import contextlib
 import logging
 import os
-from typing import Union, Any, Generator, Dict, List, Optional
+from typing import Any, Dict, Generator, List, Optional, Union
 
-import pytorch_lightning as pl
 import torch
+from torch.distributed.distributed_c10d import _get_default_group
+from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch, CPUOffload, FullyShardedDataParallel
+from torch.distributed.fsdp.wrap import enable_wrap
+
+import pytorch_lightning as pl
 from pytorch_lightning.overrides.distributed import prepare_for_backward
-from pytorch_lightning.plugins.environments.cluster_environment import (
-    ClusterEnvironment,
-)
+from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.parallel import ParallelStrategy
 from pytorch_lightning.utilities import rank_zero_only
 from pytorch_lightning.utilities.distributed import (
-    init_dist_connection,
-    sync_ddp_if_available,
-    ReduceOp,
-    group as _group,
     _get_process_group_backend_from_env,
     distributed_available,
     get_default_process_group_backend_for_device,
 )
+from pytorch_lightning.utilities.distributed import group as _group
+from pytorch_lightning.utilities.distributed import init_dist_connection, ReduceOp, sync_ddp_if_available
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_8
 from pytorch_lightning.utilities.optimizer import optimizers_to_device
 from pytorch_lightning.utilities.seed import reset_seed
-from torch.distributed.distributed_c10d import _get_default_group
-from torch.distributed.fsdp.fully_sharded_data_parallel import (
-    BackwardPrefetch,
-    CPUOffload,
-    FullyShardedDataParallel,
-)
-from torch.distributed.fsdp.wrap import enable_wrap
-
 
 log = logging.getLogger(__name__)
 
@@ -103,9 +95,7 @@ def __init__(
             precision_plugin=precision_plugin,
         )
         self._process_group = None
-        self.num_processes = (
-            len(self.parallel_devices) if self.parallel_devices is not None else 0
-        )
+        self.num_processes = len(self.parallel_devices) if self.parallel_devices is not None else 0
         self._has_loaded_state_dict: bool = False
         self._process_group_backend: Optional[str] = process_group_backend
         self.cpu_offload = cpu_offload
@@ -173,9 +163,7 @@ def setup(self, trainer: "pl.Trainer") -> None:
 
     def model_to_device(self) -> None:
         # ensure we update the device type in the lightning module
-        log.info(
-            f"{self.__class__.__name__}: moving model to device [{self.root_device}]..."
-        )
+        log.info(f"{self.__class__.__name__}: moving model to device [{self.root_device}]...")
         self.lightning_module.to(self.root_device)
 
     @contextlib.contextmanager

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -728,7 +728,10 @@ def _validate_precision_choice(self) -> None:
                 "it's not supported. Try using `amp_type='native'` instead."
             )
         if self._precision_flag in (16, "bf16") and self._amp_type_flag == AMPType.APEX:
-            if isinstance(self.strategy, (DDPShardedStrategy, DDPSpawnShardedStrategy, DDPFullyShardedStrategy, DDPFullyShardedNativeStrategy)):
+            if isinstance(
+                self.strategy,
+                (DDPShardedStrategy, DDPSpawnShardedStrategy, DDPFullyShardedStrategy, DDPFullyShardedNativeStrategy),
+            ):
                 raise MisconfigurationException(
                     "Sharded plugins are not supported with apex, please switch to `amp_backend='native'`."
                 )

tests/strategies/test_ddp_fully_sharded_native.py

@@ -4,21 +4,23 @@
 
 import pytest
 import torch
+from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
+from torch.distributed.fsdp.wrap import wrap
+
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.plugins import FullyShardedNativeMixedPrecisionPlugin
 from pytorch_lightning.strategies import DDPFullyShardedNativeStrategy
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers.boring_model import BoringModel
 from tests.helpers.runif import RunIf
-from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
-from torch.distributed.fsdp.wrap import wrap
 
 
 def test_invalid_on_cpu(tmpdir):
     """Test to ensure that to raise Misconfiguration for Native FSDP on CPU."""
     with pytest.raises(
-        MisconfigurationException, match="You selected strategy to be `ddp_fully_sharded_native`, but GPU is not available."
+        MisconfigurationException,
+        match="You selected strategy to be `ddp_fully_sharded_native`, but GPU is not available.",
     ):
         trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, strategy="fsdp_native")
         assert isinstance(trainer.strategy, DDPFullyShardedNativeStrategy)
@@ -102,7 +104,9 @@ def test_fully_sharded_native_strategy_checkpoint_multi_gpus(tmpdir):
 
     model = TestFSDPModel()
     ck = ModelCheckpoint(save_last=True)
-    trainer = Trainer(default_root_dir=tmpdir, gpus=2, strategy="fsdp_native", precision=16, max_epochs=1, callbacks=[ck])
+    trainer = Trainer(
+        default_root_dir=tmpdir, gpus=2, strategy="fsdp_native", precision=16, max_epochs=1, callbacks=[ck]
+    )
     _run_multiple_stages(trainer, model)