Minor LightningLite clean up (#15780)

src/lightning_lite/lite.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 import inspect
 import os
-from abc import ABC
 from contextlib import contextmanager, nullcontext
 from functools import partial
 from pathlib import Path
@@ -55,7 +54,7 @@
 from lightning_lite.wrappers import _LiteDataLoader, _LiteModule, _LiteOptimizer
 
 
-class LightningLite(ABC):
+class LightningLite:
     """Lite accelerates your PyTorch training or inference code with minimal changes required.
 
     - Automatic placement of models and data onto the device.

src/lightning_lite/strategies/fsdp.py

@@ -306,3 +306,9 @@ def no_backward_sync(self, module: Module) -> Generator:
             )
         with module.no_sync():
             yield
+
+
+def _optimizer_has_flat_params(optimizer: Optimizer) -> bool:
+    from torch.distributed.fsdp import FlatParameter
+
+    return any(isinstance(param, FlatParameter) for param in optimizer.param_groups[0]["params"])

src/pytorch_lightning/strategies/fully_sharded_native.py

@@ -20,7 +20,7 @@
 
 import pytorch_lightning as pl
 from lightning_lite.plugins import CheckpointIO, ClusterEnvironment
-from lightning_lite.strategies.fsdp_native import _optimizer_has_flat_params
+from lightning_lite.strategies.fsdp import _optimizer_has_flat_params
 from lightning_lite.utilities.distributed import (
     _get_default_process_group_backend_for_device,
     _init_dist_connection,