Add LoRA fine-tuning to AWQ

RonanKMcGovern commented 1 year ago

It would be fantastic if we could add the ability to do LoRA fine-tuning and merging of adapters.

Background on QLoRA

Interestingly, for many fine-tunings, the results of QLoRA are very similar to doing an unquantized LoRA fine-tune. Of course, QLoRA allows for fine-tuning with one third of the VRAM requirement (if doing 4bit).

The two common libraries I use are:

bitsandbytes: This does not allow correct merging of adapters to the dequantized base model.
gptq: also does not allow merging of adapters, plus the perplexity is worse than awq (and bnb in some cases).

Why add LoRA to AWQ

AWQ has the best perplexity and good inference speed
If it were possible to do QLoRA AND merge adapters to the base dequantized model, AWQ would be the best available solution for doing fine-tuning, at least in quantized form.

casper-hansen commented 1 year ago

I would love to add LoRA and make AutoAWQ compatible with PEFT. This is something that I have thought about but currently it’s more important for me to see what I can do a high throughput quantized model.

RonanKMcGovern commented 1 year ago

Ok cool, I think supporting QLoRA merging is underappreciated though. I don't know of any way to do this and it means there isn't a good open source way to serve QLoRA tuned open source models.

BTW, when you say high throughput, do you mean batch size larger than 8? so bf16 implementation?

s4rduk4r commented 12 months ago

I could probably look into it during next week. Maybe the autograd_4bit code from here could be adapted somehow

casper-hansen commented 12 months ago

In general, I think we should integrate with PEFT. From my understanding, this requires our WQLinear modules to generate gradients during a backward pass - so you would have to implement that functionality. It may turn out to be easy enough since autograd works pretty well - maybe look to AutoGPTQ to see how they integrated with PEFT.

K024 commented 12 months ago

@casper-hansen AutoGPTQ implements QLinear with various underlying QGEMM implementations (cuda, exllama, qigen, openai/triton) and most of them did not implement the backward kernel except for triton. The triton kernel is currently the only one could be used for training a quantized model in AutoGPTQ, though not the most optimal.

FYI the autograd_4bit mentioned above simply unpacks the weights into fp and calls torch.matmul

casper-hansen commented 12 months ago

I welcome any work on a backward pass function for AWQ. There are many ways to go about it. Just keep in mind the AWQ kernel does not scale well with larger batch sizes, above batch size 16 and it will be slower than FP16. I found some code where someone did the backward pass:

https://github.com/compressa-ai/llm-awq/tree/dev

K024 commented 12 months ago

@casper-hansen FYI the above one still unpacks and gemms everything in fp...

~And I noticed the pack order had changed in the llm-awq repo since Sep 7~

I see the changes gemmv2_forward_cuda vs gemm_forward_cuda

casper-hansen commented 12 months ago

Yes, I see that, they dequantize to run FP16. I’m pretty sure this is normal for training?

I created v2 based on their new GEMM kernel but it’s way slower and only compatible with GEMV where it processes the context. GEMV is 20% faster at small prompts but not great for high throughput or deployments.

Ph0rk0z commented 12 months ago

ime, triton was never faster for anything. exclusionary high compute requirements and slower speed, oh my.

The only one who has pulled off merging adapters into quantized models is GGUF. With that alpaca_lora_4bit repo + extensions I can merge LoRA together but not to the model.

cassianlewis commented 11 months ago

AFAIK you can merge the LoRA weights and unquantised base model (even if you fine-tuned in 4/8 bit) using model.merge_and_unload(). You can then quantise this model using AQW and run as normal.

I guess this only really applies if you don't have the VRAM to train the model without PEFT though.

RonanKMcGovern commented 11 months ago

@cassianlewis yeah in bnb but not gptq AFAIK.

Not ideal to merge to unquantified either.

sd3ntato commented 8 months ago

AFAIK you can merge the LoRA weights and unquantised base model (even if you fine-tuned in 4/8 bit) using model.merge_and_unload(). You can then quantise this model using AQW and run as normal.

I guess this only really applies if you don't have the VRAM to train the model without PEFT though.

Hi, I'm trying to do it with Mixtral, but i get the following output / error:

Downloading and preparing dataset json/mit-han-lab--pile-val-backup to /root/.cache/huggingface/datasets/mit-han-lab___json/mit-han-lab--pile-val-backup-39bc257d0ce73de2/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4...
 82 Downloading readme: 100%|██████████| 167/167 [00:00<00:00, 112kB/s]
 83 Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]
 84 Downloading data:  48%|████▊     | 225M/471M [00:01<00:01, 125MB/s]
 85 Extracting data files: 100%|██████████| 1/1 [00:02<00:00,  2.22s/it]
 86 AWQ:   0%|          | 0/32 [00:06<?, ?it/s]
 87 Dataset json downloaded and prepared to /root/.cache/huggingface/datasets/mit-han-lab___json/mit-han-lab--pile-val-backup-39bc257d0ce73de2/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4. Subsequent calls will reuse this data.
 88 AWQ:   0%|          | 0/32 [00:00<?, ?it/s]
 89 Traceback (most recent call last):
 90   File "/opt/ml/code/run_clm_awq.py", line 339, in <module>
 91     main()
 92   File "/opt/ml/code/run_clm_awq.py", line 324, in main
 93     training_function(run, args)
 94   File "/opt/ml/code/run_clm_awq.py", line 292, in training_function
 95     model.quantize(
 96   File "/opt/conda/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
 97     return func(*args, **kwargs)
 98   File "/opt/conda/lib/python3.10/site-packages/awq/models/base.py", line 95, in quantize
 99     self.quantizer.quantize()
100   File "/opt/conda/lib/python3.10/site-packages/awq/quantize/quantizer.py", line 107, in quantize
101     module_config: List[Dict] = self.awq_model.get_layers_for_scaling(
102   File "/opt/conda/lib/python3.10/site-packages/awq/models/mixtral.py", line 46, in get_layers_for_scaling
103     inp=input_feat['self_attn.q_proj'],
104 KeyError: 'self_attn.q_proj'

could anyone please help me out with this?

RonanKMcGovern commented 8 months ago

If you merge a quantized (transformers) model then it will become a 4 or 8 bit model, which you can't then do AWQ on.

Instead, you would need to reload a base model in 16 bit and merge your LoRA to that (using merge and unload). Then you can AWQ that merged model. More info in this vid.

Ph0rk0z commented 8 months ago

Having the base model becomes unmanageable with 70b+, that's part of the issue. They're 160gb+

RanchiZhao commented 5 months ago

Hi! any progress? is train LoRA modules with AWQ available now?

RMimo commented 5 months ago

Hi, I'm also interested to know whether LoRA + AWQ is already available now. Thanks!

RanchiZhao commented 5 months ago

Hi, I'm also interested to know whether LoRA + AWQ is already available now. Thanks!

@RicardoHalak see this, is runnable https://github.com/huggingface/transformers/pull/28987

casper-hansen / AutoAWQ

Add LoRA fine-tuning to AWQ #85