Questions and problems with different precisions

[gatoniel]

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.

[ethanwharris]

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?

[gatoniel]

Thanks, I was thinking about this, too. However, right now I am unable to replicate the error. It must have been the result of a specific combination of trainer arguments. I will open the issue, when I can replicate.

[ethanwharris]

Awesome! I had a go at trying to replicate it but couldn't manage to either haha. I'll do some more digging and let you know if I find anything :)