🚨🚨 `GPT2Model` StaticCache support
Author: poedator
Pull request: yes

I copied _update_causal_mask() and _prepare_4d_causal_attention_mask_with_cache_position() from `LlamaModel`

**some tests are still failing:**

1. tests/models/gpt2/test_modeling_gpt2.py::GPT2ModelTest::test_custom_4d_attention_mask
2. test_modeling_vision_encoder_decoder.py::VIT2GPT2Test::test_save_and_load_from_pretrained

both may be linked to attention implementations. So far I was enable to figure out the reasons for failures. I'd appreciate advice or help from the maintainers.

cc: @gante

Comments:
> Not entirely sure it's worth adding ...

I agree in principle, but my friends use it for [Tortoise text-to-speech ](https://github.com/neonbjb/tortoise-tts) and intend to compile it (with modifications for static shapes) to accelerate.

> Let's make sure we test cross attention path with kv cache as I am not even sure it was supported before

I made effort to patch the cross-attention parts of the code as well. The relevant tests seem to pass



@Rocketknight1  @ArthurZucker , could you, please, approve the remaining checks workflows ?

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

@Rocketknight1, could you kindly give feedback on this PR please

I'm not confident in reviewing this one, so cc @gante for final review!

Hi @poedator . Nice patch! I am trying to enable it on opt model by following your codes if you don't mind!

> Hi @poedator . Nice patch! I am trying to enable it on opt model by following your codes if you don't mind!

hi, @jiqing-feng, Good luck in updating OPT! It should be pretty straightforward. Also look at new Cache class -related code is done for Llama. (1) this is where I borrowed from  (2) it may be closer to OPT than gpt2.

sorry @poedator having a look! 

run slow: bert

(the comment above should have triggered slow tests in our CI, testing this workflow :) )

EDIT: https://github.com/huggingface/transformers/actions/runs/13306553429

This comment contains run-slow, running the specified jobs: This comment contains run-slow, running the specified jobs:

models: ['models/bert']
quantizations: [] ...

this is a friendly ping to @Rocketknight1
now that @gante approved, could you finalise you review, please?

Hey! I'm not actually the core maintainer, but if @gante is happy then pinging @arthurzucker @cyrilvallez for core maintainer review

I'll have a look this afternoon!

Also, did you run slow tests by any chance? 🤗

@Cyrilvallez 
> Also, did you run slow tests by any chance? 🤗

I did run the SLOW tests, and all of them passed before I asked for the review. There were couple of exceptions in some other models' tests which did not appear in CI upon submission. 

Will run the SLOW tests again and attach the output link here before calling for another review.

I addressed most of your comments. Please see if you could resolve some or all of them, or give me more guidance 

Status as of 2025-03-01
- Addressed comments from @Cyrilvallez (see above, need resolving)
- rebased to recent main
- ran SLOW tests locally: 2 tests failed, [see logs here](https://gist.github.com/poedator/f4bdaac321440dd4110af0ee89ae9905):
    (1) `tests/models/gpt2/test_modeling_gpt2.py::GPT2ModelLanguageGenerationTest::test_flash_attn_2_generate_padding_left` - has different text generated with FA2. reasons unclear. Generation with eager matches fine
    (2) `tests/models/gpt2/test_modeling_gpt2.py::GPT2ModelLanguageGenerationTest::test_gpt2_sample_max_time` - longer test time, reasons unclear.
    
Also [observed errors in CI](https://app.circleci.com/pipelines/github/huggingface/transformers/120136/workflows/8e6806d4-930a-4cb7-9541-00f5518882a1/jobs/1624129) :
    (3) `tests/models/vision_encoder_decoder/test_modeling_vision_encoder_decoder.py::VIT2GPT2Test::test_save_and_load_from_pretrained` - somehow loaded model uses eager, not sdpa and diverges. but difference is larger than with main. If both models use eager, generation matches perfectly.
    (4) `tests/models/vision_encoder_decoder/test_modeling_vision_encoder_decoder.py::Donut2GPT2Test::test_save_and_load_from_pretrained` - probably same issue    

calling @Cyrilvallez to resolve issues and suggest how to fix the tests problems.

Just answered on the separate comments! Let me know if something is still unclear! For the slow tests, if the particular test is not passing on main as well, it's fine! You should just make sure that this PR does not introduce any new failure!

Hi,  @Cyrilvallez !
it took me some time to fix a bug, but now this PR passes all SLOW tests locally. The CI tests have only one issue related to gated Mistral mode. So could you please, run the remaining workflows and finalize the review?

Looking forward to this PR merged!

BTW, I don't understand why we keep cross attention in gpt2 as I didn't see any GPT2 arch model have encoder hidden states. GPT is a decoder-only model. Could I have an example or an existing model in HF model hub which has encoder_hhidden_states and uses cross attention? If not, can we remove the cross attention code so we can just use static cache?

> Looking forward to this PR merged!
> 
> BTW, I don't understand why we keep cross attention in gpt2 as I didn't see any GPT2 arch model have encoder hidden states. GPT is a decoder-only model. Could I have an example or an existing model in HF model hub which has encoder_hhidden_states and uses cross attention? If not, can we remove the cross attention code so we can just use static cache?

Hi @SunMarc @ArthurZucker @Rocketknight1 . I'd like to hear your opinion, and also want to know what need to be changed before merging this PR. The static cache is essential for torch.compile. We'd like to see gpt2 gain this optimization!

> Hey @poedator! Here is a long due review. [...]

Hi, @Cyrilvallez! Hope that the merge is close. You were correct about inconsistencies with the `decision_transformer`. I fixed those, updated deprecation, rebased to fresh main. Let's see if this is good enough

> Hey @poedator! Thanks a lot for updating! Here are some final thoughts! Also, I still see some logic differences between both models, most notable the mask handling! Let's make sure both models are functionyla equivalent 🤗

I fixed what you directly suggested and also made the cache init sections between the 2 models more similar. 

As for  mask handling and the rest of logic differences -- I am less familiar with this exotic `decision transformer` model, so  I'd be glad if you point me to the specific failing tests so I'd debug the problem. Or feel free to edit this PR directly and take 50% of the glory.
the CI tests seem to fail for unrelated reasons.

> Oh sorry, I thought you wanted to add the same changes and functionalities to decision_transformer as well. If it's not the case, we can just remove the flag saying that it supports static cache (as we would need correct masking for it), then it should be good to go! 🤗

Yes, I was focused on adding support of StaticCache to GPT2, as it is a rather popular model. I will gladly leave the last bit of the Decision Transformer challenge to whoever knows that model better. Please see if this is OK now.

> Alright, LGTM!! 🤗🚀 ... good to go!

@Cyrilvallez @ArthurZucker @gante 
thank you for the valuable inputs and support along the way. Hopefully there are still quite a few people who use GPT2 and who may find this PR useful