Issue #10: Kernel Fusion using torch.jit

[sami-bg]

From the issue description:

Fuse some popular functions and automatically replace modules in an existing 🤗 transformers model with their corresponding fusion module

TODOs

• Decide on location to implement fused functions

• Decide on a way to automatically replace modules on import, i.e. whenever someone imports a 🤗transformer model, its activations are replaced with the fused implementations. @xrsrke I'm not sure I understand this part properly?

• Fuse bias addition and GeLU [link]

• Fuse bias addition and dropout [link]

• Error-check dropout code at first glance to make sure forward/backward passes are correct

• Write unit tests for fusion layers

  • For some reason, the torch Dropout is zeroing out several more things in the unit test
  • The Fused Gelu is asking me for a static forward method, which is already implemented, so not sure what's going on there

• Write test for toy model

  • Relatively inexperienced here, so might need a hand

Reading (could be ignored)

OSLO’s kernel fusion [[link]](https://github.com/tunib-ai/oslo/blob/88dcca0441a605b462bf825cb0104bc692f14c57/oslo/fused_kernels_utils.py#L259)
GPT-NeoX’s kernel fusion [[link]](https://github.com/EleutherAI/gpt-neox/blob/b02d98932f95fe0500c28698b38acb175e92e980/megatron/model/activations.py#L27)