Skip to content

[MODEL] Add support for Zamba2 models #13185

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
Merged
merged 10 commits into from
Mar 18, 2025
Merged

[MODEL] Add support for Zamba2 models #13185

merged 10 commits into from
Mar 18, 2025
+1,082 −27

Conversation

@yury-tokpanov
Copy link
Contributor

@yury-tokpanov yury-tokpanov commented Feb 13, 2025 •
edited by github-actions bot
Loading

This PR adds support for Zamba2 models (#9382), a series of mamba2-transformer hybrid models with shared attention blocks and LoRAs, applied to shared MLP and attention blocks, depending on the model. 1.2B and 7B models use RoPE for their attention blocks.

This PR is fully compatible with Zamba2 integration in HuggingFace transformers library, which was recently merged into the main branch.

  • We are able to reproduce evaluation results using evaluation harness with vllm evaluator.
  • We also inspected logits and intermediate layers outputs, comparing them with our reference implementation. We find a good agreement between the two (given numerical precision).
  • Chunked prefill appears to be working.
  • TP is supported.
  • PP is not supported, and we believe it wouldn't make much sense to use our models with PP, since it will remove the memory advantage of shared attention layers.

Unit tests pass now.

We would like to acknowledge authors of Bamba PR and Mamba2 PR (@fabianlim and @tlrmchlsmth respectively) for adding mamba2 support to vLLM and having productive discussions!

cc: @Quentin-Anthony @BerenMillidge @pglorio

@github-actions
Copy link

github-actions bot commented Feb 13, 2025

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

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 either: Add ready label to the PR or enable auto-merge.

🚀

@yury-tokpanov yury-tokpanov marked this pull request as draft February 13, 2025 01:56
@mergify mergify bot added the ci/build label Feb 22, 2025
@yury-tokpanov yury-tokpanov marked this pull request as ready for review February 22, 2025 03:19
@yury-tokpanov yury-tokpanov force-pushed the zamba2 branch 2 times, most recently from 9593333 to 0a83814 Compare February 25, 2025 00:27
@tlrmchlsmth tlrmchlsmth self-assigned this Feb 25, 2025
tlrmchlsmth
tlrmchlsmth reviewed Feb 25, 2025
View reviewed changes
tlrmchlsmth
tlrmchlsmth reviewed Mar 3, 2025
View reviewed changes
@yury-tokpanov yury-tokpanov force-pushed the zamba2 branch 6 times, most recently from 66b1112 to 2b7397c Compare March 6, 2025 23:38
@yury-tokpanov
Copy link
Contributor Author

yury-tokpanov commented Mar 7, 2025 •
edited
Loading

seems like current failures in checks are due to cv2 imports in transformers v4.49.0. This is a known issue: #13905

Other than that things work.

@yury-tokpanov
Copy link
Contributor Author

yury-tokpanov commented Mar 11, 2025

@tlrmchlsmth other than the external issue with the latest released transformers (cv2 import in 4.49.0, but I see it's fixed in their dev branch), do you have other suggestions for this PR?

tlrmchlsmth
tlrmchlsmth reviewed Mar 13, 2025
View reviewed changes
Copy link
Member

@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.

The PR looks great to me, thanks for the contribution -- I'll accept once the transformers 4.49 issue is resolved

DarkLight1337
DarkLight1337 reviewed Mar 14, 2025
View reviewed changes
DarkLight1337
DarkLight1337 reviewed Mar 15, 2025
View reviewed changes
@yury-tokpanov
Zamba2 initial commit
1716efb 
Signed-off-by: Yury Tokpanov <[email protected]>
yury-tokpanov and others added 7 commits March 18, 2025 05:49
@yury-tokpanov
TP support, unit test, docsrtrings
c91c9c4 
Signed-off-by: Yury Tokpanov <[email protected]>
@yury-tokpanov
Fix unit tests
89dab81 
Signed-off-by: Yury Tokpanov <[email protected]>
@Quentin-Anthony @tlrmchlsmth @yury-tokpanov
label transformers req as just zamba2
39105a9 
Co-authored-by: Tyler Michael Smith <[email protected]>
Signed-off-by: Quentin Anthony <[email protected]>
@yury-tokpanov
Rebase + remove kv_cache/attn_metadata args
0867d1b 
Signed-off-by: Yury Tokpanov <[email protected]>
@yury-tokpanov
MergedColumnParallel for MLP block
f3ec9ef 
Signed-off-by: Yury Tokpanov <[email protected]>
@yury-tokpanov
Zamba2LoRA class + block indexing rework
7ad96e0 
Signed-off-by: Yury Tokpanov <[email protected]>
@yury-tokpanov @DarkLight1337
Update tests/models/registry.py
54f25dd 
Co-authored-by: Cyrus Leung <[email protected]>
Signed-off-by: Yury Tokpanov <[email protected]>
@yury-tokpanov yury-tokpanov force-pushed the zamba2 branch 2 times, most recently from 4e3b75d to 71cf8f3 Compare March 18, 2025 06:17
@yury-tokpanov
revert requirements
68ffb4a 
Signed-off-by: Yury Tokpanov <[email protected]>
@DarkLight1337
Copy link
Member

DarkLight1337 commented Mar 18, 2025

Thanks, can you also update the list of supported models in the docs with this model?

@mergify mergify bot added the documentation Improvements or additions to documentation label Mar 18, 2025
@DarkLight1337 DarkLight1337 enabled auto-merge (squash) March 18, 2025 07:24
@github-actions github-actions bot added the ready ONLY add when PR is ready to merge/full CI is needed label Mar 18, 2025
@yury-tokpanov
Copy link
Contributor Author

yury-tokpanov commented Mar 18, 2025

Thanks, can you also update the list of supported models in the docs with this model?

Yep, all done.

Btw, I see a bunch of tests failed, but upon further inspection it seems they're unrelated to the PR.

@DarkLight1337
Copy link
Member

DarkLight1337 commented Mar 18, 2025

Indeed they are unrelated - merging

@vllm-bot vllm-bot merged commit 452e8fd into vllm-project:main Mar 18, 2025
37 of 42 checks passed
@DarkLight1337 DarkLight1337 added this to the v0.8.0 milestone Mar 18, 2025
simon-mo pushed a commit that referenced this pull request Mar 18, 2025
@yury-tokpanov @Quentin-Anthony
@tlrmchlsmth @DarkLight1337 @simon-mo
[MODEL] Add support for Zamba2 models (#13185)
01ca85b 
Signed-off-by: Yury Tokpanov <[email protected]>
Signed-off-by: Quentin Anthony <[email protected]>
Co-authored-by: Quentin Anthony <[email protected]>
Co-authored-by: Tyler Michael Smith <[email protected]>
Co-authored-by: Cyrus Leung <[email protected]>
lulmer pushed a commit to lulmer/vllm that referenced this pull request Apr 7, 2025
@yury-tokpanov @Quentin-Anthony
@tlrmchlsmth @DarkLight1337 @lulmer
[MODEL] Add support for Zamba2 models (vllm-project#13185)
5393436 
Signed-off-by: Yury Tokpanov <[email protected]>
Signed-off-by: Quentin Anthony <[email protected]>
Co-authored-by: Quentin Anthony <[email protected]>
Co-authored-by: Tyler Michael Smith <[email protected]>
Co-authored-by: Cyrus Leung <[email protected]>
Signed-off-by: Louis Ulmer <[email protected]>
@ckhordiasma ckhordiasma mentioned this pull request Apr 17, 2025
Closed
@tlrmchlsmth tlrmchlsmth mentioned this pull request Apr 24, 2025
Open
9 tasks
shreyankg pushed a commit to shreyankg/vllm that referenced this pull request May 3, 2025
@yury-tokpanov @Quentin-Anthony
@tlrmchlsmth @DarkLight1337 @shreyankg
[MODEL] Add support for Zamba2 models (vllm-project#13185)
cf71ae7 
Signed-off-by: Yury Tokpanov <[email protected]>
Signed-off-by: Quentin Anthony <[email protected]>
Co-authored-by: Quentin Anthony <[email protected]>
Co-authored-by: Tyler Michael Smith <[email protected]>
Co-authored-by: Cyrus Leung <[email protected]>
RichardoMrMu pushed a commit to RichardoMrMu/vllm that referenced this pull request May 12, 2025
@yury-tokpanov @Quentin-Anthony
@tlrmchlsmth @DarkLight1337
[MODEL] Add support for Zamba2 models (vllm-project#13185)
db7a893 
Signed-off-by: Yury Tokpanov <[email protected]>
Signed-off-by: Quentin Anthony <[email protected]>
Co-authored-by: Quentin Anthony <[email protected]>
Co-authored-by: Tyler Michael Smith <[email protected]>
Co-authored-by: Cyrus Leung <[email protected]>
Signed-off-by: Mu Huai <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tlrmchlsmth tlrmchlsmth tlrmchlsmth left review comments

@DarkLight1337 DarkLight1337 DarkLight1337 approved these changes

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

+1 more reviewer

@fabianlim fabianlim fabianlim left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

@tlrmchlsmth tlrmchlsmth

Labels

ci/build documentation Improvements or additions to documentation ready ONLY add when PR is ready to merge/full CI is needed

Projects

None yet

Milestone

v0.8.0

Development

Successfully merging this pull request may close these issues.

6 participants

@yury-tokpanov @DarkLight1337 @tlrmchlsmth @fabianlim @vllm-bot @Quentin-Anthony