[Kernel] Support Microsoft Runtime Kernel Lib for our Low Precision Computation #6036

LeiWang1999 · 2024-07-01T16:30:04Z

Hi all, this PR introduces support for the Microsoft Runtime Kernel Library to enhance our low precision computation capabilities.

Brief Introduction of BitBLAS

BitBLAS is a library to support mixed-precision BLAS operations on GPUs, for example, the $W_{wdtype}A_{adtype}$ mixed-precision matrix multiplication where $C_{cdtype}[M, N] = A_{adtype}[M, K] \times W_{wdtype}[N, K]$.
BitBLAS aims to support efficient mixed-precision DNN model deployment, especially the $W_{wdtype}A_{adtype}$ quantization in large language models (LLMs), for example, the $W_{UINT4}A_{FP16}$ in GPTQ, the $W_{INT2}A_{FP16}$ in BitDistiller, the $W_{INT2}A_{INT8}$ in BitNet-b1.58.

PR Overview

This PR integrates BitBLAS into vLLM by adding examples of its usage. We provide two forms:

Load from GPTQ Checkpoints: This allows the loading of models from GPTQ format checkpoints.
Load from GPTQ CKPT with BitBLAS Format: This enables the loading of models using the BitBLAS format for further optimized performance.

Below are the benchmarking results that we evaluated several months ago:

TODO ITEMS

Update and provide the latest benchmarking results.
1.58Bits Model
Provide Benchmark/Test Scripts

Any feedback and suggestions to improve this integration are appreciated.

…nearMethod constructor

robertgshaw2-neuralmagic · 2024-07-01T16:39:27Z

Nice!

LeiWang1999 · 2024-07-01T16:44:51Z

BTW, are there any tools available that can automatically resolve lint issues?

vllm/model_executor/layers/quantization/gptq_bitblas.py:28:1: E402 Module level import not at top of file
vllm/model_executor/layers/quantization/gptq_bitblas.py:28:8: F811 Redefinition of unused `bitblas` from line 21
vllm/model_executor/layers/quantization/gptq_bitblas.py:29:1: E402 Module level import not at top of file
vllm/model_executor/layers/quantization/gptq_bitblas.py:66:81: E501 Line too long (107 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:172:81: E501 Line too long (85 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:222:81: E501 Line too long (105 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:230:81: E501 Line too long (89 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:233:81: E501 Line too long (110 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:236:81: E501 Line too long (99 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:242:81: E501 Line too long (84 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:253:81: E501 Line too long (94 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:414:81: E501 Line too long (86 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:417:29: G004 Logging statement uses f-string
vllm/model_executor/layers/quantization/gptq_bitblas.py:420:17: G004 Logging statement uses f-string
vllm/model_executor/layers/quantization/gptq_bitblas.py:427:81: E501 Line too long (103 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:433:81: E501 Line too long (116 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:454:81: E501 Line too long (82 > 80)

robertgshaw2-neuralmagic · 2024-07-01T16:48:10Z

BTW, are there any tools available that can automatically resolve lint issues?

vllm/model_executor/layers/quantization/gptq_bitblas.py:28:1: E402 Module level import not at top of file
vllm/model_executor/layers/quantization/gptq_bitblas.py:28:8: F811 Redefinition of unused `bitblas` from line 21
vllm/model_executor/layers/quantization/gptq_bitblas.py:29:1: E402 Module level import not at top of file
vllm/model_executor/layers/quantization/gptq_bitblas.py:66:81: E501 Line too long (107 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:172:81: E501 Line too long (85 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:222:81: E501 Line too long (105 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:230:81: E501 Line too long (89 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:233:81: E501 Line too long (110 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:236:81: E501 Line too long (99 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:242:81: E501 Line too long (84 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:253:81: E501 Line too long (94 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:414:81: E501 Line too long (86 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:417:29: G004 Logging statement uses f-string
vllm/model_executor/layers/quantization/gptq_bitblas.py:420:17: G004 Logging statement uses f-string
vllm/model_executor/layers/quantization/gptq_bitblas.py:427:81: E501 Line too long (103 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:433:81: E501 Line too long (116 > 80)
vllm/model_executor/layers/quantization/gptq_bitblas.py:454:81: E501 Line too long (82 > 80)

./format.sh fixes whatever it can, but not everything is automated for fixing (esp line length)

mgoin · 2024-07-01T16:50:08Z

@LeiWang1999 thanks for the WIP, very cool interface with bitblas as a package. Can you explain if the GPTQ benchmarking results in vLLM were run with the base "gptq" kernels or using the "gptq_marlin" interface to take advantage of Marlin kernels? This will be important to compare with the current baseline we consider for GPTQ models in vLLM

LeiWang1999 · 2024-07-01T17:01:40Z

@LeiWang1999 thanks for the WIP, very cool interface with bitblas as a package. Can you explain if the GPTQ benchmarking results in vLLM were run with the base "gptq" kernels or using the "gptq_marlin" interface to take advantage of Marlin kernels? This will be important to compare with the current baseline we consider for GPTQ models in vLLM

Thanks, it utilized exllamav2 during our benchmarking at that time; we will examine the comparison with the Marlin kernel.

…las-intg

LeiWang1999 · 2024-07-19T04:22:02Z

Hi all, I recently update the the supports for 1.58bits model and related bitblas inference kernel for vllm.

		Token Per Second(tok/s)
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

LeiWang1999 · 2024-07-19T04:27:05Z

We will soon do benchmarking with marlin, and looks like the docs build failed because of the dependency for bitblas, do you have any ideas to fix this issue? should we put the bitblas requirements to the doc/requirements or is there some options to skip this dependency? @mgoin

…las-intg

LeiWang1999 · 2024-08-20T18:50:38Z

I think this PR is ready for review. Here is a summary of this update:

We now support BitBLAS as a quantized backend and can use vLLM to serve pretrained models from Hugging Face (in GPTQ, BitNet, or BitBLAS format) with the BitBLAS inference kernel.

We briefly tested the performance using Marlin with the throughput benchmark scripts provided by vLLM on A100:

python benchmark_throughput.py --backend vllm --num-prompts 1 --input-len 32 --output-len 512 --max-model-len 1024 --model "hxbgsyxh/llama-13b-4bit-g-1-bitblas" --quantization "bitblas" 

python benchmark_throughput.py --backend vllm --num-prompts 1 --input-len 32 --output-len 512 --max-model-len 1024 --model "hxbgsyxh/llama-13b-4bit-g-1-marlin" --quantization "marlin"

The performance results are:

Marlin: 122.67 toks/s
BitBLAS: 127.11 toks/s

Some notes:

Marlin requires a workspace for spin lock to perform global reduction, while BitBLAS doesn’t require it.
BitBLAS supports more complex cases compared to Marlin, such as sym=False or 2-bit formats.

Moreover, this PR also adds support for the 1.58-bit BitNET model.

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

All correctness checks have been evaluated with the following:

from conftest import VllmRunner
import torch

# 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:
    torch.cuda.profiler.start()
    bitbnet_outputs = bitnet_model.generate_greedy(
        ["Hi, tell me about Microsoft?"], max_tokens=128
    )
    torch.cuda.profiler.stop()
    print("bitblas:")
    print(bitbnet_outputs[0][0])
    print(bitbnet_outputs[0][1])

# Test GPTQ quantized model
with VllmRunner(
    "hxbgsyxh/opt-125m-4bit-128g",
    dtype="half",
    quantization="gptq",
    enforce_eager=True,
) as marlin_model:
    torch.cuda.profiler.start()
    bitbnet_outputs = marlin_model.generate_greedy(
        ["Hi, tell me about Microsoft?"], max_tokens=128
    )
    torch.cuda.profiler.stop()
    print("bitblas:")
    print(bitbnet_outputs[0][0])
    print(bitbnet_outputs[0][1])

torch.compiler.reset()

# Test GPTQ quantized model with BitBLAS
with VllmRunner(
    "hxbgsyxh/opt-125m-4bit-128g-bitblas",
    dtype="half",
    quantization="bitblas",
    enforce_eager=True,
) as bitblas_model:
    torch.cuda.profiler.start()
    bitbnet_outputs = bitblas_model.generate_greedy(
        ["Hi, tell me about Microsoft?"], max_tokens=128
    )
    torch.cuda.profiler.stop()
    print("bitblas:")
    print(bitbnet_outputs[0][0])
    print(bitbnet_outputs[0][1])

LeiWang1999 · 2024-08-20T19:16:28Z

any questions are welcome and please take a review when you have a moment :) @mgoin @robertgshaw2-neuralmagic

mgoin · 2024-08-20T20:44:45Z

Thanks for all the work @LeiWang1999! I have a few high-level thoughts first on how to make landing this more straightforward:

Make bitblas an optional dependency and remove from requirements-common.txt. The pypi package is ~90MB, possibly is built for a specific version of PyTorch/CUDA, and seems to include a lot of deps (TVM?). I think it is hard to require, especially for non-CUDA devices. See bitsandbytes or deepspeedfp for an example of how we usually implement lazy import with an exception message to install.
Add support for bitnet in another PR. I think it is worth looking at separately and understanding the pros/cons. I find it a bit surprising that it requires implementing a whole new model and tokenizer class.
gptq_bitblas seems a bit redundant without further benchmarking separating it from gptq_marlin. I agree it is certainly useful where marlin doesn't have support. Also we want to eventually perform a refactor to separate checkpoint formats from kernel implementations, so we will need to revisit this soon.

LeiWang1999 · 2024-08-21T06:46:45Z

Thanks for your suggestions, @mgoin!

For the first suggestion, we’ve restructured the bitblas import to be a lazy import. Additionally, while bitblas still requires a CUDA environment, it is now linked to libcuda.so instead of being tied to a specific CUDA version :).
Regarding he second suggestion, we’ve removed the bitnet-related items in this PR. Let’s continue that discussion in a separate pr.
For the third suggestion, I believe keeping gptq_bitblas is still valuable for formats that Marlin doesn’t support, such as those with dynamic zero points or lower precision. It allows us to directly repack the bitblas format from a gptq checkpoint.

with VllmRunner(
    "hxbgsyxh/llama-13b-4bit-g-1", # model with gptq format
    dtype="half",
    quantization="bitblas",
    enforce_eager=True,
) as bitblas_model:
    torch.cuda.profiler.start()
    bitbnet_outputs = bitblas_model.generate_greedy(
        ["Hi, tell me about microsoft?"], max_tokens=128
    )
    torch.cuda.profiler.stop()
    print("bitblas:")
    print(bitbnet_outputs[0][0])
    print(bitbnet_outputs[0][1])

mgoin

Thanks for splitting it up! I left a first round of clear nits/issues and will do a more in-depth pass later. There seem to be a lot of various formatting changes, for some reason

docs/source/quantization/bitblas.rst

vllm/config.py

vllm/model_executor/layers/quantization/bitblas.py

Co-authored-by: Michael Goin <[email protected]>

…tblas into bitblas-intg

…las-intg

LeiWang1999 · 2024-08-26T03:56:42Z

Hi @mgoin , are there any further updates or actions we should take?

mgoin

Hi @LeiWang1999 I'm very sorry for the delay, I lost track of this PR and didn't catch your ping.

There has been an ongoing refactor for quantization methods to use a new set of vLLMParameters (see gptq_marlin PR #7281) to simplify weight loading, but we could delay this for bitblas to make it easier to land this initial PR.

Also as mentioned in #7725 (comment), there will be a few merge conflicts with main.

If/when you have bandwidth to finish this out, I promise to get this over the line asap. Please let me know!

mgoin · 2024-09-13T17:42:01Z

vllm/model_executor/layers/quantization/gptq_bitblas.py

+        if layer.bitblas_state == GPTQBitBLASState.REPACK:
+            layer.bitblas_state = GPTQBitBLASState.READY
+
+            # Newly generated tensors need to replace existing tensors that are
+            # already registered as parameters by vLLM (and won't be freed)
+            def replace_tensor(name, new_t):
+                # It is important to use copy_() here since it ensures
+                # the same buffer is reused
+                getattr(layer, name).copy_(
+                    new_t.view(getattr(layer, name).dtype).view(
+                        getattr(layer, name).shape))
+                del new_t
+
+            # Repack weights
+            bitblas_qweight, bitblas_scales, bitblas_qzeros = (
+                self.repack_bitblas_from_gptq(
+                    layer.qweight,
+                    layer.scales,
+                    layer.qzeros,
+                ))
+            replace_tensor("qweight", bitblas_qweight)
+            replace_tensor("scales", bitblas_scales)
+            replace_tensor("qzeros", bitblas_qzeros)


It would be best to move this into a process_weights_after_loading function we have specifically for this purpose, example in gptq_marlin.py

vllm/vllm/model_executor/layers/quantization/gptq_marlin.py

Line 287 in a84e598

def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

Thanks, I'll take a look. I'm currently working on the stream-k template in bitblas :)

LeiWang1999 added 2 commits July 1, 2024 15:37

Support Repack from GPTQ.

2be6218

chore: Remove unused input_size and output_size variables in MarlinLi…

b92de92

…nearMethod constructor

LeiWang1999 added 7 commits July 16, 2024 03:19

Support BitNet Model for 1.58bits.

71ea469

Lint Fix

dfa6b2f

lint fix

8d2c635

Lint Fix for line length

41bb18e

Support Loading 1.58B Model with BitBLAS Format

29ac34d

Improve performance for bitnet

7f69aef

Merge branch 'main' of https://github.com/vllm-project/vllm into bitb…

01a789a

…las-intg

LeiWang1999 marked this pull request as ready for review July 19, 2024 04:23

LeiWang1999 added 13 commits July 19, 2024 06:00

fix lm_head for gptq model refactor

a973123

linx fix

aea1f4c

handle compressed scale weight.

17128d5

lint fix

1741ed4

remove partial weight load for sw

726a1f7

apply torch compile for uncompressed weight.

68c8052

Merge branch 'main' of https://github.com/vllm-project/vllm into bitb…

6eb2870

…las-intg

merge bug fix

52418ef

lint fix

a15ba12

fix torch compile issue

53babae

bug fix.

40a4e53

BENCHMARK SCRIPTS

d316a87

Merge branch 'main' of https://github.com/vllm-project/vllm into bitb…

4d40275

…las-intg

LeiWang1999 added 3 commits August 20, 2024 17:21

install bitblas by default to pass the doc gen.

8e1a7e8

hide the bitblas import

7fbbccf

import fix

c487e69

LeiWang1999 added 9 commits August 21, 2024 04:43

remove all bitnet related items

0d364e7

remove bitnet

d8e9448

format fix

112ef48

remove commen requirements

7a88731

remove dep

1fed839

lint fix

4c4cf59

remove red

68d8ff1

optimize tuning message

c4e818c

update

546553f

LeiWang1999 mentioned this pull request Aug 21, 2024

[Model] 1.58bits BitNet Model Support #7725

Open

1 task

mgoin reviewed Aug 21, 2024

View reviewed changes

LeiWang1999 and others added 10 commits August 22, 2024 13:38

Update docs/source/quantization/bitblas.rst

e74368a

Co-authored-by: Michael Goin <[email protected]>

Update docs/source/quantization/bitblas.rst

051e11a

Co-authored-by: Michael Goin <[email protected]>

Update docs/source/quantization/bitblas.rst

7f6708a

Co-authored-by: Michael Goin <[email protected]>

Update docs/source/quantization/bitblas.rst

e3c8159

Co-authored-by: Michael Goin <[email protected]>

Merge branch 'bitblas-intg' of https://github.com/LeiWang1999/vllm-bi…

b335ea8

…tblas into bitblas-intg

typo fix

5b8cc8c

support bfloat16

6308db5

update recommend installing command

763c71e

Merge branch 'main' of https://github.com/vllm-project/vllm into bitb…

2ad2a14

…las-intg

lint fix

31a7bbd

mgoin reviewed Sep 13, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Kernel] Support Microsoft Runtime Kernel Lib for our Low Precision Computation #6036

[Kernel] Support Microsoft Runtime Kernel Lib for our Low Precision Computation #6036

LeiWang1999 commented Jul 1, 2024 •

edited

Loading

robertgshaw2-neuralmagic commented Jul 1, 2024

LeiWang1999 commented Jul 1, 2024 •

edited

Loading

robertgshaw2-neuralmagic commented Jul 1, 2024

mgoin commented Jul 1, 2024

LeiWang1999 commented Jul 1, 2024

LeiWang1999 commented Jul 19, 2024 •

edited

Loading

LeiWang1999 commented Jul 19, 2024

LeiWang1999 commented Aug 20, 2024 •

edited

Loading

LeiWang1999 commented Aug 20, 2024

mgoin commented Aug 20, 2024 •

edited

Loading

LeiWang1999 commented Aug 21, 2024

mgoin left a comment

LeiWang1999 commented Aug 26, 2024

mgoin left a comment

mgoin Sep 13, 2024

LeiWang1999 Sep 13, 2024

[Kernel] Support Microsoft Runtime Kernel Lib for our Low Precision Computation #6036

Are you sure you want to change the base?

[Kernel] Support Microsoft Runtime Kernel Lib for our Low Precision Computation #6036

Conversation

LeiWang1999 commented Jul 1, 2024 • edited Loading

Brief Introduction of BitBLAS

PR Overview

TODO ITEMS

robertgshaw2-neuralmagic commented Jul 1, 2024

LeiWang1999 commented Jul 1, 2024 • edited Loading

robertgshaw2-neuralmagic commented Jul 1, 2024

mgoin commented Jul 1, 2024

LeiWang1999 commented Jul 1, 2024

LeiWang1999 commented Jul 19, 2024 • edited Loading

LeiWang1999 commented Jul 19, 2024

LeiWang1999 commented Aug 20, 2024 • edited Loading

LeiWang1999 commented Aug 20, 2024

mgoin commented Aug 20, 2024 • edited Loading

LeiWang1999 commented Aug 21, 2024

mgoin left a comment

Choose a reason for hiding this comment

LeiWang1999 commented Aug 26, 2024

mgoin left a comment

Choose a reason for hiding this comment

mgoin Sep 13, 2024

Choose a reason for hiding this comment

LeiWang1999 Sep 13, 2024

Choose a reason for hiding this comment

LeiWang1999 commented Jul 1, 2024 •

edited

Loading

LeiWang1999 commented Jul 1, 2024 •

edited

Loading

LeiWang1999 commented Jul 19, 2024 •

edited

Loading

LeiWang1999 commented Aug 20, 2024 •

edited

Loading

mgoin commented Aug 20, 2024 •

edited

Loading