Activation checkpointing for TFGPT2DoubleHeadsModel
### Feature request

Activation checkpointing is implemented for the PyTorch GPT2 model (and different head variants), however, this is not the case for the Tensorflow implementation of GPT2,

### Motivation

A lot of GPU memory is required to finetune GPT2. This is especially the case for TFGPT2DoubleHeadsModel, because different choices (represented by different sequences) are combined in one sample. I think [`tf.recompte_grad`](https://www.tensorflow.org/api_docs/python/tf/recompute_grad) can play a role here.

### Your contribution

I've more experience with PyTorch then Tensorflow, but I could investigate possible solution directions. If it turns out to be easy I could spent time to create a PR, however, help by others is appreciated.
cc @Rocketknight1 @gante<|||||>Hi @visionscaper - this seems like something that could work! We haven't experimented with `tf.recompute_grad` in detail but the core code for training our Keras models is in the `train_step` and `test_step` methods [here](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_tf_utils.py#L1380), so you could try adding it there. Anything you add there will control `model.fit()` for all of our models.

Bear in mind I don't know if this will work - I don't know the exact semantics of `recompute_grad` or if it plays nicely with Graph mode, but if you discover anything or you have any questions feel free to post them here!

<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.