fix jamba slow foward for multi-gpu

[SunMarc]

What does this PR do ?

This PR makes the slow forward of Jamba model compatible with multi-gpu setup with use_cache=True. The issue is that the cache is not initialized on the right device. Not sure we want to refactor the cache since the key_cache and and value_cache are handled correctly. This is just specific to ssm_states. Moving the tensor manually should only impact the first forward also. After the first forward, the cache will be correctly initialized.
Fixes #30367

Code:

import os
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig

prompt = "def quicksort(arr):\n"
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-tiny-random")
tokenizer.pad_token = tokenizer.eos_token
config = AutoConfig.from_pretrained("ai21labs/Jamba-tiny-random")
# change the following to switch between fast and slow forward 
config.use_mamba_kernels=False 
model = AutoModelForCausalLM.from_pretrained(
    "ai21labs/Jamba-tiny-random",
    device_map="auto",
    config = config
)
print(model.hf_device_map)
print(model.config)
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]))

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

Seems like it's the best tradeoff:

transformers/src/transformers/models/mamba/modeling_mamba.py

Line 63 in 9eac9c7

class MambaCache:

could be changed but would need to accomodate for every layers? (or is the PKV handled by accelerate separately? )

[SunMarc]

could be changed but would need to accommodate for every layers? (or is the PKV handled by accelerate separately? )

Yes, we would need to accommodate for every layer but that can become quite complicated if we start to introduce layers that are offloaded to the disk and cpu. Hence, I think that it is not worth investing too much effort on that. The PKV is not handled by accelerate separately.