Cannot load AWQ model on GPU with deepspeed zero3
Author: ghost
Pull request: no
Labels: name: DeepSpeed

name: Quantization

### System Info

- `transformers` version: 4.39.0.dev0
- Platform: Linux-5.15.0-92-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.21.4
- Safetensors version: 0.4.2
- Accelerate version: 0.28.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.2.0a0+81ea7a4 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>


### Who can help?

@pacman100

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Hi, I want to load from_pretrained a AWQ quantized model using zero3.
But I got the following warning:
```
[WARNING|logging.py:329] 2024-04-18 09:42:18,716 >> **You have loaded an AWQ model on CPU and have a CUDA device available, make sure to set your model on a GPU device in order to run your model.**
[WARNING|modeling_utils.py:3054] 2024-04-18 09:42:18,716 >> `low_cpu_mem_usage` was None, now set to True since model is quantized.
[INFO|modeling_utils.py:3276] 2024-04-18 09:42:18,716 >> loading weights file /workspace/lama-awq/model.safetensors.index.json
[INFO|modeling_utils.py:1413] 2024-04-18 09:42:18,717 >> Instantiating LlamaForCausalLM model under default dtype torch.float16.
[WARNING|logging.py:329] 2024-04-18 09:42:18,718 >> You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
[INFO|configuration_utils.py:927] 2024-04-18 09:42:18,720 >> Generate config GenerationConfig {
  "bos_token_id": 1,
  "eos_token_id": 2,
  "use_cache": false
}

```

the code:
```
   model = AutoModelForCausalLM.from_pretrained(
            model_args.model_name_or_path,
            torch_dtype=torch.float16,
            trust_remote_code=True,
            # device_map='auto',
            # low_cpu_mem_usage=True,
            attn_implementation="flash_attention_2",
        )
```
the deepspeed config is:
```
compute_environment: LOCAL_MACHINE                                                                                                                                           
debug: false
deepspeed_config:
  deepspeed_multinode_launcher: standard
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```
After I change the model to a standard fp16 model, the model were loaded on the GPU during 'AutoModelForCausalLM.from_pretrained'.

### Expected behavior

How to load an AWQ model on GPU? Looking forward to your reply.

Comments:
cc @SunMarc 

Hi @mxjyst, to move a model on GPU, all you need to do is to `model.to('cuda')`. Could you shared with us a minimal reproducer and the error that you are getting ? This is normal that you are getting with error since you loaded your model on cpu. Another solution would be to pass device_map="cuda" in `from_pretrained`. 

Here is a reproducer.
train.py
```
from transformers import (
    AutoModelForCausalLM,
)
import torch
def main():
    # Set seed for reproducibility
    

    model = AutoModelForCausalLM.from_pretrained(
            "Qwen/Qwen1.5-1.8B-Chat-AWQ",
            trust_remote_code=True,
            attn_implementation="flash_attention_2",
            )
    print(torch.cuda.max_memory_allocated())

if __name__ == "__main__":
    main()
```
**The pretrained model is a AWQ quantized model `Qwen1.5-1.8B-Chat-AWQ`.**
run commanr:
`accelerate launch --config_file "z3.yaml" train.py`

z3.yaml:
```
compute_environment: LOCAL_MACHINE                                                                                                                                           
debug: false
deepspeed_config:
  deepspeed_multinode_launcher: standard
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

I set gpu number is 4. The output of train.py is:
```
[2024-04-22 01:50:59,116] torch.distributed.run: [WARNING]
[2024-04-22 01:50:59,116] torch.distributed.run: [WARNING] *****************************************
[2024-04-22 01:50:59,116] torch.distributed.run: [WARNING] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
[2024-04-22 01:50:59,116] torch.distributed.run: [WARNING] *****************************************
You have loaded an AWQ model on CPU and have a CUDA device available, make sure to set your model on a GPU device in order to run your model.
`low_cpu_mem_usage` was None, now set to True since model is quantized.
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
You have loaded an AWQ model on CPU and have a CUDA device available, make sure to set your model on a GPU device in order to run your model.
`low_cpu_mem_usage` was None, now set to True since model is quantized.
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
You have loaded an AWQ model on CPU and have a CUDA device available, make sure to set your model on a GPU device in order to run your model.
`low_cpu_mem_usage` was None, now set to True since model is quantized.
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
You have loaded an AWQ model on CPU and have a CUDA device available, make sure to set your model on a GPU device in order to run your model.
`low_cpu_mem_usage` was None, now set to True since model is quantized.
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
0
0
0
0
```
It seems AWQ model didn't be loaded on GPU.

I create a test repo following PEFT example with Deepspeed Zero3. 
link: https://github.com/mxjyst/test
I add some code in utils.py line 138 to print the memory_allocated of the model.
When run [run_standard.sh](https://github.com/mxjyst/test/blob/main/run_standard.sh), the memory_allocated is about (model size // num gpu). 
When run [run_awq.sh](https://github.com/mxjyst/test/blob/main/run_awq.sh), the memory_allocated is 0. And if I continue to run the train code, the memory_allocated is not recuded compared to w/o deepspeed. It seems the model were not partitioned across the processes. 
Looking forward to your reply.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

AWQ is not compatible yet with deepspeed zero 3. If someone manage to make it work, please share it here ! 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.