Loading quantized models

[abhinavkulkarni]

Hi,

Is there a way to load quantized models using vLLM? For e.g. I have been using AWQ quantization and have released a few models here.

The model loading process looks like the following:

1. Model is first initialized with empty weights
2. Linear layers are replaced by a custom linear layer that supports zero-point quantization
3. Weights are loaded from a checkpoint onto a GPU

Please note the matrix multiplication inside the linear layer is done by a Python extension that uses custom CUDA kernels since AWQ uses 4-bit quantization. Everything else - the attention mechanism, etc. is the same.

The model otherwise supports all the HuggingFace AutoModelForCausalLM APIs.

Thanks!

[WoosukKwon]

Hi @abhinavkulkarni, thanks for exploring vLLM and requesting the feature! For now, vLLM doesn't support quantized models. But as you mentioned, we believe it is not very difficult to add it into vLLM. We will definitely look into it after finishing other urgent issues (e.g., Falcon support, and bug fixes). Besides, we are looking forward to contributions from community! Please make a PR if you have time to implement it.

[johnnysmithy123467]

Hello, I am considering adding support for 4 / 8 bit quantization into VLLM. What specifically would have to be done to implement this feature? Can I emulate the on the fly style weight quantization done in the fast chat repo?: https://github.com/lm-sys/FastChat/blob/main/fastchat/model/compression.py

[abhinavkulkarni]

@johnnysmithy123467: It would be best if the models are loaded upfront before they are passed to vLLM. Usually, quantized models modify the structure of the base model by replacing Linear and LayerNorm layers with custom ones that support zero-point quantization. They also use custom CUDA extensions to do 4-bit matrix multiplication.

You can see how GPTQ quantized models are loaded here and AWQ quantized models here.

[TheBloke]

Given bitsandbytes recent work on 4bit performance - now claiming up to 4.2x faster performance from 4bit versus fp16 - adding bitsandbytes would seem like an obvious, and perhaps easiest, first quantization format to support?

[johnnysmithy123467]

@TheBloke what would it take to add bits and bytes support? considering implementing it myself and just want to know how I would go about doing it, seems the weight loading schemes are very customized on vllm

[casper-hansen]

Given bitsandbytes recent work on 4bit performance - now claiming up to 4.2x faster performance from 4bit versus fp16 - adding bitsandbytes would seem like an obvious, and perhaps easiest, first quantization format to support?

I like the theoretical speedup but the problem is they do not support serialization/deserialization of the 4-bit models. Until that is resolved, it will not be great for production usage.

[creatorrr]

@zhuohan123 support for bnb-nf4 would be amazing. Especially for larger models. If you can point me in the right direction, I’d be happy to take a crack at implementing it

[casper-hansen]

-1 for BNB and +1 for AWQ.

I will just let my test results speak for themselves. This is for the model mosaicml/mpt-7b-8k-chat. EDIT: LLaMa-2 7B also runs at 103 tokens/s with AWQ TinyChat on a 4090 according to my own testing.

System: 1 x RTX A6000, 4 vCPU, 61 GB RAM (RunPod Community Cloud)
BNB (NF4): 7.5 tokens/s (load time: 16.41 s)
AWQ (W4A16): 28 tokens/s (load time: 14.27 s)
AWQ tinychat (W4A16): 46.86 tokens/s (load time: 22.29s)

This is torch 2.0.1, bitsandbytes 0.41.0, transformers 4.31.0, and AWQ compiled from main.

CC @TheBloke and @WoosukKwon

[casper-hansen]

117 tokens/s now on 4090 with MPT-7B model.

https://huggingface.co/casperhansen/mpt-7b-8k-chat-awq

[ri938]

Was thinking of maybe working on this myself.

My understanding is that I only need to replace the linear layers. And I dont need to touch the attention code. Its a lot to simpler to implement if I dont need to get into the attention code with its custom C++ kernels and paged attention. Can someone confirm my understanding there?

I was going to use GPTQ and use text-generation-inference library as a template for how to implement this.

[johnnysmithy123467]

@ri938 That sounds right. Let us know how it goes, eager to hear from you. at least someone's taking a crack at it!

[ri938]

Update on this: I have a hacky proof of concept quantisation working which reduces memory use and inference quality looks high quality.

However inference speed is much slower. Even with a single batch size. I tried the default and exllama kernels from TGI repo but both were about 60% the inference speed as with fp16 models. Its too slow for my particular use case and I think its important to get it to comparable inference speeds as fp16.

One option is I could open a PR for quantisation and then allow others with more experience in CUDA optimisations to then contribute to this. Otherwise I probably will continue to work on this.

[casper-hansen]

Update on this: I have a hacky proof of concept quantisation working which reduces memory use and inference quality looks high quality.

However inference speed is much slower. Even with a single batch size. I tried the default and exllama kernels from TGI repo but both were about 60% the inference speed as with fp16 models. Its too slow for my particular use case and I think its important to get it to comparable inference speeds as fp16.

One option is I could open a PR for quantisation and then allow others with more experience in CUDA optimisations to then contribute to this. Otherwise I probably will continue to work on this.

GPTQ is pretty slow in general unless you build hyper optimized kernels on top. Did you try AWQ, might give you better results?

[ri938]

Not tried AWQ. But for GPTQ I did use the kernels from text-generation-inference (including a port of exllama) so I thought this should be good performance.