Introduce Stateful PrecisionPlugin

pytorch_lightning/plugins/precision/apex_amp.py

@@ -96,5 +96,4 @@ def state_dict(self) -> Dict[str, Any]:
         return amp.state_dict()
 
     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        if state_dict:
-            amp.load_state_dict(state_dict)
+        amp.load_state_dict(state_dict)

pytorch_lightning/plugins/precision/native_amp.py

@@ -109,7 +109,9 @@ def forward_context(self) -> Generator[None, None, None]:
     def state_dict(self) -> Dict[str, Any]:
         if self.scaler is not None:
             return self.scaler.state_dict()
+        else:
+            return {}
 
     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        if self.scaler is not None and state_dict:
+        if self.scaler is not None:
             self.scaler.load_state_dict(state_dict)

pytorch_lightning/plugins/precision/precision_plugin.py

@@ -244,7 +244,17 @@ def teardown(self) -> None:
         """
 
     def state_dict(self) -> Dict[str, Any]:
+        """Called when saving a checkpoint, implement to generate and save precision plugin state.
+
+        Returns:
+            A dictionary containing precision plugin state.
+        """
         return {}
 
     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        """Called when loading a checkpoint, implement to reload precision plugin state given precision plugin state_dict.
+
+        Args:
+            state_dict: the precision plugin state returned by ``state_dict``.
+        """
         pass

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -194,8 +194,8 @@ def restore_training_state(self) -> None:
         # restore precision plugin (scaler etc.)
         prec_plugin = self.trainer.precision_plugin
         prec_plugin.on_load_checkpoint(self._loaded_checkpoint)
-        if prec_plugin.__class__.__name__ in self._loaded_checkpoint:
-            prec_plugin.load_state_dict(self._loaded_checkpoint[prec_plugin.__class__.__name__])
+        if prec_plugin.__class__.__qualname__ in self._loaded_checkpoint:
+            prec_plugin.load_state_dict(self._loaded_checkpoint[prec_plugin.__class__.__qualname__])
 
         # old checkpoints compatibility
         # should we raise error and force user to run utilities/upgrade_checkpoint instead?
@@ -384,7 +384,10 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
 
             # precision plugin
             prec_plugin = self.trainer.precision_plugin
-            checkpoint[prec_plugin.__class__.__name__] = self.trainer.precision_plugin.state_dict()
+            prec_plugin_state_dict = prec_plugin.state_dict()
+            if prec_plugin_state_dict:
+                checkpoint[prec_plugin.__class__.__qualname__] = prec_plugin_state_dict
+            prec_plugin.on_save_checkpoint(checkpoint)
 
         # dump hyper-parameters
         if model.hparams:
@@ -401,8 +404,6 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
         model.on_save_checkpoint(checkpoint)
         if self.trainer.datamodule is not None:
             self.trainer.datamodule.on_save_checkpoint(checkpoint)
-        if not weights_only:
-            self.trainer.precision_plugin.on_save_checkpoint(checkpoint)
 
         # TODO: remove this in v1.8.
         environment = self.trainer._accelerator_connector.cluster_environment