RWKV v4 not working with device_map auto and 4 GPUs
Author: Epliz
Pull request: no

### System Info

- `transformers` version: 4.33.3
- Platform: Linux-4.18.0-477.13.1.el8_8.x86_64-x86_64-with-glibc2.2.5
- Python version: 3.8.16
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (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?: yes, 4 GPUs (Nvidia A10)
- Using distributed or parallel set-up in script?: nothing custom

### Who can help?

@ArthurZucker , @younesbelkada , @pacman100 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Try the following scripts, either loading the model with the first approach or the second:
case 1 (works):
```
from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "RWKV/rwkv-4-169m-pile"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
# Case 1: works
model = AutoModelForCausalLM.from_pretrained(checkpoint)
model.to("cuda")

from transformers import pipeline
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, do_sample=False,max_new_tokens=1024, use_cache=True, device=0)

pipe("What is life?")
```
case 2 (doesn't work):
```
from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "RWKV/rwkv-4-169m-pile"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto")


from transformers import pipeline
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, do_sample=False,max_new_tokens=1024, use_cache=True)

pipe("What is life?")
```

error (when setting CUDA_LAUNCH_BLOCKING):
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>                                                                                                                                                                     
  File "<stdin>", line 4, in time_func                                                                                                                                                                    
  File "<stdin>", line 1, in <lambda>                                                                                                                                                                     
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/pipelines/text_generation.py", line 205, in __call__                                                                     
    return super().__call__(text_inputs, **kwargs)                                                                                                                                                        
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/pipelines/base.py", line 1140, in __call__                                                                               
    return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)                                                                                                                 
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/pipelines/base.py", line 1147, in run_single                                                                             
    model_outputs = self.forward(model_inputs, **forward_params)                                                                                                                                          
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/pipelines/base.py", line 1046, in forward                                                                                
    model_outputs = self._forward(model_inputs, **forward_params)                                                                                                                                         
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/pipelines/text_generation.py", line 268, in _forward                                                                     
    generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)                                                                                       
  File "<DIR>venv38/lib64/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context                                                                            
    return func(*args, **kwargs)                                                                                                                                                                          
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/generation/utils.py", line 1602, in generate                                                                             
    return self.greedy_search(                                                                                                                                                                            
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/generation/utils.py", line 2450, in greedy_search                                                                        
    outputs = self(                                                                                                                                                                                          File "<DIR>venv38/lib64/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl                                                                                  
    return forward_call(*args, **kwargs)        
  File "<DIR>venv38/lib64/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)                                                                                                                                                                 
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/models/rwkv/modeling_rwkv.py", line 815, in forward                                                                       
    rwkv_outputs = self.rwkv(                                            
  File "<DIR>venv38/lib64/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl                                                                                 
    return forward_call(*args, **kwargs)                                                                                                                                                                   
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/models/rwkv/modeling_rwkv.py", line 690, in forward                                                                       
    hidden_states, state, attentions = block(                                                                                                                                                              
  File "<DIR>venv38/lib64/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl                                                                                  
    return forward_call(*args, **kwargs)                                                                                                                                                                   
  File "<DIR>venv38/lib64/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward                                                                                         
    output = old_forward(*args, **kwargs)                                                                                                                                                                  
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/models/rwkv/modeling_rwkv.py", line 384, in forward                                                                       
    attention, state = self.attention(self.ln1(hidden), state=state, use_cache=use_cache)                                                                                                                  
  File "<DIR>venv38/lib64/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl                                                                                  
    return forward_call(*args, **kwargs)                                                                                                                                                                  
  File "<DIR>venv38/lib64/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward                                                                                         
    output = old_forward(*args, **kwargs)                                                                                                                                                                  
  File "<DIR>venv38/lib64/python3.8/site-packages/transformers/models/rwkv/modeling_rwkv.py", line 320, in forward                                                                       
    state[2][:, :, self.layer_id] = layer_state[0]                                                                                                                                                         
RuntimeError: CUDA error: an illegal memory access was encountered                                                                                                                                         
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

### Expected behavior

Both should work

Comments:
Maybe of interest to @SunMarc when back from leave :)

Hi @Epliz

I have ran:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "RWKV/rwkv-4-169m-pile"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto")


from transformers import pipeline
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, do_sample=False,max_new_tokens=1024, use_cache=True)

print(pipe("What is life?"))
```
On 2x NVIDIA T4 GPUs and it seemed to work fine. I am using the same accelerate version as the one you shared. The only difference is that I am using transformers main branch - perhaps you can try to switch to transformers main and see if the issue persists?

Can you also print `model.hf_device_map` and `set(model.hf_device_map.values())` ?

Thank you @younesbelkada  for commenting.
I think we are getting somewhere.
Here is what I observe: it doesn't work with the released & git versions of HuggingFace Transformers for 4 GPUs.
What I can see from the device maps is that the GPU0 is not used? The device maps in both cases are the following:
```
>>>model.hf_device_map
{'rwkv.embeddings': 1, 'rwkv.blocks.0': 1, 'rwkv.blocks.1': 2, 'rwkv.blocks.2': 2, 'rwkv.blocks.3': 2, 'rwkv.blocks.4': 2, 'rwkv.blocks.5': 2, 'rwkv.blocks.6': 2, 'rwkv.blocks.7': 3, 'rwkv.blocks.8': 3, 'rwkv.blocks.9': 3, 'rwkv.blocks.10': 3, 'rwkv.blocks.11': 3, 'rwkv.ln_out': 3, 'head': 3}

>>>set(model.hf_device_map.values())
{1, 2, 3}
```
When I use only 2 GPUs (by changing CUDA_VISIBLE_DEVICES), the generation works, and the device map uses GPU0:
```
>>> model.hf_device_map
{'rwkv.embeddings': 0, 'rwkv.blocks.0': 0, 'rwkv.blocks.1': 0, 'rwkv.blocks.2': 1, 'rwkv.blocks.3': 1, 'rwkv.blocks.4': 1, 'rwkv.blocks.5': 1, 'rwkv.blocks.6': 1, 'rwkv.blocks.7': 1, 'rwkv.blocks.8': 1, 'rwkv.blocks.9': 1, 'rwkv.blocks.10': 1, 'rwkv.blocks.11': 1, 'rwkv.ln_out': 1, 'head': 1}
>>> set(model.hf_device_map.values())
{0, 1}
```
But it seems to get weirder when trying to use 3 GPUs as the generation works, but the device map seems to show that GPU0 is not used:
```
>>> model.hf_device_map
{'rwkv.embeddings': 1, 'rwkv.blocks.0': 1, 'rwkv.blocks.1': 1, 'rwkv.blocks.2': 1, 'rwkv.blocks.3': 2, 'rwkv.blocks.4': 2, 'rwkv.blocks.5': 2, 'rwkv.blocks.6': 2, 'rwkv.blocks.7': 2, 'rwkv.blocks.8': 2, 'rwkv.blocks.9': 2, 'rwkv.blocks.10': 2, 'rwkv.blocks.11': 2, 'rwkv.ln_out': 2, 'head': 2}
>>> set(model.hf_device_map.values())
{1, 2}
```

Even if the model is only on device 1 and 2, it should work. This is a problem on pytorch [side](https://github.com/pytorch/pytorch/issues/21819). If the model fits in two GPUs, I would suggest to only use 2 GPUs as there is no reason to split it across 3 GPUs. It will be faster because there will be less communication overhead. 

Thank you @SunMarc  for your reply.

Do you think it might be a pytorch issue even though it works with transformers==4,29.1 but not 4.29.2 and afterwards (all other dependencies staying the same)?
I see that in between those versions, there were changes to make sure RWKV CUDA kernels were included in the package.

Best regards,
Epliz

Thanks for your investigation @Epliz. This is indeed strange. I will try to reproduce this with 4 gpus. Maybe the number of gpus matters. IN the meantime, i would suggest using only two gpus since the model is not that big. 

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.