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.

Sign up for GitHub

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

Jump to bottom

enable custom activation functions #65

Merged
merged 3 commits into from
Dec 19, 2023
Merged

enable custom activation functions #65

merged 3 commits into from
Dec 19, 2023

Conversation

vchiley
Copy link
Collaborator

@vchiley vchiley commented Dec 19, 2023

No description provided.

tgale96
tgale96 reviewed Dec 19, 2023
View reviewed changes
Copy link
Contributor

@tgale96 tgale96 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks awesome!

megablocks/layers/glu.py Outdated Show resolved Hide resolved
megablocks/layers/glu.py Show resolved Hide resolved
megablocks/layers/mlp.py Outdated Show resolved Hide resolved
megablocks/layers/glu.py Outdated Show resolved Hide resolved
megablocks/layers/mlp.py Outdated Show resolved Hide resolved
megablocks/layers/mlp.py Outdated Show resolved Hide resolved
megablocks/layers/mlp.py Outdated Show resolved Hide resolved
megablocks/layers/mlp.py Outdated Show resolved Hide resolved
megablocks/layers/mlp.py Outdated Show resolved Hide resolved
megablocks/layers/arguments.py Outdated Show resolved Hide resolved
@tgale96
Copy link
Contributor

tgale96 commented Dec 19, 2023

LGTM! Tests passing? Ok to merge?

@vchiley vchiley marked this pull request as ready for review December 19, 2023 18:19
@vchiley vchiley force-pushed the custom_act_fn branch 2 times, most recently from 825671d to ff9c2de Compare December 19, 2023 18:30
@vchiley
Copy link
Collaborator Author

vchiley commented Dec 19, 2023

LGTM! Tests passing? Ok to merge?

the tests you have seem to pass
It runs in llm-foundry for configs

  1. mlp_type: glu + grouped_mlp: true
  2. mlp_type: glu + grouped_mlp: false
  3. mlp_type: mlp + grouped_mlp: false + memory_optimized_mlp: false
  4. mlp_type: mlp + grouped_mlp: true + memory_optimized_mlp: false
  5. mlp_type: mlp + grouped_mlp: false + memory_optimized_mlp: true
  6. mlp_type: mlp + grouped_mlp: true + memory_optimized_mlp: true

@vchiley vchiley requested a review from tgale96 December 19, 2023 18:33
@vchiley
Copy link
Collaborator Author

vchiley commented Dec 19, 2023

The only functionality being used from megablocks.layers import gelu is gelu.gelu_backward_
@tgale96 how much do you think that buy us?

@tgale96
Copy link
Contributor

tgale96 commented Dec 19, 2023

The only functionality being used from megablocks.layers import gelu is gelu.gelu_backward_ @tgale96 how much do you think that buy us?

Probably not much if anything. I think this was a holdover from when the approx gelu wasn't supported in Torch, which I then reused for gelu remat.

@tgale96 tgale96 merged commit 44d9743 into databricks:main Dec 19, 2023
@vchiley vchiley deleted the custom_act_fn branch December 19, 2023 20:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@tgale96 tgale96 Awaiting requested review from tgale96

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@vchiley @tgale96