Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Issue #10: Kernel Fusion using torch.jit #36

Open
wants to merge 21 commits into
base: main
Choose a base branch
from

Conversation

sami-bg
Copy link

@sami-bg sami-bg commented Nov 19, 2023

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)

@sami-bg sami-bg mentioned this pull request Nov 19, 2023
3 tasks
tests/nn/test_fusion.py Outdated Show resolved Hide resolved
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants