Fixes DynamicCache export issues due to control flow and inplace modifications
Author: xadupre
Pull request: yes

# What does this PR do?

Some pieces of codes cannot be exported when the torch is compiling the model but can still be exported with ``strict=False`` which can handle inplace modifications. This PR restores a couple of pieces of code introducing a failure at export time if removed (default values set to None instead of 0). With PR #36311 DynamicCache can be serialized. This PR fixes the remaining export issues.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] 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).
- [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.

<!-- 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: @eustlb
- graph models: @clefourrier

Library:

- flax: @gante and @Rocketknight1
- 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: @Rocketknight1
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page).

cc @gante @zucchini-nlp 

> * We've dropped support for compilation of the whole `generate()` function, as it was barely used and was slowing development. Does `torch.export` require `torch.compile` support? If so, this PR may require additional thinking, to enable support without adding much overhead ([[Generation] remove leftover code from end-to-end compilation #36685](https://github.com/huggingface/transformers/pull/36685) -- worst case scenario, this PR gets reverted).

Right now, there is no easy way to export a LLM with ``torch.export.export`` if the model has control flow depending on the input. A control flow can be rewritten with ``torch.cond`` but technically, [torch.cond](https://github.com/pytorch/pytorch/blob/main/torch/_higher_order_ops/cond.py#L136) is executed like eager mode when the model is not being compiled or exported (pytorch has two flags but they are enabled at the same time). I chose to have a flag and keep both implementations because the one using `torch.cond` is more difficult to read. In that sense, there is no needed support for `torch.compile`, it is just about capturing a model with control flows and not only one execution path. Sometimes, ``torch.export.export`` complains about some inconsistencies in the code (like the definition of an empty cache), it is more strict about types and shapes, but that's all there is to it and these are changes other developpers willing to export can suggest like I just did.

The other option for me would be to patch but patching caches is more difficult than patching models. As I mentioned, this PR goes with #36311, both are needed to export a LLM with a dynamic cache. 

> * A few models overwrite `prepare_inputs_for_generation`. Can we update their logic accordingly?

I'll look into it. Do you want me to update them as well in this PR or to do it separately? My goal is to export many LLM so I'll have to do it anyway. I'll have to extend #36311 to other caches as well.


As an example, the following script [export tiny-llm](https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_tiny_llm.html) tries to export a model using dynamic cache. It fails without this PR.

@xadupre are you happy with the PR, i.e. can I merge it?

> I'll look into it. Do you want me to update them as well in this PR or to do it separately? My goal is to export many LLM so I'll have to do it anyway. I'll have to extend https://github.com/huggingface/transformers/pull/36311 to other caches as well.

Let's make it a separate PR. As you've probably noticed, I may be slow to respond due to workload. PRs open for long often result in merge conflicts, which are annoying for everyone 🤗 

I'm happy. Thanks a lot.