Flax Remat for LongT5
Author: KMFODA
Pull request: yes

# 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 

Comments:
_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!