Accelerator model state dict

CHANGELOG.md

@@ -28,6 +28,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Changed `resolve_training_type_plugins` to allow setting `num_nodes` and `sync_batchnorm` from `Trainer` setting ([7026](https://github.com/PyTorchLightning/pytorch-lightning/pull/7026))
 
 
+- Changed `model.state_dict()` in `CheckpointConnector` to allow `training_type_plugin` to customize the model's `state_dict()` ([7474](https://github.com/PyTorchLightning/pytorch-lightning/pull/7474))
+
+
 ### Deprecated
 
 

pytorch_lightning/accelerators/accelerator.py

@@ -420,6 +420,12 @@ def optimizer_state(self, optimizer: Optimizer) -> Dict[str, Tensor]:
  """
  return getattr(self.training_type_plugin, 'optimizer_state', lambda x: x.state_dict())(optimizer)
 
+ def lightning_module_state_dict(self) -> Dict[str, Union[Any, Tensor]]:
+ """
+ Returns state of model. Allows for syncing/collating model state from processes in custom plugins.
+ """
+ return self.training_type_plugin.lightning_module_state_dict()
+
  def on_save(self, checkpoint: Dict[str, Union[Any, Tensor]]) -> Dict[str, Union[Any, Tensor]]:
  return self.training_type_plugin.on_save(checkpoint)
 

pytorch_lightning/plugins/training_type/training_type_plugin.py

@@ -16,6 +16,7 @@
 from typing import Any, Callable, Dict, Generator, Iterable, Optional, Tuple, TypeVar, Union
 
 import torch
+from torch import Tensor
 from torch.nn import Module
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
@@ -241,6 +242,11 @@ def update_global_step(self, total_batch_idx: int, current_global_step: int) ->
  """
  return current_global_step + 1
 
+ def lightning_module_state_dict(self) -> Dict[str, Union[Any, Tensor]]:
+ """Returns model state."""
+ model = self.lightning_module
+ return model.state_dict()
+
  def save_checkpoint(self, checkpoint: Dict[str, Any], filepath: str) -> None:
  """Save model/training states as a checkpoint file through state-dump and file-write.
 

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -273,7 +273,7 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
  'epoch': current_epoch,
  'global_step': global_step,
  'pytorch-lightning_version': pytorch_lightning.__version__,
- 'state_dict': model.state_dict(),
+ 'state_dict': self.trainer.accelerator.lightning_module_state_dict(),
  }
 
  if not weights_only: