Unbreak torch export with static cache
Author: guangy10
Pull request: yes

# What does this PR do?

Fix the issue introduced in #32543, which breaks the `torch.export` use-cases. The fix makes sense for `torch.export` as we are always exporting on "cpu" devices atm, hence avoid calling `.to(device=key_states.device)` on cache for export.

<!-- Remove if not applicable -->

Fixes # (issue)


## Before submitting
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@SunMarc
@ArthurZucker
@dvrogozh


Comments:
Verified `RUN_SLOW=1 pytest tests/utils/test_cache_utils.py -k test_static_cache_exportability` works

Documentation added

The repo consistency failure is irrelevant to this PR. Running `make fix-copies` from latest `main` (`ecd61c62862f925a18b4f063dc17fcaf01826e25`) will generate same patches 

About the CI failure, can you rebase on the latest commit ? This should be solved ! Thanks 

@guangy10 (sorry, I don't want to be your antagonist in your PRs, but this might be a source of future related problems. Please bear with me 🙏 )

In `torch`'s original `to` [docs](https://pytorch.org/docs/stable/generated/torch.Tensor.to.html), we can read: "If the `self` Tensor already has the correct `torch.dtype` and `torch.device`, then `self` is returned.". In other words, a harmless operation, `self = self`.

We are observing that the same doesn't happen in `torch.export`. As such, this is an ad hoc fix, and we will likely need similar fixes in other places. On top of that, because the change is not intuitive versus the base `torch`, it is (correctly) accompanied by a long comment explaining why it is needed. Both problems apply to outside of `transformers` as well.

Ideally, a more scalable solution would be implemented 🤗 I'm wondering whether one of the options below is possible:
1. Handle this edge case in `torch.export` directly. If the device is correctly set before the `to` call this would become a no-op
2. Alternatively, this edge case is documented on `torch.export`, and we can link to your docs instead of having a 4-line comment every time we apply this pattern

WDYT?
(If that is not possible/too hard, of course we will accept this PR)
_______________________________________

On top of this, I think we (`transformers`) can do a better job at setting devices at cache initialization time -- we could initialize the cache given how the model is mapped across devices, as opposed to taking a single `device` argument. 

@SunMarc this should enable us to get rid of the `to` operations after a deprecation cycle, correct?



> @SunMarc this should enable us to get rid of the to operations after a deprecation cycle, correct?

That right ! Yeah it would be better to initialize the cache at initialization time. This should be feasible  as long as we can get the `model` or the `device_map` at the init. 

This WA works for me. Note reply from @bdhirsh here: https://github.com/huggingface/transformers/pull/33179#discussion_r1744250051. They plan to look into this issue on pytorch side to avoid any WA, but it's unclear how long it will take.

Rebased on latest `main`.

> @guangy10 (sorry, I don't want to be your antagonist in your PRs, but this might be a source of future related problems. Please bear with me 🙏 )
> 
> In `torch`'s original `to` [docs](https://pytorch.org/docs/stable/generated/torch.Tensor.to.html), we can read: "If the `self` Tensor already has the correct `torch.dtype` and `torch.device`, then `self` is returned.". In other words, a harmless operation, `self = self`.
> 
> We are observing that the same doesn't happen in `torch.export`. As such, this is an ad hoc fix, and we will likely need similar fixes in other places. On top of that, because the change is not intuitive versus the base `torch`, it is (correctly) accompanied by a long comment explaining why it is needed. Both problems apply to outside of `transformers` as well.
> 
> Ideally, a more scalable solution would be implemented 🤗 I'm wondering whether one of the options below is possible:
> 
> 1. Handle this edge case in `torch.export` directly. If the device is correctly set before the `to` call this would become a no-op
> 2. Alternatively, this edge case is documented on `torch.export`, and we can link to your docs instead of having a 4-line comment every time we apply this pattern
> 
> WDYT? (If that is not possible/too hard, of course we will accept this PR)
> 
> On top of this, I think we (`transformers`) can do a better job at setting devices at cache initialization time -- we could initialize the cache given how the model is mapped across devices, as opposed to taking a single `device` argument.
> 
> @SunMarc this should enable us to get rid of the `to` operations after a deprecation cycle, correct?

Well first it's not an edge case for `torch.export`. Setting the device explicitly doesn't seem to be helpful. As you can see in [#32830](https://github.com/huggingface/transformers/pull/32830/files#r1744318330), the `device` is already set at init time. To unblock I think we should just merge this workaround as it passes both export and compile tests, and left a TODO for proper fix in the future. Your thoughts? CC: @SunMarc @gante 

@guangy10 happy to merge with a TODO, if you believe #33303 wouldn't fix this issue :)

(our current initialization doesn't handle multi-device setups, so that's why we have a `to` operation. #33303 should fix it and remove the need to have a `to` op. Nevertheless, there are other places in our codebase where we assume `to` operations to be harmless, and may clash with `torch.export` 🤔 )

> @guangy10 happy to merge with a TODO, if you believe #33303 wouldn't fix this issue :)
> 
> (our current initialization doesn't handle multi-device setups, so that's why we have a `to` operation. #33303 should fix it and remove the need to have a `to` op. Nevertheless, there are other places in our codebase where we assume `to` operations to be harmless, and may clash with `torch.export` 🤔 )

@gante @SunMarc If the export test is fixed in #33303 and it can be landed any time soon, we can discard this workaround and merge #33303 instead. What matters is that we want to start the `transformers` integration work in `executorch` once #32830 is merged which requires the export test to be working again.

@gante If #33303 needs more time to review, can we merge this workaround to unblock?

#33303 will be merged today!