Cache in different devices when use split model with dispatch_model() function and model.generate()
Author: ZetangForward
Pull request: no
Labels: name: Generation

name: Cache

### System Info

- `transformers` version: 4.40.1
- Platform: Linux-5.15.0-67-generic-x86_64-with-glibc2.31
- Python version: 3.10.4
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.1
- Accelerate version: 0.27.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2+cu121 (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?

@ArthurZucker @younesbelkada @gante @pacman100 @SunMarc 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

1. Load Model and Split the model on 4 Devices
```python3
import transformers
from accelerate.utils import infer_auto_device_map, get_balanced_memory
model = transformers.LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
device_map = "balanced_low_0"
no_split_module_classes = transformers.LlamaForCausalLM._no_split_modules
max_memory = get_balanced_memory(
    model,
    max_memory=None,
    no_split_module_classes=no_split_module_classes,
    dtype=None,
    low_zero=(device_map=="balanced_low_0"),
)
device_map = infer_auto_device_map(
    model, max_memory=max_memory, no_split_module_classes=no_split_module_classes, dtype=None
)
model = dispatch_model(model, device_map=device_map)
```
2. Conduct Model.generate()
```
input_ids = tokenizer(batch['context_str'], return_tensors='pt').input_ids.to(model.device)
output = model.generate(input_ids, min_new_tokens=1, max_new_tokens=max_gen_len, use_cache=True)
```

### Expected behavior

Return a BUG:
```
  0%|                                                                                          | 0/5120 [00:16<?, ?it/s]
Traceback (most recent call last):
  File "/nvme1/zecheng/modelzipper/projects/state-space-model/src/test_dev2.py", line 200, in <module>
    main(config)
  File "/nvme1/zecheng/modelzipper/projects/state-space-model/src/test_dev2.py", line 147, in main
    output = model.generate(input_ids, min_new_tokens=1, max_new_tokens=max_gen_len, use_cache=True)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/transformers/generation/utils.py", line 1576, in generate
    result = self._greedy_search(
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/transformers/generation/utils.py", line 2494, in _greedy_search
    outputs = self(
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1208, in forward
    outputs = self.model(
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1018, in forward
    layer_outputs = decoder_layer(
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 741, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 454, in forward
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
  File "/home/amax/anaconda3/envs/zecheng/lib/python3.10/site-packages/transformers/cache_utils.py", line 146, in update
    self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:2! (when checking argument for argument tensors in method wrapper_CUDA_cat)
```

Comments:
When I used transformers=4.40.1, I also encountered this issue. However, in an earlier version, such as 4.32.0, this problem doesn’t occur. In version 4.32.0, the llama model code in the modeling_llama.py file runs normally. It seems to be related to the use of the cache class, especially the part involving 'from ...cache_utils import Cache, DynamicCache, StaticCache'. Can anyone provide some help? Thank you.

> When I used transformers=4.40.1, I also encountered this issue. However, in an earlier version, such as 4.32.0, this problem doesn’t occur. In version 4.32.0, the llama model code in the modeling_llama.py file runs normally. It seems to be related to the use of the cache class, especially the part involving 'from ...cache_utils import Cache, DynamicCache, StaticCache'. Can anyone provide some help? Thank you.

I have the same problem. I would be grateful if you tell me when you solve it :) 

cc @SunMarc and @muellerzr not sure if we broke something or is it accelerate? 😅 

Hi there @ZetangForward @BlackWaterDrop @Miiyuu 👋 

I've tried to reproduce the error with the following stand-alone script (running on 2 GPUs):
```py
import transformers
from accelerate.utils import infer_auto_device_map, get_balanced_memory
from accelerate import dispatch_model
model = transformers.LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", cache_dir="/raid/joao/")
device_map = "balanced_low_0"
no_split_module_classes = transformers.LlamaForCausalLM._no_split_modules
max_memory = get_balanced_memory(
    model,
    max_memory=None,
    no_split_module_classes=no_split_module_classes,
    dtype=None,
    low_zero=(device_map=="balanced_low_0"),
)
device_map = infer_auto_device_map(
    model, max_memory=max_memory, no_split_module_classes=no_split_module_classes, dtype=None
)
model = dispatch_model(model, device_map=device_map)

demo_text = ["foo", "bar", "baz", "quc", "this is a longer text to force padding", "potato"]
tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", cache_dir="/raid/joao/", padding_side="left")
tokenizer.pad_token = tokenizer.eos_token
inputs = tokenizer(demo_text, return_tensors='pt', padding=True).to(model.device)
output = model.generate(**inputs, min_new_tokens=1, max_new_tokens=10, use_cache=True)
```

with
```
- `transformers` version: 4.42.0.dev0 (commit = eae2b6b89e72dbee531e5b96acc94b98161aa8a4)
- Platform: Linux-5.4.0-166-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.23.1
- Safetensors version: 0.4.3
- Accelerate version: 0.29.3
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.2+cu121 (True)
- Tensorflow version (GPU?): 2.13.1 (True)
- Flax version (CPU?/GPU?/TPU?): 0.7.0 (cpu)
- Jax version: 0.4.13
- JaxLib version: 0.4.13
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>
```

and I couldn't reproduce the issue. Could you try to update your libraries (namely `accelerate` and `transformers`), and letting us know if you still have the same issue? 🤗 

> Hi there @ZetangForward @BlackWaterDrop @Miiyuu 👋
> 
> I've tried to reproduce the error with the following stand-alone script (running on 2 GPUs):
> 
> ```python
> import transformers
> from accelerate.utils import infer_auto_device_map, get_balanced_memory
> from accelerate import dispatch_model
> model = transformers.LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", cache_dir="/raid/joao/")
> device_map = "balanced_low_0"
> no_split_module_classes = transformers.LlamaForCausalLM._no_split_modules
> max_memory = get_balanced_memory(
>     model,
>     max_memory=None,
>     no_split_module_classes=no_split_module_classes,
>     dtype=None,
>     low_zero=(device_map=="balanced_low_0"),
> )
> device_map = infer_auto_device_map(
>     model, max_memory=max_memory, no_split_module_classes=no_split_module_classes, dtype=None
> )
> model = dispatch_model(model, device_map=device_map)
> 
> demo_text = ["foo", "bar", "baz", "quc", "this is a longer text to force padding", "potato"]
> tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", cache_dir="/raid/joao/", padding_side="left")
> tokenizer.pad_token = tokenizer.eos_token
> inputs = tokenizer(demo_text, return_tensors='pt', padding=True).to(model.device)
> output = model.generate(**inputs, min_new_tokens=1, max_new_tokens=10, use_cache=True)
> ```
> 
> with
> 
> ```
> - `transformers` version: 4.42.0.dev0 (commit = eae2b6b89e72dbee531e5b96acc94b98161aa8a4)
> - Platform: Linux-5.4.0-166-generic-x86_64-with-glibc2.29
> - Python version: 3.8.10
> - Huggingface_hub version: 0.23.1
> - Safetensors version: 0.4.3
> - Accelerate version: 0.29.3
> - Accelerate config: 	not found
> - PyTorch version (GPU?): 2.1.2+cu121 (True)
> - Tensorflow version (GPU?): 2.13.1 (True)
> - Flax version (CPU?/GPU?/TPU?): 0.7.0 (cpu)
> - Jax version: 0.4.13
> - JaxLib version: 0.4.13
> - Using GPU in script?: <fill in>
> - Using distributed or parallel set-up in script?: <fill in>
> ```
> 
> and I couldn't reproduce the issue. Could you try to update your libraries (namely `accelerate` and `transformers`), and letting us know if you still have the same issue? 🤗

Thx, it maybe the version issue... I have updated my transformers and accelerate to the latest version and haven't encountered those issues again.  I plan to close this issue until anyone encounter such an issue again.

其实这个就是use_cache的问题啦，后面的代码执行时需要用到前面的cache，因此导致了设备不匹配的问题，把他设为false就行了

> > 你好呀@ZetangForward @BlackWaterDrop @Miiyuu👋
> > 我尝试使用以下独立脚本（在 2 个 GPU 上运行）重现该错误：
> > ```python
> > import transformers
> > from accelerate.utils import infer_auto_device_map, get_balanced_memory
> > from accelerate import dispatch_model
> > model = transformers.LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", cache_dir="/raid/joao/")
> > device_map = "balanced_low_0"
> > no_split_module_classes = transformers.LlamaForCausalLM._no_split_modules
> > max_memory = get_balanced_memory(
> >     model,
> >     max_memory=None,
> >     no_split_module_classes=no_split_module_classes,
> >     dtype=None,
> >     low_zero=(device_map=="balanced_low_0"),
> > )
> > device_map = infer_auto_device_map(
> >     model, max_memory=max_memory, no_split_module_classes=no_split_module_classes, dtype=None
> > )
> > model = dispatch_model(model, device_map=device_map)
> > 
> > demo_text = ["foo", "bar", "baz", "quc", "this is a longer text to force padding", "potato"]
> > tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", cache_dir="/raid/joao/", padding_side="left")
> > tokenizer.pad_token = tokenizer.eos_token
> > inputs = tokenizer(demo_text, return_tensors='pt', padding=True).to(model.device)
> > output = model.generate(**inputs, min_new_tokens=1, max_new_tokens=10, use_cache=True)
> > ```
> > 
> > 
> >     
> >       
> >     
> > 
> >       
> >     
> > 
> >     
> >   
> > 和
> > ```
> > - `transformers` version: 4.42.0.dev0 (commit = eae2b6b89e72dbee531e5b96acc94b98161aa8a4)
> > - Platform: Linux-5.4.0-166-generic-x86_64-with-glibc2.29
> > - Python version: 3.8.10
> > - Huggingface_hub version: 0.23.1
> > - Safetensors version: 0.4.3
> > - Accelerate version: 0.29.3
> > - Accelerate config: 	not found
> > - PyTorch version (GPU?): 2.1.2+cu121 (True)
> > - Tensorflow version (GPU?): 2.13.1 (True)
> > - Flax version (CPU?/GPU?/TPU?): 0.7.0 (cpu)
> > - Jax version: 0.4.13
> > - JaxLib version: 0.4.13
> > - Using GPU in script?: <fill in>
> > - Using distributed or parallel set-up in script?: <fill in>
> > ```
> > 
> > 
> >     
> >       
> >     
> > 
> >       
> >     
> > 
> >     
> >   
> > 我无法重现该问题。您可以尝试更新您的库（即`accelerate`和`transformers`），然后让我们知道您是否仍然遇到相同的问题？🤗
> 
> 谢谢，这可能是版本问题……我已经将我的 transformers 和加速器更新到最新版本，并且没有再遇到这些问题。我打算关闭此问题，直到有人再次遇到此类问题。

其实这个就是use_cache的问题啦，后面的代码执行时需要用到前面的cache，因此导致了设备不匹配的问题，把他设为false就行了