Make `Cache` a subclass of `torch.Tensor`
Author: IlyasMoutawwakil
Pull request: yes

# What does this PR do?

Both torch script tracing and torch dynamo/fx have restrictions on input types (torch script has more) which makes the export fail as one torch module (the model) is passing another (the cache) around as its input. Having Cache be a subclass of torch.Tensor bypasses these issues and imo makes more sense as the Cache class has no forward and is just a container of torch tensors.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] 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.
- [ ] 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).
- [ ] 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.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
@guangy10 as requested on Slack, have a look if you're available 🙏 

For the correctness testing, no extensive testing, but we do have some correctness guarantee for supported models `test_export_static_cache` ([pointer](https://github.com/search?q=repo%3Ahuggingface%2Ftransformers+test_export_static_cache&type=code)). Can you run slow tests on this PR?

Also I'm not exactly sure if the `StaticCache` will be functioning as expected. Because with `nn.Module` the Cache is registered as a mutable buffer and lifted to the graph input during `export`. I'm curious how it works with tensor subclass. It seems like tensor subclasses do not directly support buffer registration like nn.Module does. Can we compare the graph between using the `nn.Module` solution vs. the tensor subclass solution.

Alternatively, since the motivation is to handle the legacy torch script tracing (I assume the traffic to this path will be lower and lower over time), would it be a cleaner separation if we create a dedicated Cache subclass for it but keeping the one for pytorch2.0+ as `nn.Module`? No need to maintain compatibility to the torch script solution.

> Question: Cache object holds a list of tensors, usually with a pair of tensors per layer. On some cases, we can have different tensors of a cache on different devices. Would this conflict with the new inheritance?

Shouldn't be an issue as we're not using the `_make_subclass()` but rather `_make_wrapper_subclass()`, the difference is explained by @albanD:

> These two functions do quite different things. The main difference is that when you do _make_subclass(), the current object is a honest to goodness Tensor with data in its storage and everything. When you do _make_wrapper_subclass(), the current object has no data and it is expected that some field on the Tensor will be another Tensor (hence the outer one being called wrapper) that contains real data.
in https://dev-discuss.pytorch.org/t/whats-the-difference-between-torch-tensor-make-subclass-and-torch-tensor-make-wrapper-subclass/1839

One example is the [QuantizedTensor subclass](https://docs.google.com/presentation/d/1piuv9nBzyoqdH49D1SoE5OZUPSMpOOFqfSKOhr-ab2c/edit#slide=id.p1) which has two dtypes (a public one `qt.dtype` and an internal one `qt._data.dtype`

> Have you confirmed that slow llama tests and slow cache tests have no regressions with respect to main? (RUN_SLOW=1 py.test tests/models/llama/test_modeling_lama.py -vv and RUN_SLOW=1 py.test tests/utils/test_cache_utils.py -vv)
Have you confirmed that llama + static cache + compilation preserves throughput? (can share a script if needed :) )

Running them right now (btw is there a way to trigger them on the CI ?), I was only running llama fast tests and llama+executorch integration tests.

Edit: confirmed these two tests fail on main as well

Running `RUN_SLOW=1 pytest tests/models/llama/test_modeling_llama.py -vv` give two errors which I guess are related the machine I'm testing on (A100 vs the A10 that's used in the CI) ;
```
FAILED tests/models/llama/test_modeling_llama.py::LlamaIntegrationTest::test_llama_3_1_hard - AssertionError: 'Tell[74 chars]ical social and political upheaval in France t[557 chars]s.\n' != 'Tell[74 chars]ical political...
FAILED tests/models/llama/test_modeling_llama.py::LlamaIntegrationTest::test_model_7b_logits_bf16 - AssertionError: False is not true
```
in the first `social and political` is reversed to `political and social `:
```
E       AssertionError: 'Tell[74 chars]ical social and political upheaval in France t[557 chars]s.\n' != 'Tell[74 chars]ical political and social upheaval in France t[557 chars]s.\n'
E       Diff is 1259 characters long. Set self.maxDiff to None to see it.
```
in the second the assertion is not verbose enough:
```
>       self.assertTrue(
            torch.allclose(
                EXPECTED_MEAN[self.cuda_compute_capability_major_version].to(torch_device),
                out.logits.float().mean(-1),
                atol=1e-2,
                rtol=1e-2
            )
        )
E       AssertionError: False is not true
```
adding some verbosity:
```
E       AssertionError: False is not true : Expected: tensor([[-6.5208, -4.1218, -4.9377, -3.2536,  0.8127, -2.9811,  1.2918, -3.3848]],
E              device='cuda:0')
E       Got: tensor([[-6.5081, -4.1175, -4.9761, -3.1678,  0.8199, -3.0029,  1.2809, -3.3309]],
E              device='cuda:0')
```

Thanks everyone, I'm replacing this PR with https://github.com/huggingface/transformers/pull/35873 that's less restrictive.

> Thanks everyone, I'm replacing this PR with #35873 that's less restrictive.

Yeah, this one looks much cleaner. Do you mind rerun the slow export/executorch tests on this PR? 

After that, can you run cross-repo integration tests in Optimum (running it locally is fine) ? @echarlaix recently moved ExecuTorch to a new repo `huggingface/optimum-executorch`, however, all documentations/tutorials are deleted unintentionally after the move.

It used to be as simple as:
```
pip install optimum[exporters-executorch]
```
Override the installed `transformers` version to your dev version including this PR, then simply run
```
pytest executorch/*/test_*.py -s -vvvv --durations=0
```

@echarlaix @michaelbenayoun can you guide @IlyasMoutawwakil how to run executorch e2e tests?

> Do you mind rerun the slow export/executorch tests on this PR?

They pass, I made sure both llama and cache slow tests are passing on this PR before dropping it so that it could be used for future references when subclassing the tensor class.

> After that, can you run cross-repo integration tests in Optimum (running it locally is fine) ?

Can do that later.