Fuse some popular functions and automatically replace modules in an existing 🤗 transformers model with their corresponding fusion module
TODOs
[x] Decide on location to implement fused functions
[x] 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?
From the issue description:
Fuse some popular functions and automatically replace modules in an existing 🤗 transformers model with their corresponding fusion module
TODOs
[x] Decide on location to implement fused functions
[x] 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?
[x] Fuse bias addition and GeLU [link]
[x] Fuse bias addition and dropout [link]
[x] Error-check dropout code at first glance to make sure forward/backward passes are correct
[x] Write unit tests for fusion layers
[ ] Write test for toy model
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)