Fix the initialization of the cache when we have multi gpu

[SunMarc]

What does this PR do ?

Fixes #33287 (comment)
This PR initializes the cache on the right device when we are in multi-gpu setup. Before this PR, we would move the tensors to the right device during update() which created a few issues with export(). This is also a cleaner solution in general.

cc @gante I would love to have a quick feedback from you
cc @ArthurZucker as you were also interested in

Tested with:

RUN_SLOW=True CUDA_VISIBLE_DEVICES=0,1 pytest tests/generation/test_utils.py -k "test_generate_with_static_cache_multi_gpu"

RUN_SLOW=True CUDA_VISIBLE_DEVICES=0,1 pytest tests/generation/test_utils.py -k "test_init_static_cache_multi_gpu"

RUN_SLOW=1 pytest tests/utils/test_cache_utils.py -k test_static_cache_exportability

and:

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)
print(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)

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

[gante]

Nice, thank you for the quick jump to the solution! In general it looks great :D

I would add three more things:

1. make device and layer_device_mapping mutually exclusive OR make the device argument also accept a device map (whichever is the most consistent across HF libraries)
2. make sure we throw an exception somewhere: if a user is using multi GPU, then a layer map has to be passed
3. tests 🤗

[SunMarc]

make sure we throw an exception somewhere: if a user is using multi GPU, then a layer map has to be passed

The only way to know if we are using multi-gpu is with the device_map. Since we are computing the layer_device_map using device_map in get_layer_device_map, that should be good enough no ? If we want to do the check inside the cache, we would have to pass the device_map arg.

[gante]

@SunMarc

That's a good point... 👀 My question at the moment is the following: if a user decides to instantiate a cache manually, is using multi-gpu, and doesn't pass the new argument, how can we let the user know that they should have used the new argument?

[SunMarc]

Refactored a bit + added an integration test to check the device on the cache ! Note that we do the check after generation as we need to initialize the cache, meaning that we need the model. We already have tests for multi-gpu static cache so i don't think we need to add these. Let me know if there are other tests you would like to see.

RUN_SLOW=True CUDA_VISIBLE_DEVICES=0,1 pytest tests/generation/test_utils.py -k "test_generate_with_static_cache_multi_gpu"

That's a good point... 👀 My question at the moment is the following: if a user decides to instantiate a cache manually, is using multi-gpu, and doesn't pass the new argument, how can we let the user know that they should have used the new argument?

When passing the cache to the model in generate, is there a way to post initialize the cache again ? It seems to be the best angle to fix this issue. If one is using multi-gpu, we would have to initialize the model before creating the cache object if we want to pass a relevant args such as the device_map. However, I prefer not passing the device_map as we would have to recreate layer_device_map afterwards in the cache init + not friendly to the user.

[dvrogozh]

Works for me to address #33178.

[gante]

is there a way to post initialize the cache again ?

There is not atm, and I think we would hit a few issues if we do so 🤔 We would have to either:

1. check whether there is data to copy across devices (data checks = not compile-friendly)
2. never copy data when creating tensors in new devices (unexpected behavior from a user perspective)
3. always copy data (we would have to use to to copy data across devices, which is essentially the issue Unbreak torch export with static cache #33287 tries to solve)

---

What about a simple device check at update time? If cache tensor device != new cache data device then throw informative exception

[SunMarc]

What about a simple device check at update time? If cache tensor device != new cache data device then throw informative exception

Sounds good @gante ! I've updated the PR with the check ! Let me know if this is better ! I'll add quick test to check if the warning is raised correctly !

[gante]

LGTM, thank you for iterating 💛

[guangy10]

Thanks for the fix. 🚀 🚀 🚀

[guangy10]

@SunMarc can you include RUN_SLOW=1 pytest tests/utils/test_cache_utils.py -k test_static_cache_exportability in the Test Plan to ensure the export test is fixed?

[SunMarc]

@SunMarc can you include RUN_SLOW=1 pytest tests/utils/test_cache_utils.py -k test_static_cache_exportability in the Test Plan to ensure the export test is fixed?

Test passed !

[guangy10]

@gante @LysandreJik can we prioritize to get this fix merged? I will need this one to unblock ExecuTorch integration. Thanks!

[LysandreJik]

Thanks, for cache related PRs I recommend pinging @ArthurZucker for a review (pinging him now)

[ArthurZucker]

LGTM but I don't think we should check in the update. Maybe in generate?

[ArthurZucker]

I am really unsure this is worth it for us to run this at every forward pass. I know we want to help our users but would need to make sur it does not cost us anything

[gante]

there is a long discussion above, but tl;dr the options are:

• don't warn at all
• check devices in update (this implementation)

with torch.compile, these lines should get ignored anyway when called correctly (at tracing they have the same device). We should benchmark compile to confirm, though. Assuming they have no throughput cost, I think it's a win to have the error

[gante]

We can also wrap update with a try/except, rather than using an if/else

[ArthurZucker]

Let's remove it for now and merge

[ArthurZucker]

cc @SunMarc

[SunMarc]

I ran a quick benchmark to see what the impact on generate with NUM_TOKENS_TO_GENERATE = 2000.

• without the check : ~145.5s
• with the check : ~145s
  -> We have a very small overhead

But yeah, let's remove that for the release and I can do a follow-up pr with either the same patch or use try except as joao suggested. Feel free to merge the PR for the release @ArthurZucker if you are fine with the modification!

[ArthurZucker]

Thanks, tests looks great