Add support for loading GPTQ models on CPU
Author: vivekkhandelwal1
Pull request: yes

Right now, we can only load the GPTQ Quantized model on the CUDA device. The flag `load_gptq_on_cpu` adds the support to load the GPTQ models on the CPU. The larger variants of the model are hard to load/run/trace on the GPU and that's the rationale behind adding this flag.

Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>

Comments:
@SunMarc @younesbelkada Please take a look at this PR.

>
> Can you share a small snippet you used to test out your implementation?

Hi @younesbelkada, here's a code snippet for the `Falcon-180B-Chat-GPTQ` model. Right now, I'm working on this and I get it running through this.
```
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model_name_or_path = "TheBloke/Falcon-180B-Chat-GPTQ"

model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
                                             device_map="auto",
                                             revision="main",
                                             load_gptq_on_cpu=True)
model = model.to(torch.float32)

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)

prompt = "Tell me about AI"
prompt_template=f'''User: {prompt}
Assistant: '''

print("\n\n*** Generate:")

input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids
output = model.generate(inputs=input_ids, do_sample=True, temperature=0.7, max_new_tokens=512)
print(tokenizer.decode(output[0]))

print("*** Pipeline:")
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    temperature=0.7,
    do_sample=True,
    top_p=0.95,
    repetition_penalty=1.15
)

print(pipe(prompt_template)[0]['generated_text'])
```

> Thanks for the great contribution, I did not managed to succesfully run a GPTQ model using your branch:
> 
> ```python
> import torch
> from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig
> 
> model_id = "ybelkada/opt-125m-gptq-4bit"
> 
> quantization_config = GPTQConfig(
>     bits=4,
>     disable_exllama=True
> )
> 
> model = AutoModelForCausalLM.from_pretrained(model_id, load_gptq_on_cpu=True, quantization_config=quantization_config, torch_dtype=torch.float32)
> tokenizer = AutoTokenizer.from_pretrained(model_id)
> 
> text = "Hello my name is"
> input_ids = tokenizer(text, return_tensors="pt").input_ids
> 
> out = model.generate(input_ids)
> print(tokenizer.decode(out[0], skip_special_tokens=True))
> ```
> 
> Can you share a small snippet you used to test out your implementation?

You need to add `model = model.to(torch.float32)` before doing the inference. Additionally, you would require these changes: https://github.com/PanQiWei/AutoGPTQ/pull/367

@younesbelkada, here's a smaller falcon GPTQ variant loading and executing fine, with the changes suggested above:
```
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, GPTQConfig
import torch

model_name_or_path = "TheBloke/falcon-7b-instruct-GPTQ"
quantization_config = GPTQConfig(
    bits=4,
    disable_exllama=True
)

model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
                                             device_map="auto",
                                             trust_remote_code=True,
                                             revision="main",
                                             quantization_config=quantization_config,
                                             load_gptq_on_cpu=True)
model = model.to(torch.float32)

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)

prompt = "Tell me about AI"
prompt_template=f'''User: {prompt}
Assistant: '''

print("\n\n*** Generate:")

input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids
output = model.generate(inputs=input_ids, do_sample=True, temperature=0.7, max_new_tokens=512)
print(tokenizer.decode(output[0]))

print("*** Pipeline:")
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    temperature=0.7,
    do_sample=True,
    top_p=0.95,
    repetition_penalty=1.15
)

print(pipe(prompt_template)[0]['generated_text'])
```

OK thanks, I'll run some tests and report back here

> OK thanks, I'll run some tests and report back here

Hi @younesbelkada, did you get a chance to look at this?

Hi @vivekkhandelwal1 
I did not managed to run an inference with your PR on CPU, I have used another model since I don't have access to an instance that can support falcon-180B

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig

checkpoint = "marcsun13/opt-350m-gptq-4bit"
device = "cpu" # for GPU usage or "cpu" for CPU usage

tokenizer = AutoTokenizer.from_pretrained(checkpoint)

quantization_config = GPTQConfig(bits=4, disable_exllama=True)

model = AutoModelForCausalLM.from_pretrained(checkpoint, low_cpu_mem_usage=True, quantization_config=quantization_config, torch_dtype=torch.bfloat16).to(device)

inputs = tokenizer.encode("Hello how are you?", return_tensors="pt").to(device)
outputs = model.generate(inputs, max_new_tokens=4, do_sample=False)
print(tokenizer.decode(outputs[0]))
```

I get:

```bash
    out = torch.matmul(x.half(), weight)
RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
``` 

Can you let me know if you manage to run the snippet above?

> Hi @vivekkhandelwal1 I did not managed to run an inference with your PR on CPU, I have used another model since I don't have access to an instance that can support falcon-180B
> 
> ```python
> import torch
> from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig
> 
> checkpoint = "marcsun13/opt-350m-gptq-4bit"
> device = "cpu" # for GPU usage or "cpu" for CPU usage
> 
> tokenizer = AutoTokenizer.from_pretrained(checkpoint)
> 
> quantization_config = GPTQConfig(bits=4, disable_exllama=True)
> 
> model = AutoModelForCausalLM.from_pretrained(checkpoint, low_cpu_mem_usage=True, quantization_config=quantization_config, torch_dtype=torch.bfloat16).to(device)
> 
> inputs = tokenizer.encode("Hello how are you?", return_tensors="pt").to(device)
> outputs = model.generate(inputs, max_new_tokens=4, do_sample=False)
> print(tokenizer.decode(outputs[0]))
> ```
> 
> I get:
> 
> ```shell
>     out = torch.matmul(x.half(), weight)
> RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
> ```
> 
> Can you let me know if you manage to run the snippet above?

Hi @younesbelkada, I think you missed my comments above. Actually, to run the model on the CPU, you also need to convert the model to float. After loading the model, you have to do `model = model.to(torch.float32)`. After doing this, you should be able to run the model.

Thanks @vivekkhandelwal1 , indeed I forgot to cast it to fp32, however running this script:

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig

checkpoint = "marcsun13/opt-350m-gptq-4bit"
device = "cpu" # for GPU usage or "cpu" for CPU usage

tokenizer = AutoTokenizer.from_pretrained(checkpoint)

quantization_config = GPTQConfig(bits=4, disable_exllama=True)

model = AutoModelForCausalLM.from_pretrained(checkpoint, low_cpu_mem_usage=True, quantization_config=quantization_config).to(device)

model = model.to(torch.float32)

inputs = tokenizer.encode("Hello how are you?", return_tensors="pt").to(device)
outputs = model.generate(inputs, max_new_tokens=4, do_sample=False)
print(tokenizer.decode(outputs[0]))
```

Led to the same error, what version of optimum and `auto-gptq` are you using?

> Thanks @vivekkhandelwal1 , indeed I forgot to cast it to fp32, however running this script:
> ....
> Led to the same error, what version of optimum and `auto-gptq` are you using?

Hi @younesbelkada, can you please add this flag `load_gptq_on_cpu=True` while loading the model? The below code runs fine for me:
```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig

checkpoint = "marcsun13/opt-350m-gptq-4bit"
device = "cpu" # for GPU usage or "cpu" for CPU usage

tokenizer = AutoTokenizer.from_pretrained(checkpoint)

quantization_config = GPTQConfig(bits=4, disable_exllama=True)

model = AutoModelForCausalLM.from_pretrained(checkpoint, low_cpu_mem_usage=True, quantization_config=quantization_config, load_gptq_on_cpu=True).to(device)

model = model.to(torch.float32)

inputs = tokenizer.encode("Hello how are you?", return_tensors="pt").to(device)
outputs = model.generate(inputs, max_new_tokens=4, do_sample=False)
print(tokenizer.decode(outputs[0]))
```

The result:
```
Downloading pytorch_model.bin: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 125M/125M [00:03<00:00, 37.4MB/s]
CUDA extension not installed.
Downloading (…)neration_config.json: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 137/137 [00:00<00:00, 1.10MB/s]

Hello how are you?????
```

@younesbelkada, now I got it. I think making these changes https://github.com/PanQiWei/AutoGPTQ/pull/367 locally would fix the issue for you.

> Can you also confirm that loading the model in float32 works for you instead of first loading your model in fp16 then cast it back?

@younesbelkada, first I tried loading the model in float32 only but that didn't work. Even after passing the `torch_dtype` as `torch.float32`, the model loaded had the weight in fp16. The same was the issue with your code snippet, and that's why we have to do the casting explicitly. If you think there's an issue, and it could be fixed then I would be happy to drop this patch.

Edit: Sorry, we can't drop this patch, since it enables loading the model on the CPU, but yeah `.to` is required if the `torch_dtype` is not working fine.

Hi @vivekkhandelwal1 

Hmm this is a bit strange, running: 

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig

checkpoint = "marcsun13/opt-350m-gptq-4bit"
device = "cpu" # for GPU usage or "cpu" for CPU usage

tokenizer = AutoTokenizer.from_pretrained(checkpoint)

quantization_config = GPTQConfig(bits=4, disable_exllama=True)

model = AutoModelForCausalLM.from_pretrained(checkpoint, low_cpu_mem_usage=True, quantization_config=quantization_config, torch_dtype=torch.float32).to(device)

for n, p in model.named_parameters():
    print(n, p.dtype)
```

returns:

```bash
model.decoder.embed_tokens.weight torch.float32
model.decoder.embed_positions.weight torch.float32
model.decoder.final_layer_norm.weight torch.float32
model.decoder.final_layer_norm.bias torch.float32
model.decoder.layers.0.self_attn_layer_norm.weight torch.float32
model.decoder.layers.0.self_attn_layer_norm.bias torch.float32
model.decoder.layers.0.final_layer_norm.weight torch.float32
model.decoder.layers.0.final_layer_norm.bias torch.float32
model.decoder.layers.1.self_attn_layer_norm.weight torch.float32
model.decoder.layers.1.self_attn_layer_norm.bias torch.float32
model.decoder.layers.1.final_layer_norm.weight torch.float32
model.decoder.layers.1.final_layer_norm.bias torch.float32
model.decoder.layers.2.self_attn_layer_norm.weight torch.float32
model.decoder.layers.2.self_attn_layer_norm.bias torch.float32
model.decoder.layers.2.final_layer_norm.weight torch.float32
model.decoder.layers.2.final_layer_norm.bias torch.float32
model.decoder.layers.3.self_attn_layer_norm.weight torch.float32
model.decoder.layers.3.self_attn_layer_norm.bias torch.float32
model.decoder.layers.3.final_layer_norm.weight torch.float32
model.decoder.layers.3.final_layer_norm.bias torch.float32
model.decoder.layers.4.self_attn_layer_norm.weight torch.float32
model.decoder.layers.4.self_attn_layer_norm.bias torch.float32
model.decoder.layers.4.final_layer_norm.weight torch.float32
model.decoder.layers.4.final_layer_norm.bias torch.float32
model.decoder.layers.5.self_attn_layer_norm.weight torch.float32
model.decoder.layers.5.self_attn_layer_norm.bias torch.float32
model.decoder.layers.5.final_layer_norm.weight torch.float32
model.decoder.layers.5.final_layer_norm.bias torch.float32
model.decoder.layers.6.self_attn_layer_norm.weight torch.float32
model.decoder.layers.6.self_attn_layer_norm.bias torch.float32
model.decoder.layers.6.final_layer_norm.weight torch.float32
model.decoder.layers.6.final_layer_norm.bias torch.float32
model.decoder.layers.7.self_attn_layer_norm.weight torch.float32
model.decoder.layers.7.self_attn_layer_norm.bias torch.float32
model.decoder.layers.7.final_layer_norm.weight torch.float32
model.decoder.layers.7.final_layer_norm.bias torch.float32
model.decoder.layers.8.self_attn_layer_norm.weight torch.float32
model.decoder.layers.8.self_attn_layer_norm.bias torch.float32
model.decoder.layers.8.final_layer_norm.weight torch.float32
model.decoder.layers.8.final_layer_norm.bias torch.float32
model.decoder.layers.9.self_attn_layer_norm.weight torch.float32
model.decoder.layers.9.self_attn_layer_norm.bias torch.float32
model.decoder.layers.9.final_layer_norm.weight torch.float32
model.decoder.layers.9.final_layer_norm.bias torch.float32
model.decoder.layers.10.self_attn_layer_norm.weight torch.float32
model.decoder.layers.10.self_attn_layer_norm.bias torch.float32
model.decoder.layers.10.final_layer_norm.weight torch.float32
model.decoder.layers.10.final_layer_norm.bias torch.float32
model.decoder.layers.11.self_attn_layer_norm.weight torch.float32
model.decoder.layers.11.self_attn_layer_norm.bias torch.float32
model.decoder.layers.11.final_layer_norm.weight torch.float32
model.decoder.layers.11.final_layer_norm.bias torch.float32
```

+ inspecting some `QuantLinear` modules I can't see any attribute in fp16. Let me get back to you on this

Hi @vivekkhandelwal1 when building auto-gptq using https://github.com/PanQiWei/AutoGPTQ/pull/367 + running your script above I am facing.

```bash
  File "../qlinear_exllamav2.py", line 137, in post_init
    assert self.qweight.device.type == "cuda"
AssertionError
```

Do you think there is anything I missed on my end?

> Hi @vivekkhandelwal1 when building auto-gptq using [PanQiWei/AutoGPTQ#367](https://github.com/PanQiWei/AutoGPTQ/pull/367) + running your script above I am facing.
> 
> ```shell
>   File "../qlinear_exllamav2.py", line 137, in post_init
>     assert self.qweight.device.type == "cuda"
> AssertionError
> ```
> 
> Do you think there is anything I missed on my end?

Can you please try building the auto_gptq with:
```
BUILD_CUDA_EXT=0 pip install -v .
```

Hi @younesbelkada, my Auto-GPTQ PR https://github.com/PanQiWei/AutoGPTQ/pull/367 is merged. Can we go ahead with this?

Hi @vivekkhandelwal1 let me have another look and get back to you

Hi @vivekkhandelwal1 , thanks again! 
I made https://github.com/PanQiWei/AutoGPTQ/pull/376 in order for this script: 
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig

checkpoint = "marcsun13/opt-350m-gptq-4bit"
device = "cpu" # for GPU usage or "cpu" for CPU usage

tokenizer = AutoTokenizer.from_pretrained(checkpoint)

quantization_config = GPTQConfig(bits=4, disable_exllama=True)

model = AutoModelForCausalLM.from_pretrained(checkpoint, low_cpu_mem_usage=True, quantization_config=quantization_config, torch_dtype=torch.float32).to(device)

inputs = tokenizer.encode("Hello how are you?", return_tensors="pt").to(device)
outputs = model.generate(inputs, max_new_tokens=4, do_sample=False)
print(tokenizer.decode(outputs[0]))
```
to work. However, I found it quite slow to get generations on CPU, out of curiosity how long does it takes to generate text on falcon 180B on your end?

> However, I found it quite slow to get generations on CPU, out of curiosity how long does it takes to generate text on falcon 180B on your end?

Hi @younesbelkada, the performance of Falcon-180B-GPTQ on the CPU is quite slow. We lower this model through Torch-MLIR and then compile via IREE to run over the CUDA and the Rocm backend. The Torch-MLIR doesn't support cuda tensors that's why we need to have this model loaded over the CPU.

Also, were you getting some errors even after my patch was merged in Auto-GPTQ? I mean, what's the reason behind this https://github.com/PanQiWei/AutoGPTQ/pull/376/files.

Hi @vivekkhandelwal1 
Yes I was getting some errors even with your patch, in my script I don't call `to()` with a torch dtype but directly load in fp32 with `torch_dtype=torch.float32`, for some reasons some weights still remained in fp32. Note we don't support explicit model casting for quantized weights since: https://github.com/huggingface/transformers/pull/26761

> Hi @vivekkhandelwal1 Yes I was getting some errors even with your patch, in my script I don't call `to()` with a torch dtype but directly load in fp32 with `torch_dtype=torch.float32`, for some reasons some weights still remained in fp32. Note we don't support explicit model casting for quantized weights since: #26761

I'm also not doing the explicit model casting and without that it worked for me.

> Author

I have made two more changes(https://github.com/PanQiWei/AutoGPTQ/pull/367/files#diff-c7731808a14c99106c4a0e48729c7435181a1085978ca567433d28a5f2473b1dL272) in the Auto-GPTQ if you try the complete changes of this https://github.com/PanQiWei/AutoGPTQ/pull/367/files, then your patch is not required. I have tried it just now for the script you shared.

Hi @vivekkhandelwal1 
I cloned again auto-gptq and I can confirm I am on the main branch `git log` gives me your commit and I uninstalled and installed again auto-gptq with `BUILD_CUDA_EXT=0 pip install -U  -v .` and ran into:

```bash
s/qlinear/qlinear_cuda_old.py", line 269, in forward
    out = torch.matmul(x.to(weight.dtype), weight)
RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
```

Sharing again the snippet I use: 

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig

checkpoint = "marcsun13/opt-350m-gptq-4bit"
device = "cpu" # for GPU usage or "cpu" for CPU usage

tokenizer = AutoTokenizer.from_pretrained(checkpoint)

quantization_config = GPTQConfig(bits=4, disable_exllama=True)

model = AutoModelForCausalLM.from_pretrained(checkpoint, low_cpu_mem_usage=True, quantization_config=quantization_config, torch_dtype=torch.float32).to(device)

inputs = tokenizer.encode("Hello how are you?", return_tensors="pt").to(device)
outputs = model.generate(inputs, max_new_tokens=4, do_sample=False)
print(tokenizer.decode(outputs[0]))
``` 

Let me know if I missed anything!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26719). All of your documentation changes will be reflected on that endpoint.

> Hi @vivekkhandelwal1 I cloned again auto-gptq and I can confirm I am on the main branch `git log` gives me your commit and I uninstalled and installed again auto-gptq with `BUILD_CUDA_EXT=0 pip install -U -v .` and ran into:
> 
> ```shell
> s/qlinear/qlinear_cuda_old.py", line 269, in forward
>     out = torch.matmul(x.to(weight.dtype), weight)
> RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
> ```
> 
> Sharing again the snippet I use:
> 
> ```python
> import torch
> from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig
> 
> checkpoint = "marcsun13/opt-350m-gptq-4bit"
> device = "cpu" # for GPU usage or "cpu" for CPU usage
> 
> tokenizer = AutoTokenizer.from_pretrained(checkpoint)
> 
> quantization_config = GPTQConfig(bits=4, disable_exllama=True)
> 
> model = AutoModelForCausalLM.from_pretrained(checkpoint, low_cpu_mem_usage=True, quantization_config=quantization_config, torch_dtype=torch.float32).to(device)
> 
> inputs = tokenizer.encode("Hello how are you?", return_tensors="pt").to(device)
> outputs = model.generate(inputs, max_new_tokens=4, do_sample=False)
> print(tokenizer.decode(outputs[0]))
> ```
> 
> Let me know if I missed anything!

Hi @younesbelkada, I just tried your script on another machine with a fresh setup. I cloned the auto_gptq repo built it and made the changes in transformers same as this PR, and it worked for me. If you still have some issues locally which are fixed by your changes in Auto-GPTQ then we can do that as well. But, can we now go ahead with this patch? If you want then we can discuss this on a call, I can help you set this up and try.

Thanks a lot @vivekkhandelwal1 
The changes look great to me, I tried it again :D and still facing the same issue :/ I suggest we wait https://github.com/PanQiWei/AutoGPTQ/pull/376 to be merged before merging this PR
Can you also add one line in the documentation explaining that you can perform CPU inference with auto-gptq? the new paragraph can come right after: https://github.com/huggingface/transformers/blob/main/docs/source/en/main_classes/quantization.md#exllama-kernels-for-faster-inference - what do you think?


> Thanks a lot @vivekkhandelwal1 The changes look great to me, I tried it again :D and still facing the same issue :/ I suggest we wait [PanQiWei/AutoGPTQ#376](https://github.com/PanQiWei/AutoGPTQ/pull/376) to be merged before merging this PR Can you also add one line in the documentation explaining that you can perform CPU inference with auto-gptq? the new paragraph can come right after: https://github.com/huggingface/transformers/blob/main/docs/source/en/main_classes/quantization.md#exllama-kernels-for-faster-inference - what do you think?

@younesbelkada done! Also, I'm still not sure if your patch is needed or not, also the GPTQ guys don't seem convinced about getting your patch in. Can we get this in so that I'm unblocked on this, and it works for me as well?

Hi @vivekkhandelwal1 
Thanks! I think @SunMarc managed to reproduce the same issue that I had on his end, so I prefer to wait before https://github.com/PanQiWei/AutoGPTQ/pull/385 gets merged, let me ping @fxmarty internally and see if we can get this merged ASAP
Thanks again!

> Hi @vivekkhandelwal1 Thanks! I think @SunMarc managed to reproduce the same issue that I had on his end, so I prefer to wait before [PanQiWei/AutoGPTQ#385](https://github.com/PanQiWei/AutoGPTQ/pull/385) gets merged, let me ping @fxmarty internally and see if we can get this merged ASAP Thanks again!

Sure @younesbelkada!

@younesbelkada, can we now merge this PR since the Auto-GPTQ PRs are merged now? https://github.com/PanQiWei/AutoGPTQ/pull/385
https://github.com/huggingface/optimum/pull/1496

> Thank you very much for your great work! @vivekkhandelwal1

Thanks, @younesbelkada for the support and feedback. I don't have merge access to the repo, can you please merge it for me?

Thanks again for all your work on this @vivekkhandelwal1 !