fix fsdp strateg pytorch

Lightning-AI · Sep 11, 2024 · d3a527d · d3a527d
1 parent e3f530a
commit d3a527d
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 6 deletions.
diff --git a/src/lightning_habana/pytorch/strategies/fsdp.py b/src/lightning_habana/pytorch/strategies/fsdp.py
@@ -53,7 +53,6 @@
 else:
     raise ModuleNotFoundError("You are missing `lightning` or `pytorch-lightning` package, please install it.")
 
-from lightning_habana.pytorch.accelerator import HPUAccelerator
 from lightning_habana.pytorch.plugins.fsdp_precision import HPUFSDPPrecision
 from lightning_habana.pytorch.plugins.io_plugin import HPUCheckpointIO
 from lightning_habana.pytorch.strategies.parallel import HPUParallelStrategy
@@ -101,10 +100,7 @@ def __init__(
     ) -> None:
         if not _LIGHTNING_GREATER_EQUAL_2_3_0:
             raise OSError("HPUFSDPStrategy requires `lightning>=2.3.0 or pytorch-lightning >= 2.3.0`.")
-        if parallel_devices is None:
-            parallel_devices = [torch.device("hpu", torch.hpu.current_device())] * HPUAccelerator.auto_device_count()
-        elif torch.device("hpu") in parallel_devices:
-            parallel_devices = [torch.device("hpu", torch.hpu.current_device())] * len(parallel_devices)
+
         super().__init__(
             accelerator=accelerator,
             parallel_devices=parallel_devices,

diff --git a/src/lightning_habana/pytorch/strategies/parallel.py b/src/lightning_habana/pytorch/strategies/parallel.py
@@ -137,7 +137,6 @@ def setup_distributed(self) -> None:
         _init_dist_connection(self.cluster_environment, self._process_group_backend, timeout=self._timeout)
 
     def _get_process_group_backend(self) -> str:
-        assert self.root_device.type == "hpu"
         return "hccl"
 
     def set_world_ranks(self) -> None: