question: Questions and problems with different precisions
context: Hi,
  I try to implement an autoencoder where the latent space is not euclidean but has hyperbolic and spherical geometry. I use geoopt for this. They recommend to use double precision for numerical stability. When I run my LightningModule with precision=32 I encounter NaNs in the training. When I change to precision=64 I get this pickling error:
  Traceback (most recent call last):
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 599, in run_train
      self.train_loop.run_training_epoch()
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 494, in run_training_epoch
      self.check_checkpoint_callback(True)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 149, in check_checkpoint_callback
      cb.on_validation_end(self.trainer, model)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 238, in on_validation_end
      self.save_checkpoint(trainer)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 281, in save_checkpoint
      self._save_none_monitor_checkpoint(trainer, monitor_candidates)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 660, in _save_none_monitor_checkpoint
      self._save_model(trainer, filepath)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 410, in _save_model
      self._do_save(trainer, filepath)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 422, in _do_save
      self.save_function(filepath, self.save_weights_only)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/properties.py", line 336, in save_checkpoint
      self.checkpoint_connector.save_checkpoint(filepath, weights_only)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py", line 393, in save_checkpoint
      self.trainer.accelerator.save_checkpoint(_checkpoint, filepath)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 490, in save_checkpoint
      self.training_type_plugin.save_checkpoint(checkpoint, filepath)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 254, in save_checkpoint
      atomic_save(checkpoint, filepath)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/utilities/cloud_io.py", line 63, in atomic_save
      torch.save(checkpoint, bytesbuffer)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/torch/serialization.py", line 372, in save
      _save(obj, opened_zipfile, pickle_module, pickle_protocol)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/torch/serialization.py", line 476, in _save
      pickler.dump(obj)
  _pickle.PicklingError: Can't pickle <function LinkPredictionModule.training_step at 0x7f08220f0ee0>: it's not the same object as stereographic_link_prediction.Models.Modules.LinkPredictionModule.training_step
  
  During handling of the above exception, another exception occurred:
  
  Traceback (most recent call last):
    File "main.py", line 40, in <module>
      trainer.fit(module, datamodule=datamodule)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 480, in fit
      self.dispatch()
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 523, in dispatch
      self.accelerator.start_training(self)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 95, in start_training
      self.training_type_plugin.start_training(trainer)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 142, in start_training
      self._results = trainer.run_stage()
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 535, in run_stage
      self.run_train()
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 634, in run_train
      self.train_loop.on_train_end()
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 117, in on_train_end
      self.check_checkpoint_callback(should_update=True, is_last=True)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 149, in check_checkpoint_callback
      cb.on_validation_end(self.trainer, model)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 238, in on_validation_end
      self.save_checkpoint(trainer)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 281, in save_checkpoint
      self._save_none_monitor_checkpoint(trainer, monitor_candidates)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 660, in _save_none_monitor_checkpoint
      self._save_model(trainer, filepath)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 410, in _save_model
      self._do_save(trainer, filepath)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 422, in _do_save
      self.save_function(filepath, self.save_weights_only)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/properties.py", line 336, in save_checkpoint
      self.checkpoint_connector.save_checkpoint(filepath, weights_only)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py", line 393, in save_checkpoint
      self.trainer.accelerator.save_checkpoint(_checkpoint, filepath)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 490, in save_checkpoint
      self.training_type_plugin.save_checkpoint(checkpoint, filepath)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 254, in save_checkpoint
      atomic_save(checkpoint, filepath)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/pytorch_lightning/utilities/cloud_io.py", line 63, in atomic_save
      torch.save(checkpoint, bytesbuffer)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/torch/serialization.py", line 372, in save
      _save(obj, opened_zipfile, pickle_module, pickle_protocol)
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/torch/serialization.py", line 476, in _save
      pickler.dump(obj)
  _pickle.PicklingError: Can't pickle <function LinkPredictionModule.training_step at 0x7f08220f0ee0>: it's not the same object as stereographic_link_prediction.Models.Modules.LinkPredictionModule.training_step
  Exception ignored in: <function tqdm.__del__ at 0x7f086346a430>
  Traceback (most recent call last):
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/tqdm/std.py", line 1145, in __del__
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/tqdm/std.py", line 1299, in close
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/tqdm/std.py", line 1492, in display
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/tqdm/std.py", line 1148, in __str__
    File "/home/netter/.cache/pypoetry/virtualenvs/stereographic-link-prediction-ra14Y8Aq-py3.8/lib/python3.8/site-packages/tqdm/std.py", line 1450, in format_dict
  TypeError: cannot unpack non-iterable NoneType object
  
  I think this error comes from inconsistencies in my Module. But I do not know how to find them. Do you have any tips to track down the inconsistencies that disturb pickling?
  However, when I run with precision=16 I do not run into NaNs. As this precision should be more numerically unstable than precision=32 my only explanation for this is that there is automatic gradient clipping going on or so in 16-bit precision. Is this correct? And where do I find these values? Then I could apply them to 32-bit precision training.
answer: Hi @gatoniel, not sure about the NaNs but the pickle error definitely looks like a bug on our end. The double precision plugin patches the training_step which would explain that error. Mind opening an issue for it?