`model.generate` fail for Jamba with `device_map="auto"`
Author: TechxGenus
Pull request: no
Labels: name: Big Model Inference

### System Info

Name: transformers
Version: 4.40.0

### Who can help?

@ArthurZucker 

### 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

When running Jamba with multiple GPUs, an error is reported:
```
RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered
```
After disabling mamba kernel (`use_mamba_kernels=False`), getting the following error:
```
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!
```
The minimal code to run is as follows:
```
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
from transformers import AutoTokenizer, AutoModelForCausalLM

prompt = "def quicksort(arr):\n"
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-tiny-random")
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    "ai21labs/Jamba-tiny-random",
    device_map="auto",
)
inputs = tokenizer.encode(prompt, return_tensors="pt")
outputs = model.generate(input_ids=inputs.to(model.device), max_new_tokens=512, do_sample=False)
print(tokenizer.decode(outputs[0]))
```

When using a single GPU or using CPU, it works well

### Expected behavior

Work well when using multiple GPUs

Comments:
I can reproduce, would you like to open a pr for a fix? 
(cc @SunMarc otherwise, maybe the no_split_modules should be updated ) 

@ArthurZucker not seem to be a bug with accelerate or no_split_modules
In https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py#L1729-L1731:
```
            past_key_values = HybridMambaAttentionDynamicCache(
                self.config, input_ids.shape[0], self.dtype, device=self.device
            )
```
The hybrid cache here is initialized to the wrong device. Since `ssm_state` of the mamba layer is not empty at the beginning, it may be on a different device than the specific layer, so an error occured during the first forward.
I'm not sure what the best fix is, moving the tensor manually is not good. Some refactoring of the hybrid cache implementation may be required.

Hi @TechxGenus, the above PR should fix your issue. However, I wasn't able to reproduce it when `use_mamba_kernels=True `. I think that the simplest way is to move the tensor, as it is only going to affect the first forward.  See the PR for more details. 

@SunMarc I just reproduced this bug with `mamba_kernels` is on.

I have

```
hidden_states.device
>device(type='cuda', index=2)
cache_params.conv_states[self.layer_idx].device
>device(type='cuda', index=0)
```

at here: https://github.com/huggingface/transformers/blob/9578c2597e2d88b6f0b304b5a05864fd613ddcc1/src/transformers/models/jamba/modeling_jamba.py#L672

Traceback:
```
Traceback (most recent call last):
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/aiscuser/.vscode-server/extensions/ms-python.debugpy-2024.10.0-linux-x64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/__main__.py", line 39, in <module>
    cli.main()
  File "/home/aiscuser/.vscode-server/extensions/ms-python.debugpy-2024.10.0-linux-x64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 430, in main
    run()
  File "/home/aiscuser/.vscode-server/extensions/ms-python.debugpy-2024.10.0-linux-x64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 284, in run_file
    runpy.run_path(target, run_name="__main__")
  File "/home/aiscuser/.vscode-server/extensions/ms-python.debugpy-2024.10.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 321, in run_path
    return _run_module_code(code, init_globals, run_name,
  File "/home/aiscuser/.vscode-server/extensions/ms-python.debugpy-2024.10.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 135, in _run_module_code
    _run_code(code, mod_globals, init_globals,
  File "/home/aiscuser/.vscode-server/extensions/ms-python.debugpy-2024.10.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 124, in _run_code
    exec(code, run_globals)
  File "/scratch/MInference/t-1.py", line 25, in <module>
    output = model.generate(input_ids, do_sample=False)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/aiscuser/.local/lib/python3.10/site-packages/transformers/generation/utils.py", line 2024, in generate
    result = self._sample(
  File "/home/aiscuser/.local/lib/python3.10/site-packages/transformers/generation/utils.py", line 2982, in _sample
    outputs = self(**model_inputs, return_dict=True)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/aiscuser/.local/lib/python3.10/site-packages/accelerate/hooks.py", line 169, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/home/aiscuser/.local/lib/python3.10/site-packages/transformers/models/jamba/modeling_jamba.py", line 1482, in forward
    outputs = self.model(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/aiscuser/.local/lib/python3.10/site-packages/transformers/models/jamba/modeling_jamba.py", line 1304, in forward
    layer_outputs = decoder_layer(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/aiscuser/.local/lib/python3.10/site-packages/accelerate/hooks.py", line 169, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/home/aiscuser/.local/lib/python3.10/site-packages/transformers/models/jamba/modeling_jamba.py", line 1040, in forward
    hidden_states = self.mamba(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/aiscuser/.local/lib/python3.10/site-packages/accelerate/hooks.py", line 169, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/home/aiscuser/.local/lib/python3.10/site-packages/transformers/models/jamba/modeling_jamba.py", line 830, in forward
    return self.cuda_kernels_forward(hidden_states, cache_params)
  File "/home/aiscuser/.local/lib/python3.10/site-packages/transformers/models/jamba/modeling_jamba.py", line 672, in cuda_kernels_forward
    hidden_states = causal_conv1d_update(
  File "/home/aiscuser/.local/lib/python3.10/site-packages/causal_conv1d/causal_conv1d_interface.py", line 187, in causal_conv1d_update
    return causal_conv1d_cuda.causal_conv1d_update(
RuntimeError: CUDA error: an illegal memory access was encountered
```

@SunMarc should reopen this issue.

Hey! Can you make sure you have the latest version of transformers (`transformers.__version__ == "4.44.2"`)