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

Use Llama RMSNorm for Gemma #2974

Merged
merged 3 commits into from
Feb 22, 2024
Merged

Use Llama RMSNorm for Gemma #2974

merged 3 commits into from
Feb 22, 2024

Conversation

WoosukKwon
Copy link
Collaborator

@WoosukKwon WoosukKwon commented Feb 22, 2024

Gemma's RMSNorm is only slightly different from Llama's RMSNorm. Thus, we can use the existing custom op for it. This optimization leads to ~10% latency reduction.

@WoosukKwon WoosukKwon requested a review from Yard1 February 22, 2024 01:03
Copy link
Collaborator

@Yard1 Yard1 left a comment

Choose a reason for hiding this comment

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

Nice!

@Yard1
Copy link
Collaborator

Yard1 commented Feb 22, 2024

Before we merge, let's make sure it doesn't change the outputs (maybe we could add a test like we have for other models, using transformers as a reference).

@WoosukKwon
Copy link
Collaborator Author

For a note, using the custom op brings a slight numerical difference in handling the residual connection.

While the original implementation uses the current dtype (f16 of bf16) in hidden_states + residual, the fused RMSNorm op upcasts both to FP32 before addition:

x = x + residual.to(torch.float32)

@WoosukKwon WoosukKwon merged commit 95529e3 into main Feb 22, 2024
4 of 6 checks passed
@WoosukKwon WoosukKwon deleted the optimize-gemma branch February 22, 2024 02:28
xjpang pushed a commit to xjpang/vllm that referenced this pull request Mar 4, 2024
Temirulan pushed a commit to Temirulan/vllm-whisper that referenced this pull request Sep 6, 2024
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.

2 participants