[Flax] Add remat (gradient checkpointing)
Author: sanchit-gandhi
Pull request: yes

# What does this PR do?

Adds gradient checkpointing in Flax (_c.f._ #17399). The API currently takes the form of a method:

```python
from transformers import BertConfig, FlaxBertModel

model = FlaxBertModel(BertConfig())
model.enable_gradient_checkpointing()
```

Note: checkpointing has currently only been implemented for FlaxBert. Implementing for all Flax models is a TODO.

TODO:

- [x] Add checkpointing to `init`
- [x] Add checkpointing to `from_pretrained`
- [x] Add model tests for FlaxBert in `test_modeling_flax_bert`
- [x] Decide on API: checkpointing with a kwarg (`gradient_checkpointing=True`) or a method (`model.gradient_checkpointing_enable()`)?
- [x] Add API functionality for remat policies (c.f. https://github.com/google/jax/blob/636345fd67758c19c5345bee2301df34b6f1c540/jax/_src/ad_checkpoint.py#L44)
- [ ] Copy checkpointing logic to all Flax models
- [ ] Move model tests to `test_modeling_flax_common`

## 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?

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.

cc @borisdayma 


Comments:
_The documentation is not available anymore as the PR was closed or merged._

Is there an inconvenient in adding it to all layers?

In my case I used it only on transformers blocks (attention + feed forward).

> Is there an inconvenient in adding it to all layers?

By wrapping `FlaxBertLayer` in a `remat` operation, each Bert layer (attention, intermediate FF, final FF + optional cross-attention layers) has `remat` applied to it:
https://github.com/huggingface/transformers/blob/ea8150a5f932ab28efbe2e7a31fee1ca77c289a5/src/transformers/models/bert/modeling_flax_bert.py#L555
We then use this remat'd layer to construct the Transformer block (layers collection):
https://github.com/huggingface/transformers/blob/ea8150a5f932ab28efbe2e7a31fee1ca77c289a5/src/transformers/models/bert/modeling_flax_bert.py#L559-L562
Meaning that each component of the Bert layer is checkpointed, and that **all** Bert layers in the Transformer block (layers collection) are checkpointed.

Would you like to see `remat` on the embeddings and pooler layers too? Imagine this wouldn't make a huge difference to performance at train time vs just checkpointing the entire Transformer block? 

> Would you like to see `remat` on the embeddings and pooler layers too? Imagine this wouldn't make a huge difference to performance at train time vs just checkpointing the entire Transformer block?

No actually I thought it was on all layers but the way you did is great!

Cool! Once the tests are green, happy to merge it here :-)