adapt pytorch lighting 2.0 AKA lightning

Author: qmpzzpmq

Diff for: nni/compression/pytorch/utils/evaluator.py

@@ -21,6 +21,13 @@
 else:
     LIGHTNING_INSTALLED = True
 
+try:
+    import lightning as L
+except ImportError:
+    LIGHTNING2_INSTALLED = False
+else:
+    LIGHTNING2_INSTALLED = True
+
 try:
     from transformers.trainer import Trainer as HFTrainer
 except ImportError:
@@ -161,7 +168,7 @@ def _init_optimizer_helpers(self, pure_model: Module | pl.LightningModule):
         """
         raise NotImplementedError
 
-    def bind_model(self, model: Module | pl.LightningModule, param_names_map: Dict[str, str] | None = None):
+    def bind_model(self, model: Module | pl.LightningModule | L.LightningModule, param_names_map: Dict[str, str] | None = None):
         """
         Bind the model suitable for this ``Evaluator`` to use the evaluator's abilities of model modification,
         model training, and model evaluation.
@@ -312,25 +319,27 @@ class LightningEvaluator(Evaluator):
     If the the test metric is needed by nni, please make sure log metric with key ``default`` in ``LightningModule.test_step()``.
     """
 
-    def __init__(self, trainer: pl.Trainer, data_module: pl.LightningDataModule,
+    def __init__(self, trainer: pl.Trainer | L.trainer, data_module: pl.LightningDataModule,
                  dummy_input: Any | None = None):
         assert LIGHTNING_INSTALLED, 'pytorch_lightning is not installed.'
         err_msg_p = 'Only support traced {}, please use nni.trace({}) to initialize the trainer.'
         err_msg = err_msg_p.format('pytorch_lightning.Trainer', 'pytorch_lightning.Trainer')
-        assert isinstance(trainer, pl.Trainer) and is_traceable(trainer), err_msg
+        lighting2_check = not LIGHTNING2_INSTALLED and isinstance(trainer, L.Trainer)
+        assert (isinstance(trainer, pl.Trainer) or lighting2_check) and is_traceable(trainer), err_msg
         err_msg = err_msg_p.format('pytorch_lightning.LightningDataModule', 'pytorch_lightning.LightningDataModule')
-        assert isinstance(data_module, pl.LightningDataModule) and is_traceable(data_module), err_msg
+        lighting2_check = not LIGHTNING2_INSTALLED and isinstance(data_module, L.LightningDataModule)
+        assert (isinstance(data_module, pl.LightningDataModule) or lighting2_check) and is_traceable(data_module), err_msg
         self.trainer = trainer
         self.data_module = data_module
         self._dummy_input = dummy_input
 
-        self.model: pl.LightningModule | None = None
+        self.model: pl.LightningModule | L.LightningModule | None = None
         self._ori_model_attr = {}
         self._param_names_map: Dict[str, str] | None = None
 
         self._initialization_complete = False
 
-    def _init_optimizer_helpers(self, pure_model: pl.LightningModule):
+    def _init_optimizer_helpers(self, pure_model: pl.LightningModule | L.LightningModule):
         assert self._initialization_complete is False, 'Evaluator initialization is already complete.'
 
         self._optimizer_helpers = []
@@ -395,10 +404,14 @@ def _init_optimizer_helpers(self, pure_model: pl.LightningModule):
 
         self._initialization_complete = True
 
-    def bind_model(self, model: pl.LightningModule, param_names_map: Dict[str, str] | None = None):
+    def bind_model(
+        self,
+        model: pl.LightningModule | L.LightningModule,
+        param_names_map: Dict[str, str] | None = None
+    ):
         err_msg = 'Evaluator initialization is not complete, please call `_init_optimizer_helpers` before bind model.'
         assert self._initialization_complete is True, err_msg
-        assert isinstance(model, pl.LightningModule)
+        assert isinstance(model, pl.LightningModule) or isinstance(model, L.LightningModule)
         if self.model is not None:
             _logger.warning('Already bound a model, will unbind it before bind a new model.')
             self.unbind_model()
@@ -425,7 +438,7 @@ def unbind_model(self):
             _logger.warning('Did not bind any model, no need to unbind model.')
 
     def _patch_configure_optimizers(self):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
 
         if self._opt_returned_dicts:
             def new_configure_optimizers(_):  # type: ignore
@@ -452,11 +465,11 @@ def new_configure_optimizers(_):
         self.model.configure_optimizers = types.MethodType(new_configure_optimizers, self.model)
 
     def _revert_configure_optimizers(self):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         self.model.configure_optimizers = self._ori_model_attr['configure_optimizers']
 
     def patch_loss(self, patch: Callable[[Tensor], Tensor]):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         old_training_step = self.model.training_step
 
         def patched_training_step(_, *args, **kwargs):
@@ -470,19 +483,28 @@ def patched_training_step(_, *args, **kwargs):
         self.model.training_step = types.MethodType(patched_training_step, self.model)
 
     def revert_loss(self):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         self.model.training_step = self._ori_model_attr['training_step']
 
     def patch_optimizer_step(self, before_step_tasks: List[Callable], after_step_tasks: List[Callable]):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
 
         class OptimizerCallback(Callback):
-            def on_before_optimizer_step(self, trainer: pl.Trainer, pl_module: pl.LightningModule,
-                                         optimizer: Optimizer, opt_idx: int) -> None:
+            def on_before_optimizer_step(
+                self,
+                trainer: pl.Trainer | L.Trainer,
+                pl_module: pl.LightningModule | L.LightningModule,
+                optimizer: Optimizer, opt_idx: int
+            ) -> None:
                 for task in before_step_tasks:
                     task()
 
-            def on_before_zero_grad(self, trainer: pl.Trainer, pl_module: pl.LightningModule, optimizer: Optimizer) -> None:
+            def on_before_zero_grad(
+                self,
+                trainer: pl.Trainer | L.trainer,
+                pl_module: pl.LightningModule | L.LightningModule,
+                optimizer: Optimizer,
+            ) -> None:
                 for task in after_step_tasks:
                     task()
 
@@ -496,13 +518,13 @@ def patched_configure_callbacks(_):
         self.model.configure_callbacks = types.MethodType(patched_configure_callbacks, self.model)
 
     def revert_optimizer_step(self):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         self.model.configure_callbacks = self._ori_model_attr['configure_callbacks']
 
     def train(self, max_steps: int | None = None, max_epochs: int | None = None):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         # reset trainer
-        trainer: pl.Trainer = self.trainer.trace_copy().get()  # type: ignore
+        trainer: pl.Trainer | L.Trainer = self.trainer.trace_copy().get()  # type: ignore
         # NOTE: lightning may dry run some steps at first for sanity check in Trainer.fit() by default,
         # If we want to record some information in the forward hook, we may get some additional information,
         # so using Trainer.num_sanity_val_steps = 0 disable sanity check.
@@ -529,9 +551,9 @@ def evaluate(self) -> Tuple[float | None, List[Dict[str, float]]]:
         E.g., if ``Trainer.test()`` returned ``[{'default': 0.8, 'loss': 2.3}, {'default': 0.6, 'loss': 2.4}]``,
         NNI will take the final metric ``(0.8 + 0.6) / 2 = 0.7``.
         """
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         # reset trainer
-        trainer: pl.Trainer = self.trainer.trace_copy().get()  # type: ignore
+        trainer: pl.Trainer | L.Trainer = self.trainer.trace_copy().get()  # type: ignore
         original_results = trainer.test(self.model, self.data_module)
         # del trainer reference, we don't want to dump trainer when we dump the entire model.
         self.model.trainer = None
@@ -831,7 +853,7 @@ def __init__(self, trainer: HFTrainer, dummy_input: Any | None = None) -> None:
 
         self._initialization_complete = False
 
-    def _init_optimizer_helpers(self, pure_model: Module | pl.LightningModule):
+    def _init_optimizer_helpers(self, pure_model: Module | pl.LightningModule | L.LightningModule):
         assert self._initialization_complete is False, 'Evaluator initialization is already complete.'
 
         if self.traced_trainer.optimizer is not None and is_traceable(self.traced_trainer.optimizer):
@@ -862,7 +884,7 @@ def patched_get_optimizer_cls_and_kwargs(args) -> Tuple[Any, Any]:
 
         self._initialization_complete = True
 
-    def bind_model(self, model: Module | pl.LightningModule, param_names_map: Dict[str, str] | None = None):
+    def bind_model(self, model: Module | pl.LightningModule | L.LightningModule, param_names_map: Dict[str, str] | None = None):
         err_msg = 'Evaluator initialization is not complete, please call `_init_optimizer_helpers` before bind model.'
         assert self._initialization_complete is True, err_msg
         assert isinstance(model, Module)

Diff for: nni/contrib/compression/utils/evaluator.py

@@ -22,6 +22,13 @@
 else:
     LIGHTNING_INSTALLED = True
 
+try:
+    import lightning as L
+except ImportError:
+    LIGHTNING2_INSTALLED = False
+else:
+    LIGHTNING2_INSTALLED = True
+
 try:
     from transformers.trainer import Trainer as HFTrainer
 except ImportError:
@@ -149,7 +156,7 @@ class Evaluator:
     _initialization_complete: bool
     _hook: List[Hook]
 
-    def _init_optimizer_helpers(self, pure_model: Module | pl.LightningModule):
+    def _init_optimizer_helpers(self, pure_model: Module | pl.LightningModule | L.LightningModule):
         """
         This is an internal API, ``pure_model`` means the model is the original model passed in by the user,
         it should not be the modified model (wrapped, hooked, or patched by NNI).
@@ -164,7 +171,7 @@ def _init_optimizer_helpers(self, pure_model: Module | pl.LightningModule):
         """
         raise NotImplementedError
 
-    def bind_model(self, model: Module | pl.LightningModule, param_names_map: Dict[str, str] | None = None):
+    def bind_model(self, model: Module | pl.LightningModule | L.LightningModule, param_names_map: Dict[str, str] | None = None):
         """
         Bind the model suitable for this ``Evaluator`` to use the evaluator's abilities of model modification,
         model training, and model evaluation.
@@ -186,8 +193,12 @@ def unbind_model(self):
         """
         raise NotImplementedError
 
-    def _optimizer_add_param_group(self, model: Union[torch.nn.Module, pl.LightningModule],
-                                   module_name_param_dict: Dict[str, List[Tensor]], optimizers: Optimizer | List[Optimizer]):
+    def _optimizer_add_param_group(
+        self,
+        model: Union[torch.nn.Module, pl.LightningModule, L.LightningModule],
+        module_name_param_dict: Dict[str, List[Tensor]],
+        optimizers: Optimizer | List[Optimizer]
+    ):
         # used in the bind_model process
         def find_param_group(param_groups: List[Dict], module_name: str):
             for i, param_group in enumerate(param_groups):
@@ -367,25 +378,33 @@ class LightningEvaluator(Evaluator):
     If the the test metric is needed by nni, please make sure log metric with key ``default`` in ``LightningModule.test_step()``.
     """
 
-    def __init__(self, trainer: pl.Trainer, data_module: pl.LightningDataModule,
+    def __init__(self, trainer: pl.Trainer, data_module: pl.LightningDataModule | L.LightningDataModule,
                  dummy_input: Any | None = None):
         assert LIGHTNING_INSTALLED, 'pytorch_lightning is not installed.'
         err_msg_p = 'Only support traced {}, please use nni.trace({}) to initialize the trainer.'
-        err_msg = err_msg_p.format('pytorch_lightning.Trainer', 'pytorch_lightning.Trainer')
-        assert isinstance(trainer, pl.Trainer) and is_traceable(trainer), err_msg
-        err_msg = err_msg_p.format('pytorch_lightning.LightningDataModule', 'pytorch_lightning.LightningDataModule')
+        err_msg = err_msg_p.format(
+            'pytorch_lightning.Trainer or lightning.Trainer', 
+            'pytorch_lightning.Trainer or lightning.Trainer',
+        )
+        lighting2_check = not LIGHTNING2_INSTALLED and isinstance(trainer, L.Trainer)
+        assert (isinstance(trainer, pl.Trainer) or lighting2_check)and is_traceable(trainer), err_msg
+        err_msg = err_msg_p.format(
+            'pytorch_lightning.LightningDataModule or lightning.LightningDataModule', 
+            'pytorch_lightning.LightningDataModule or lightning.LightningDataModule',
+        )
+        lighting2_check = not LIGHTNING2_INSTALLED and isinstance(data_module, L.LightningDataModule)
         assert isinstance(data_module, pl.LightningDataModule) and is_traceable(data_module), err_msg
-        self.trainer = trainer
-        self.data_module = data_module
+        self.trainer: pl.Trainer | L.Trainer = trainer
+        self.data_module: pl.LightningDataModule | L.LightningDataModule = data_module
         self._dummy_input = dummy_input
 
-        self.model: pl.LightningModule | None = None
+        self.model: pl.LightningModule | L.LightningModule | None = None
         self._ori_model_attr = {}
         self._param_names_map: Dict[str, str] | None = None
 
         self._initialization_complete = False
 
-    def _init_optimizer_helpers(self, pure_model: pl.LightningModule):
+    def _init_optimizer_helpers(self, pure_model: pl.LightningModule | L.LightningModule):
         assert self._initialization_complete is False, 'Evaluator initialization is already complete.'
 
         self._optimizer_helpers = []
@@ -450,7 +469,7 @@ def _init_optimizer_helpers(self, pure_model: pl.LightningModule):
 
         self._initialization_complete = True
 
-    def bind_model(self, model: pl.LightningModule, param_names_map: Dict[str, str] | None = None):
+    def bind_model(self, model: pl.LightningModule | L.LightningModule, param_names_map: Dict[str, str] | None = None):
         err_msg = 'Evaluator initialization is not complete, please call `_init_optimizer_helpers` before bind model.'
         assert self._initialization_complete is True, err_msg
         assert isinstance(model, pl.LightningModule)
@@ -514,7 +533,7 @@ def new_configure_optimizers(_):
         self.model.configure_optimizers = types.MethodType(new_configure_optimizers, self.model)
 
     def _patch_configure_optimizers(self):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         if self._opt_returned_dicts:
             def new_configure_optimizers(_):  # type: ignore
                 optimizers = [opt_helper.call(self.model, self._param_names_map) for opt_helper in self._optimizer_helpers]  # type: ignore
@@ -559,11 +578,11 @@ def patched_training_step(_, *args, **kwargs):
         self.model.training_step = types.MethodType(patched_training_step, self.model)
 
     def revert_loss(self):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         self.model.training_step = self._ori_model_attr['training_step']
 
     def patch_optimizer_step(self, before_step_tasks: List[Callable], after_step_tasks: List[Callable]):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         old_configure_optimizers = self.model.configure_optimizers
 
         def patched_step_factory(old_step):
@@ -599,13 +618,13 @@ def new_configure_optimizers(_):
         self.model.configure_optimizers = types.MethodType(new_configure_optimizers, self.model)
 
     def revert_optimizer_step(self):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         self.model.configure_callbacks = self._ori_model_attr['configure_callbacks']
 
     def train(self, max_steps: int | None = None, max_epochs: int | None = None):
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         # reset trainer
-        trainer: pl.Trainer = self.trainer.trace_copy().get()  # type: ignore
+        trainer: pl.Trainer | L.Trainer = self.trainer.trace_copy().get()  # type: ignore
         # NOTE: lightning may dry run some steps at first for sanity check in Trainer.fit() by default,
         # If we want to record some information in the forward hook, we may get some additional information,
         # so using Trainer.num_sanity_val_steps = 0 disable sanity check.
@@ -632,9 +651,9 @@ def evaluate(self) -> Tuple[float | None, List[Dict[str, float]]]:
         E.g., if ``Trainer.test()`` returned ``[{'default': 0.8, 'loss': 2.3}, {'default': 0.6, 'loss': 2.4}]``,
         NNI will take the final metric ``(0.8 + 0.6) / 2 = 0.7``.
         """
-        assert isinstance(self.model, pl.LightningModule)
+        assert isinstance(self.model, pl.LightningModule) or isinstance(self.model, L.LightningModule)
         # reset trainer
-        trainer: pl.Trainer = self.trainer.trace_copy().get()  # type: ignore
+        trainer: pl.Trainer | L.trainer = self.trainer.trace_copy().get()  # type: ignore
         original_results = trainer.test(self.model, self.data_module)
         # del trainer reference, we don't want to dump trainer when we dump the entire model.
         self.model.trainer = None