Updated lightning version requirement for FSDP (#182)

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added support for LightningCLI. ([#173](https://github.com/Lightning-AI/lightning-Habana/pull/173))
 - Added experimental support for FSDP on HPU. ([#174](https://github.com/Lightning-AI/lightning-Habana/pull/174))
 - Added support for FP8 inference with DeepSpeed. ([#176](https://github.com/Lightning-AI/lightning-Habana/pull/176))
+- Updated the lightning version check for using FSDP. ([#182](https://github.com/Lightning-AI/lightning-Habana/pull/182))
 
 
 ### Changed

docs/source/advanced.rst

@@ -546,6 +546,12 @@ Limitations of FSDP on HPU
 
 For more details on the supported FSDP features and functionalities, and limitations refer to `Using Fully Sharded Data Parallel (FSDP) with Intel Gaudi <https://docs.habana.ai/en/latest/PyTorch/PyTorch_FSDP/Pytorch_FSDP.html>`_.
 
+.. note::
+
+    This is an experimental feature.
+    This feature requires lightning/pytorch-lightning >= 2.3.0 or install nightly from the source.
+
+
 ----
 
 Using HPU Graphs

src/lightning_habana/pytorch/strategies/fsdp.py

@@ -60,7 +60,7 @@
 from lightning_habana.pytorch.plugins.io_plugin import HPUCheckpointIO
 from lightning_habana.pytorch.strategies.parallel import HPUParallelStrategy, _hpu_broadcast_object_list
 from lightning_habana.utils.hpu_distributed import _sync_ddp_if_available
-from lightning_habana.utils.imports import _HABANA_FRAMEWORK_AVAILABLE, _LIGHTNING_LESSER_EQUAL_2_2_3
+from lightning_habana.utils.imports import _HABANA_FRAMEWORK_AVAILABLE, _LIGHTNING_GREATER_EQUAL_2_3_0
 
 if _HABANA_FRAMEWORK_AVAILABLE:
     import habana_frameworks.torch.distributed.hccl as hpu_dist
@@ -105,8 +105,8 @@ def __init__(
         state_dict_type: Literal["full", "sharded"] = "full",
         **kwargs: Any,
     ) -> None:
-        if _LIGHTNING_LESSER_EQUAL_2_2_3:
-            raise OSError("HPUFSDPStrategy requires `pytorch-lightning > 2.2.3`.")
+        if not _LIGHTNING_GREATER_EQUAL_2_3_0:
+            raise OSError("HPUFSDPStrategy requires `lightning>=2.3.0 or pytorch-lightning >= 2.3.0`.")
         super().__init__(
             accelerator=accelerator,
             parallel_devices=parallel_devices,

src/lightning_habana/utils/imports.py

@@ -27,8 +27,8 @@
 _LIGHTNING_GREATER_EQUAL_2_0_0 = compare_version("lightning", operator.ge, "2.0.0") or compare_version(
     "pytorch_lightning", operator.ge, "2.0.0"
 )
-_LIGHTNING_LESSER_EQUAL_2_2_3 = compare_version("lightning", operator.le, "2.2.3") or compare_version(
-    "pytorch_lightning", operator.le, "2.2.3"
+_LIGHTNING_GREATER_EQUAL_2_3_0 = compare_version("lightning", operator.ge, "2.3.0", True) or compare_version(
+    "pytorch_lightning", operator.ge, "2.3.0", True
 )
 _TORCHVISION_AVAILABLE = RequirementCache("torchvision")
 _KINETO_AVAILABLE = torch.profiler.kineto_available()