fix multi-gpu with static cache

[SunMarc]

What does this PR do?

This PR fixes mutli-gpu inference when using static cache.

Without fix:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument index in method wrapper_CUDA_index_copy_)

With fix:

['Simply put, the theory of relativity states that 1) the speed of light is constant in all inertial reference frames, and 2) the laws of physics are the same for all inertial reference frames.\nThe theory of relativ', 'My favorite all time favorite condiment is ketchup. I love it on everything. I love it on my eggs, my fries, my chicken, my burgers, my hot dogs, my sandwiches, my salads, my p']

Reproduction:

from transformers import LlamaTokenizer, LlamaForCausalLM
import torch

NUM_TOKENS_TO_GENERATE = 40
# Note on `EXPECTED_TEXT_COMPLETION`'s diff: the current value matches the original test if the original test
# was changed to have a cache of 53 tokens (as opposed to 4096), on Ampere GPUs.
EXPECTED_TEXT_COMPLETION = [
    "Simply put, the theory of relativity states that 1) the speed of light is constant in all inertial "
    "reference frames, and 2) the laws of physics are the same for all inertial reference frames.\nThe "
    "theory of relativ",
    "My favorite all time favorite condiment is ketchup. I love it on everything. I love it on my eggs, "
    "my fries, my chicken, my burgers, my hot dogs, my sandwiches, my salads, my p",
]

prompts = [
    "Simply put, the theory of relativity states that ",
    "My favorite all time favorite condiment is ketchup.",
]
tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", pad_token="</s>", padding_side="right")
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf", device_map="balanced", torch_dtype=torch.float16
)
inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)

# Dynamic Cache
generated_ids = model.generate(**inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False)
dynamic_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
# self.assertEqual(EXPECTED_TEXT_COMPLETION, dynamic_text)

# Static Cache
generated_ids = model.generate(
    **inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False, cache_implementation="static"
)
static_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
# We get the same expected output ! 
print(static_text)

[SunMarc]

Same as what we have in hybrid cache class

[ArthurZucker]

Potentially related to #32485, can you add a dummy multi gpu test

[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.

[SunMarc]

Potentially related to #32485, can you add a dummy multi gpu test

This is already tested here for example: https://github.com/huggingface/transformers/actions/runs/10259247273/job/28438834978#:~:text=self%20%3D%20%3Ctests.models.llama.test_modeling_llama.LlamaIntegrationTest%20testMethod%3Dtest_compile_static_cache%3E
and we get the same error that I experienced.

Is this good enough ?

[effortprogrammer]

#32485

@SunMarc

I guess my issue is relevant with this issue.

[SunMarc]

I guess my issue is relevant with this issue.

Could you try this PR to see if this fixes your issue ?

[effortprogrammer]

I guess my issue is relevant with this issue.

Could you try this PR to see if this fixes your issue ?

This solved my issue, too. I will update my issue with this pull request for people who also having this issue.
Thanks for the fix!

[ArthurZucker]

Can you make sure slow CI are passing (mostly compile for llama?) 🤗

[guangy10]

Hi @SunMarc, this PR breaks the export test with static cache. You can reproduce the issue by running RUN_SLOW=1 pytest tests/utils/test_cache_utils.py -k test_static_cache_exportability

[dvrogozh]

Hi @SunMarc, this PR breaks the export test with static cache

@guangy10, @SunMarc : I have noticed this separately and filed #33178 for this. Test fails with RuntimeError: cannot mutate tensors with frozen storage which seems to be reported in pytorch/pytorch#127571. This pytorch issue suggests a workaround (or a fix?) to add .clone(). This works - see #33179. Let's continue discussion in the issue.

[guangy10]

@SunMarc could you help review #33287?