Updating 1.0.3 -> 1.0.4 breaks pytorch module hooks

[ludwigwinkler]

1.0.3 -> 1.0.4 breaks pytorch-hooks

What is your question?

I implemented a custom optimizer that stores the input and the incoming gradients for each layer and recomputes the gradients in a per-sample fashion (wich I do because I want to use K-FAC for my optimization).

I'm using the pytorch hooks on modules (i.e. module123.register_forward_pre_hook(_save_input_function) which will store the inputs into the state dictionary of the optimizer.
I use the same mod123.register_backward_hook(_save_gradient_function) for the backward gradients) and when calling optimizer.step().

As I updated from 1.0.1 to 1.1.0 it broke the hooks that are required in the KFAC optimizer.
I checked different versions and from 1.0.3 to 1.0.4 something changed such that the Trainer ignores/steps over the hooks in the forward and backward pass.

I checked the release notes but couldn't identify anything essential that would warrant these breaking changes.

Additionally they were issue popping up with manually modifying gradients in other sampling-based types of optimizers.

These issue do not exist for the standard optimizers straight from PyTorch optim class.

Code

It would be a huge codebase as the implementation of KFAC is a bit sophisticated.
But if nothing helps I could write an minimal working example from scratch.

[williamFalcon]

@tchaton

[williamFalcon]

@ludwigwinkler i suspect using the boring model with an optimizer that overrides a hook will replicate this problem

[ludwigwinkler]

Hi,

this is a boring model with an Optimizer that uses forward and backward hooks to do stuff with parameters
https://colab.research.google.com/drive/1z69hmWUybxW58iGMR9EQ3wlT5jwxJxzt?usp=sharing

Let me know if something is unclear/needs more comments or needs polishing.

Thanks for all the effort you put into this project. =)

[ludwigwinkler]

@tchaton Did you have time to look at this?

[lizhitwo]

I had the same bug, but it turns out to be: when retiring, pytorch_lightning.core.optimizer.LightningOptimizer calls its __del__, which in turn calls KFAC's __del__, which deletes all the hooks.

From pytorch lightning's stand point, it is very horrifying that it automatically calls my optimizer's __del__ even though I am still using my optimizer somewhere.

[ludwigwinkler]

@lizhitwo thank you for your insight!
where can I find the deletion call in the source code?
Or should I overwrite my custom optimizers del function?

[lizhitwo]

There’s no explicit calling but you can try it by constructing a lightning optimizer that wraps KFAC, and then set it to None while retaining KFAC, and then call garbage collector (gc.collect()). I think you should then see KFAC del called although you still have a KFAC object.  On Fri, Jan 22, 2021 at 9:12 AM, ludwigwinkler <notifications@github.com> wrote: @lizhitwo thank you for your insight! where can I find the deletion call in the source code? Or should I overwrite my custom optimizers del function? —You are receiving this because you were mentioned.Reply to this email directly, view it on GitHub, or unsubscribe.