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

[Model] Add has_weight to RMSNorm and re-enable weights loading tracker for Mamba #10739

Merged
merged 3 commits into from
Dec 10, 2024

Conversation

Isotr0py
Copy link
Collaborator

@Isotr0py Isotr0py commented Nov 28, 2024

This PR is a following PR for #10456

Copy link

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

  • Add ready label to the PR
  • Enable auto-merge.

🚀

Comment on lines 33 to 35
self.weight = torch.ones(hidden_size)
if self.elementwise_affine:
self.weight = nn.Parameter(self.weight)
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Just curious: can we also gain performance improvement from a rms_norm kernel without weight multiply for elementwise_affine=False?

Copy link
Collaborator

Choose a reason for hiding this comment

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

My take on this is that the opportunity is pretty small, with RMSNorm typically taking maybe 2% of the time in the forward pass -- you would save on not having to load self.weight from GPU memory, but since it's just a vector this is probably not a significant cost. And since we have custom layernorm kernels, modifying them to skip the weight application is probably not worth the effort.

However, I do think it's worth checking out if we can get speed up by skipping the weight application in the native forward method (using torch.compile for fusion)

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

I tested the performance of weight application skipped in the native forward with torch.compile, and there is a minor performance speedup (about 4%) on long sequence compared to no skip. So I think we can skip the weight application in native forward, although the speedup is small. :)

RMS_norm_performance
(Tested on a Tesla T4 16G)

@Isotr0py Isotr0py requested a review from tlrmchlsmth November 28, 2024 09:01
@Isotr0py
Copy link
Collaborator Author

Isotr0py commented Dec 4, 2024

@tlrmchlsmth Could you please take a look to this PR? Thanks!

Copy link
Collaborator

@tlrmchlsmth tlrmchlsmth left a comment

Choose a reason for hiding this comment

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

Looks good to me, thanks for the ping and sorry I missed it earlier!

vllm/model_executor/layers/layernorm.py Outdated Show resolved Hide resolved
Comment on lines 33 to 35
self.weight = torch.ones(hidden_size)
if self.elementwise_affine:
self.weight = nn.Parameter(self.weight)
Copy link
Collaborator

Choose a reason for hiding this comment

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

My take on this is that the opportunity is pretty small, with RMSNorm typically taking maybe 2% of the time in the forward pass -- you would save on not having to load self.weight from GPU memory, but since it's just a vector this is probably not a significant cost. And since we have custom layernorm kernels, modifying them to skip the weight application is probably not worth the effort.

However, I do think it's worth checking out if we can get speed up by skipping the weight application in the native forward method (using torch.compile for fusion)

@Isotr0py Isotr0py changed the title [Model] Add elementwise_affine to RMSNorm and re-enable weights loading tracker for Mamba [Model] Add has_weight to RMSNorm and re-enable weights loading tracker for Mamba Dec 5, 2024
@Isotr0py Isotr0py added the ready ONLY add when PR is ready to merge/full CI is needed label Dec 7, 2024
@Isotr0py Isotr0py merged commit d1f6d1c into vllm-project:main Dec 10, 2024
52 of 53 checks passed
@Isotr0py Isotr0py deleted the rmsnorm-affine branch December 10, 2024 02:23
sleepwalker2017 pushed a commit to sleepwalker2017/vllm that referenced this pull request Dec 13, 2024
BKitor pushed a commit to BKitor/vllm that referenced this pull request Dec 30, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ready ONLY add when PR is ready to merge/full CI is needed
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants