Exllama kernels support for AWQ models

[IlyasMoutawwakil]

What does this PR do?

Following casper-hansen/AutoAWQ#313
ExllamaV2 offers up to 2x speedup compared to GEMM, while also compatible with AMD ROCm.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@SunMarc and @younesbelkada

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

Thanks for making exllama kernel compatible with AWQ models ! This will make AWQ so much faster ! I've left a few minor comments.

[IlyasMoutawwakil]

I guess all points are addressed.
@casper-hansen when is 0.1.9 planned ?

[younesbelkada]

Thanks a lot for adding the ex-llama v2 support ! 🔥
Let's add autoawq==0.1.9 in the Dockerfile:

transformers/docker/transformers-all-latest-gpu/Dockerfile

Line 59 in bbe30c6

RUN python3 -m pip install --no-cache-dir https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.8/autoawq-0.1.8+cu118-cp38-cp38-linux_x86_64.whl

- @casperhansen could you confirm the 0.1.9 is planned sometime soon?
Do you know also if this feature is supported on NVIDIA T4 GPUs? If that's the case can you add a simple generation test in: https://github.com/huggingface/transformers/blob/main/tests/quantization/autoawq/test_awq.py
Thanks !

[casper-hansen]

@younesbelkada The next release will be 0.2.0 🤗. For T4 support, I have not tested it. If AutoGPTQ supports T4 with ExLlama v1 and v2 kernels, AutoAWQ should too as the kernels are the same.

EDIT: To answer the timeline question. There is no set-in-stone plan for the next release. PRs to be merged before release include AMD support, Marlin support, Qwen2 support, and hopefully PEFT support. I expect this could be done in <1-2 weeks.

[younesbelkada]

Awesome! Per my understanding ex-llama + AutoGPTQ should be supported on T4 so it should be all good !
Let me know whenever you have some progress for the PEFT support so that I'll dive in to add AWQ + PEFT support directly in PEFT

[younesbelkada]

@IlyasMoutawwakil - #26610 being merged would you be happy to transfer the logic inside transformers/src/quantizers/quantizer_awq.py's post-processing method?

[younesbelkada]

Thanks very much @IlyasMoutawwakil ! I left one suggestion - what do you think?

[younesbelkada]

couldn't we make max_input_len configurable through AwqConfig - wdyt?

[IlyasMoutawwakil]

Marc suggested we leave it as is for now #28634 (comment)

[younesbelkada]

oh okay ! I think it would makes sense to directly expose a exllama_config I think - wdyt @SunMarc ?

[SunMarc]

Yes, it would make more sense to expose it in a exllama_config !

[IlyasMoutawwakil]

I guess in another PR right ?

[younesbelkada]

Hmmm I think it should be better to add it now and not leave the main branch with hardcoded config values, it shouldn't be super complex as you can just copy over the existing logic in GptqConfig right?

[ArthurZucker]

Looks already really nice thanks to the integration refactor!

[younesbelkada]

Thanks, clean ! Let's merge this PR right after the next release of autoawq
@casper-hansen do you have any ETA for the next release?

[younesbelkada]

Thanks again !

[SunMarc]

I see that @casper-hansen already made the release of awq with exllama kernel. Can you check that everything works fine with the latest release @IlyasMoutawwakil ? Then, we are good to merge !

[IlyasMoutawwakil]

@SunMarc on it!

[IlyasMoutawwakil]

works on rocm5.6 with torch 2.2

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AwqConfig

quantization_config = AwqConfig(version="exllama")
model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Mistral-7B-Instruct-v0.1-AWQ",
    quantization_config=quantization_config,
    device_map="auto" or torch.device("cuda"),
)

input_ids = torch.randint(0, 100, (1, 128), dtype=torch.long, device="cuda")
output = model(input_ids)
print(output.logits)

tokenizer = AutoTokenizer.from_pretrained("TheBloke/Mistral-7B-Instruct-v0.1-AWQ")
input_ids = tokenizer.encode("How to make a cake", return_tensors="pt").to(model.device)
output = model.generate(input_ids, do_sample=True, max_length=50, pad_token_id=50256)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The device_map is mandatory since exllamav2_post_init scratch space tensors allocation needs to check which qweights are on which cuda devices. I believe this is a requirement in GPTQ as well.

[younesbelkada]

Thanks LGTM with one nit !

[ArthurZucker]

Thank you for adding exllama support 🔥