Implementation of activations as subclasses of the torch.nn.Module
Author: eldarkurtic
Pull request: no

# 🚀 Feature request

<!-- A clear and concise description of the feature proposal.
     Please provide a link to the paper and code in case they exist. -->
I guess it would be best to ask this first: are there some specific reasons why activation functions in https://github.com/huggingface/transformers/blob/master/src/transformers/activations.py are not subclasses of the `torch.nn.Module` ?

If there are, then we can probably ignore everything else below :) .
If there aren't, then it might be interesting to consider implementing them that way (I would be happy to work on a PR for it). 
A few advantages (that I'm aware of) with activations as subclasses of the `torch.nn.Module`:
1. it's easy to check which activations are used in the model by just running: `print(my_bert)`. Currently one has to check the config file for it, which is also not that bad but this just makes it a bit more convenient. Just like printing the torchvision models `print(resnet50)`, one can immediately see which activations are being used in the model.
2. composing layers with for example `nn.Sequential` would be possible (I'm not sure if this is possible when activations are implemented as python functions)
3. attaching pytorch hooks to activation modules would be possible (I think this is the most important advantage)


Comments:
Pinging @sgugger @patrickvonplaten @patil-suraj for consideration

Also pinging @mfuntowicz as I believe this may affect the graph of the models and would like to hear feedback on the request

No strong opinion on my side, apart from the fact we would need to be careful not to break anything in any existing model.

Ok for me to change if it doesn't break anything. I agree that seeing the activations when printing out the model is a considerably improvement!

okay for me as well as long as it doesn't break anything!

Also cc @stas00, @NielsRogge would love your input.

If everyone is on board, let's open a good first issue and open this to community contributions.

> * it's easy to check which activations are used in the model by just running: `print(my_bert)`. Currently one has to check the config file for it, which is also not that bad but this just makes it a bit more convenient. Just like printing the torchvision models `print(resnet50)`, one can immediately see which activations are being used in the model.

check, but this is trivial to do by modifying `__repr__` to print out the activation function.

> * composing layers with for example `nn.Sequential` would be possible (I'm not sure if this is possible when activations are implemented as python functions)

this is not a problem with functions, unless of course you need to have a super-fine "granularity". But you will have to change a ton of code for this to work. And making  `nn.Sequential` Transformers models is super-complicated , especially for some models - think conditional control flow.
 
> * attaching pytorch hooks to activation modules would be possible (I think this is the most important advantage)

check, this is indeed helpful when debugging training instabilities.

--------------

One issue is backward compatibility, if the models start using `nn.Module` activations - I don't think we could remove the function versions as others surely use those in their modified models.

-------------

the other issue is that we are likely to start having more complex activation functions soon which require `torch.autograd.Function` - basically if you want to manually tweak the `backward` function. Here is an example:

https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/dfb5d6873a0825ddbfecda009b2ac1f493f09a97/megatron/model/fused_bias_gelu.py#L47

and it's likely we will already need to use something similar for the Megatron-Deepspeed models BigScience has been training. 

I tried to search how to make `torch.autograd.Function` into `nn.Module` - but so far the only answers I found is that one has to use the former if they want a custom `backward` function. But surely there should be a way to then again fold it into `nn.Module` - I'd ask you to research this by asking on pytorch forums or slack.
update: I think this is the answer as someone already asked this https://discuss.pytorch.org/t/defining-backward-function-in-nn-module/5047/2 so you call a  `torch.autograd.Function` from `nn.Module`.

So I think it should be fine. It's just that I haven't tried doing it myself yet.

`torch.autograd.Function` docs: 

- https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html
- https://pytorch.org/docs/stable/notes/extending.html#extending-torch-autograd