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

Apply quantization on megablox kernel; support both training and serving #1100

Merged
merged 1 commit into from
Dec 17, 2024
Merged

Apply quantization on megablox kernel; support both training and serving #1100

merged 1 commit into from
Dec 17, 2024

Conversation

lenscloth
Copy link
Collaborator

@lenscloth lenscloth commented Dec 16, 2024 •
edited
Loading

Description

Support quantization on megablox for both training and serving

Tests

End to end tests on mixtral-8x22b model with max engine.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed.

@lenscloth lenscloth force-pushed the quantize_megablox_squashed branch from 7878fca to e4bee90 Compare December 16, 2024 20:31
RissyRan
RissyRan reviewed Dec 16, 2024
View reviewed changes
MaxText/layers/linears.py Outdated Show resolved Hide resolved
MaxText/layers/linears.py Outdated Show resolved Hide resolved
MaxText/layers/linears.py Show resolved Hide resolved
MaxText/layers/linears.py Show resolved Hide resolved
RissyRan
RissyRan approved these changes Dec 16, 2024
View reviewed changes
Copy link
Collaborator

@RissyRan RissyRan left a comment

Choose a reason for hiding this comment

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

Thank you! Others LGTM

@lenscloth lenscloth force-pushed the quantize_megablox_squashed branch 3 times, most recently from 8fa9666 to 9dd6476 Compare December 16, 2024 23:57
@lenscloth lenscloth force-pushed the quantize_megablox_squashed branch from 9dd6476 to 41367a3 Compare December 17, 2024 00:03
@lenscloth lenscloth force-pushed the quantize_megablox_squashed branch from e9c3d46 to 3197833 Compare December 17, 2024 01:33
@copybara-service copybara-service bot merged commit 8f2a60a into main Dec 17, 2024
14 checks passed
@copybara-service copybara-service bot deleted the quantize_megablox_squashed branch December 17, 2024 18:06
@ZhiyuLi-goog ZhiyuLi-goog mentioned this pull request Dec 27, 2024
Merged
4 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ZhiyuLi-goog ZhiyuLi-goog ZhiyuLi-goog approved these changes

@gobbleturk gobbleturk gobbleturk approved these changes

@RissyRan RissyRan RissyRan approved these changes

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

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

Successfully merging this pull request may close these issues.
4 participants
@lenscloth @RissyRan @gobbleturk @ZhiyuLi-goog