Make StaticCache configurable at model construct time
Author: guangy10
Pull request: yes
Labels: name: ExecuTorch

# What does this PR do?

This PR is to address #32500 for ["Export to ExecuTorch"](https://github.com/huggingface/transformers/issues/32253)

 - Enable the ability to load a model with options to statically config the model using `StaticCache`
 - Create a new integration point for `ExecuTorch` at `transformers/integrations/executorch.py` and hosts the wrapper module class and util `convert_and_export_with_cache` there 
 - Expose the module so that it can be used in other repos
 - Improve test

```
from transformers import convert_and_export_with_cache

model = AutoModelForCausalLM.from_pretrained(
    hf_model_repo,
    attn_implementation="sdpa",
    generation_config=GenerationConfig(
        use_cache=True,
        cache_implementation=cache_implementation,
        max_length=max_cache_len,
        cache_config={
            "batch_size": batch_size,
            "max_cache_len": max_cache_len,
        },
    ),
)

exported_prog = convert_and_export_with_cache(model)

# Further lower the exported program to ExecuTorch with delegates
```

**A real world example/demo:**
The test model `gemma-2b` is naturally exportable via `convert_and_export_with_cache`. 
The test model `gemma-2b` is also lowerable and runnable via `ExecuTorch` w/ 15 tokens/s on `XNNPACK` backend. Checkout https://github.com/pytorch/executorch/pull/4723 for details in `ExecuTorch` repo.


## 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/)? #32500
- [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.

@ArthurZucker
@amyeroberts
@gante


Comments:
Is the test failure relevant?
```
FAILED tests/generation/test_beam_search.py::ConstrainedBeamSearchTest::test_constrained_beam_scorer_finalize - 
ValueError: Each list in `nested_token_ids` can't be a complete subset of another list, but is [[19, 40], [19, 40]].
```
I don't get what this test does and what the error means

> As we've discussed in other places, the cache config should live inside `GenerationConfig`, not in `Config` 🤗 There should be no changes in the model architecture files (e.g. `modeling_gemma.py`), otherwise we expose ourselves to a very lengthy feature propagation cycle -- we would have to edit hundreds of files.
> 
> In a nutshell, to enable this feature:
> 
> 1. `generation_config` can hold and serialize the desired cache config
> 2. `PreTrainedModel.from_pretrained` should accept `generation_config` of type `GenerationConfig`, and overwrite the default `generation_config`
> 3. `generate` uses that field when it is set


I have no preference where the `cache_config` lives, so I can move it to `GenerationConfig` per your feedback.

I explained why some modeling code change is necessary in order to support `torch.export()` [abave here](https://github.com/huggingface/transformers/pull/32830#discussion_r1722305907), and the while work should be similar to support `torch.compile`. Because `torch.export` is more restricted due to its boarder use-cases in a non-python env, there are limitations to the graph input/outputs. The idea is that if we could make such changes in the fundamental models like llama, bert, clip, etc. many other models based on those could leverage the work. @amyeroberts could share more contexts just in case.

> Hey! Allowing configuration of the `staticCache` is most welcome! the other changes IMO do not belong in transformers and are way to heavy, while in `optimum` we could simply wrap arround transformers models as it's something we are more likely to push than change each modeling file

It makes sense to have `optimum` being the place where the actual lowering to on-device happens. This would align it well with other on-device workflows like onnx and tflite. Here in this PR what I'm trying to enable is purely for `torch.export`, similarly to the work for `torch.compile`. One difference between `torch.compile` and `torch.export` is that export is the entry point for the PyTorch 2.x based on-device workflow ie `ExecuTorch`. So getting a transformer model compatible with `torch.export` in a way that can be further lowered to `ExecuTorch` later via  `optimum` is what I'm trying to accomplish in the `transformers` repo.

Totally understand the motivations! 
For now the changes are too "drastic" and constraining for general use-cases. 
The changes done for torch compile were actually also benefic for non compile case and fairly "minimal". 
+ #29753 where we had a LOT of issues with registering a causal mask

@guangy10 my apologies for the acute review above -- the PR already had a positive review, I wanted to make sure it didn't go in at least without exploring alternatives 🤗 

As Arthur wrote, let's try to explore options that do not require modeling changes. We've learned from recent experience that seemingly harmless changes can quickly escalate into difficult issues -- for us and for our users 

> @guangy10 my apologies for the acute review above -- the PR already had a positive review, I wanted to make sure it didn't go in at least without exploring alternatives 🤗
> 
> As Arthur wrote, let's try to explore options that do not require modeling changes. We've learned from recent experience that seemingly harmless changes can quickly escalate into difficult issues -- for us and for our users

Thank you for reviewing this PR. All feedback is appreciated!

@ArthurZucker @gante This PR is updated according to comments about generation config plus what we've discussed in the meeting last week:
 1. Passed `generation_config` via `from_pretrained` and override the default `generation_config`.
 2. Implement the [preferred option](https://docs.google.com/document/d/1A1ieHBVMT0q6yE2q7Mz9-_RVfF5x9xMFu0RKa_vzQSc/edit#bookmark=id.xwm0p2hbn3o2) we discussed last week. That is, continue using the forward() adapter and define it in transformers where lives closest to the source of truth for all transformers models and use in all places where "Export to ExecuTorch" is suitable.
 
I was planning to split this PR into two parts then I realized it's easier to keep it in one so that you can see the full picture about how the mode is configured and exported.

cc: @amyeroberts @qubvel for additional eyes on reviewing.

Okay, all CIs are green 🚀 🚀 

> Looks good! Would be nice to have a real world example

Yeah, there is an e2e demo with PR https://github.com/pytorch/executorch/pull/4723 in ExecuTorch repo. You can see `gemma2b` is working out-of-the-box with ~15 tokens/s using `convert_and_export` from `huggingface/transformers/integrations/executorch.py` we added in this PR .

After this PR is merged, I can start adding an integration test for gemma2b utilizing the exported program (not lowered to ExecuTorch since no ExecuTorch integration in transformers repo) to generate sequences

The examples_tensorflow failure doesn't seem to be relevant

haha I see @amyeroberts was reviewing this at the same time as I was, there are a few things we both suggested/mentioned

(@amyeroberts's comments take precedence, as she's in charge of overall library maintenance) 

Cool! Thanks everyone for all the valuable feedback! Please take a look and let me know if you have further comments

Just updated the copyright headers and rebased to latest main to resolve conflicts. I start seeing new failures on this PR. @amyeroberts @gante Any idea why, or it's irrelevant and safe to merge?

Hi @guangy10 👋 

I wasn't sure what was missing at first, so I went into the code, found the issue, and pushed the fix (tl;dr it was missing the lazy import structure for `transformers.integrations`). I hope you don't mind 🤗 

EDIT: it was also missing an entry for the new doc page in our doc's toc

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_32830). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

> Hi @guangy10 👋
> 
> I wasn't sure what was missing at first, so I went into the code, found the issue, and pushed the fix (tl;dr it was missing the lazy import structure for `transformers.integrations`). I hope you don't mind 🤗
> 
> EDIT: it was also missing an entry for the new doc page in our doc's toc

@gante Thanks! If no more comments, can we merge this PR?

@guangy10 done! Thank you for iterating with us 🤗 

This PR was not universally good - it requires unnecessary `cache_config` for cases when users just want to indicate `cache_implementation=static` and let the rest of cache config be set by `.generate()`. 
suggestion: undo this PR and set custom logic just for the ExecuTorch users.

more details in https://github.com/huggingface/transformers/issues/35026
@gante