Flax Remat for LongT5
# What does this PR do?

Following on from the discussion in https://github.com/huggingface/transformers/issues/17399 this PR adds remat (flax's gradient checkpointing) to LongT5. Most of the code is copied from https://github.com/huggingface/transformers/pull/17843/files and then altered to be applicable for LongT5. 

## 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#start-contributing-pull-requests),
      Pull Request section?
- [x] 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?
@sanchit-gandhi or @patrickvonplaten 
_The documentation is not available anymore as the PR was closed or merged._<|||||>Hey @KMFODA, looks like a great start on this PR! I would advocate for adding gradient checkpointing to Flax T5 first. Once we're confident that the implementation is correct, we can run the command `make fix-copies` to have the changes reflected across all Flax T5 based models.

What remains to do is wrapping the Flax T5 Block Collection in the `remat` transformation when `gradient_checkpointing` is passed as `True`. This is analogous to wrapping the Flax Bert Layer Collection in the `remat` operation in the reference PR https://github.com/huggingface/transformers/pull/17843.

We just need to update the following lines:
https://github.com/huggingface/transformers/blob/f0982682bd6fd0b438dda79ec45f3a8fac83a985/src/transformers/models/t5/modeling_flax_t5.py#L648-L651

To reflect what is done in Flax Bert:
https://github.com/huggingface/transformers/blob/f0982682bd6fd0b438dda79ec45f3a8fac83a985/src/transformers/models/bert/modeling_flax_bert.py#L553-L558

As a pointer, the `static_argnums` kwarg is used to denote any boolean arguments passed to the `remat` layer. Inspecting the `__call__` method of FlaxT5LayerCollection, we see that the final four arguments are booleans:
https://github.com/huggingface/transformers/blob/f0982682bd6fd0b438dda79ec45f3a8fac83a985/src/transformers/models/t5/modeling_flax_t5.py#L616-L628
Thus, we have `static_argnums=(6,7,8,9)` (corresponding to `output_attentions`, `return_dict`, `deterministic` and `init_cache`).<|||||>Hey @sanchit-gandhi. Sorry this is taking so long, adding your changes was relatively easy but I'm a bit stuck trying to pass a few failing tests which I believe are caused by my addition of `key_value_states`, `encoder_hidden_states`, `encoder_attention_mask`, `use_cache` and `init_cache` into most calls as I remember in your PR you mentioned that remat doesn't take kwargs.

I believe somewhere in the code `key_value_states` is being incorrectly set to a non `None` figure which is causing some of these errors but I'm not sure where. <|||||>of course that makes sense. Apologies for the misunderstanding. I'll work on the gradient checkpointing part using your suggestions and remove the `key_value_states` for this PR.<|||||>hey @sanchit-gandhi. I beleive this is now ready for review. The PR passes all the tests except ones related to inconsistencies between t5 and long_t5. If you're happy with this let me know and then I think running `make fix-copies` should fix the remaining errors?<|||||>Thanks @sanchit-gandhi for all the helpful comments. Addressed them all and ran `make fix-copies`. Hopefully these changes should be reflected properly for `LongT5 `as well now.<|||||>Thanks for this PR @KMFODA! Awesome stuff, excited to hear how you get on with improved memory usage with the aid of checkpointing! <|||||>Amazing, will keep you posted. Thanks for all the help getting this merged!