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] 1.58bits BitNet Model Support #7725

Open
wants to merge 27 commits into
base: main
Choose a base branch
from

Conversation

LeiWang1999
Copy link

@LeiWang1999 LeiWang1999 commented Aug 21, 2024

This pull request is a follow-up to PR #6036. In this PR, we introduce the BitNet model and provide an efficient inference kernel with the BitBLAS backend. Here are the performance benchmarks:

Model Framework BS16IN32OUT128 BS1IN512OUT1024 B32IN32OUT128
BitNet-3B-1.58bits PyTorch 106.83 49.34 209.03
BitNet-3B-1.58bits PyTorch-BitBLAS 240.33 103.09 493.31
BitNet-3B-1.58bits vLLM-BitBLAS 379.25 117.43 752.55
BitNet-3B-1.58bits vLLM-BitBLAS-CUDA-Graph 2543.58 1621.08 2731.79

To answer the question raised by @mgoin in PR #6036, I believe a new BitNet model is necessary because the open-source BitNet implementation provides a unique tokenizer and model architecture, which includes an additional RMS layer compared to LLaMA. Additionally, the BitNet integration example with llama.cpp also introduces a new model architecture (refer to: llama.cpp.pr.7931).

Example Usage:

from conftest import VllmRunner

# Test BitNET model with BitBLAS quantization
with VllmRunner(
    "hxbgsyxh/bitnet_b1_58-3B",
    dtype="half",
    quantization="bitnet_bitblas",
    enforce_eager=True,
    gpu_memory_utilization=0.5,
) as bitnet_model:
    bitbnet_outputs = bitnet_model.generate_greedy(
        ["Hi, tell me about Microsoft?"], max_tokens=128
    )
    print("bitnet_bitblas:")
    print(bitbnet_outputs[0][0])
    print(bitbnet_outputs[0][1])

# Test another BitBLAS model
with VllmRunner(
    "hxbgsyxh/bitnet_b1_58-3B_bitblas",
    dtype="half",
    quantization="bitblas",
    enforce_eager=True,
) as bitnet_model:
    bitbnet_outputs = bitnet_model.generate_greedy(
        ["Hi, tell me about Microsoft?"], max_tokens=128
    )
    print("bitblas:")
    print(bitbnet_outputs[0][0])
    print(bitbnet_outputs[0][1])

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 consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

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

🚀

@DarkLight1337
Copy link
Member

Sorry for the long delay, @mgoin can you follow up on this and the previous PR?

A quick heads-up that the new locations of the model tests have been adjusted in #7820, so please merge from main.

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