feat: support force downcast after FastRMSNorm multiply for Gemma #1658
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Narsil
reviewed
Mar 20, 2024
| @@ -687,7 +687,7 @@ def load(cls, prefix, weights, eps=1e-6): | |||
| weight = weights.get_tensor(f"{prefix}.weight") | |||
| return cls(weight, eps) | |||
|
|
|||
| def forward(self, hidden_states, residual=None): | |||
| def forward(self, hidden_states, residual=None, force_downcast_after=False): | |||
There was a problem hiding this comment.
I'd personally use a different method forward_downcast_after.
Also this only triggers when hidden_size > 8192 meaning it won't trigger for gemma.
Having a dedicated method seems simpler.
There was a problem hiding this comment.
that makes alot of sense thanks for the comments. I've removed the branching logic from FastRMSNorm and added forward_downcast_after into the gemma modeling code
kdamaszk
pushed a commit
to kdamaszk/tgi-gaudi
that referenced
this pull request
Apr 29, 2024
…ggingface#1658) This PR adds `force_downcast_after` to `FastRMSNorm.forward` which is used in the Gemma model. References huggingface/transformers#29402 and huggingface/transformers#29729 Setting `force_downcast_after=True` will perform the `hidden_states * weight` multiplication in f32 and then downcast to half. This differs slightly from the current implementation which first casts the `hidden_states` to a half and then multiples.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR adds
force_downcast_aftertoFastRMSNorm.forwardwhich is used in the Gemma model. References huggingface/transformers#29402 and huggingface/transformers#29729Setting
force_downcast_after=Truewill perform thehidden_states * weightmultiplication in f32 and then downcast to half. This differs slightly from the current implementation which first casts thehidden_statesto a half and then multiples.