[`core`/ `GC` / `tests`] Stronger GC tests
Author: younesbelkada
Pull request: yes

# What does this PR do?

While working on the gradient checkpointing refactor I realized that GC was not strongly battle tested on our models. 
I propose a simple way to check if GC has correctly done its effect by simulating a training and checking if all parameters that requires grad have a non-None gradient. 

The good news is that all most used models have a proper working implementation of gradient checkpointing whereas for ~20 architectures this new test fails. 

```bash
FAILED tests/models/autoformer/test_modeling_autoformer.py::AutoformerModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : model.encoder.value_embedding.value_projection.weight has no gradient!
FAILED tests/models/beit/test_modeling_beit.py::BeitModelTest::test_training_gradient_checkpointing - RuntimeError: GET was unable to find an engine to execute this computation
FAILED tests/models/big_bird/test_modeling_big_bird.py::BigBirdModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : bert.pooler.weight has no gradient!
FAILED tests/models/canine/test_modeling_canine.py::CanineModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : canine.projection.conv.weight has no gradient!
FAILED tests/models/clipseg/test_modeling_clipseg.py::CLIPSegModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : clip.logit_scale has no gradient!
FAILED tests/models/data2vec/test_modeling_data2vec_vision.py::Data2VecVisionModelTest::test_training_gradient_checkpointing - RuntimeError: GET was unable to find an engine to execute this computation
FAILED tests/models/dinov2/test_modeling_dinov2.py::Dinov2ModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : dinov2.embeddings.mask_token has no gradient!
FAILED tests/models/dpt/test_modeling_dpt_hybrid.py::DPTModelTest::test_training_gradient_checkpointing - RuntimeError: GET was unable to find an engine to execute this computation
FAILED tests/models/flava/test_modeling_flava.py::FlavaForPreTrainingTest::test_training_gradient_checkpointing - AssertionError: False is not true : flava.text_model.pooler.dense.weight has no gradient!
FAILED tests/models/fnet/test_modeling_fnet.py::FNetModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : fnet.pooler.dense.weight has no gradient!
FAILED tests/models/gpt2/test_modeling_gpt2.py::GPT2ModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : multiple_choice_head.summary.weight has no gradient!
FAILED tests/models/imagegpt/test_modeling_imagegpt.py::ImageGPTModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : transformer.h.0.ln_1.weight has no gradient!
FAILED tests/models/informer/test_modeling_informer.py::InformerModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : model.encoder.embed_positions.weight has no gradient!
FAILED tests/models/layoutlm/test_modeling_layoutlm.py::LayoutLMModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : layoutlm.pooler.dense.weight has no gradient!
FAILED tests/models/lilt/test_modeling_lilt.py::LiltModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : lilt.embeddings.word_embeddings.weight has no gradient!
FAILED tests/models/luke/test_modeling_luke.py::LukeModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : luke.pooler.dense.weight has no gradient!
FAILED tests/models/marian/test_modeling_marian.py::MarianModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : model.encoder.embed_positions.weight has no gradient!
FAILED tests/models/mra/test_modeling_mra.py::MraModelTest::test_training_gradient_checkpointing - ValueError: sequence length must be divisible by the block_size.
FAILED tests/models/pegasus/test_modeling_pegasus.py::PegasusModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : model.encoder.embed_positions.weight has no gradient!
FAILED tests/models/roformer/test_modeling_roformer.py::RoFormerModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : roformer.encoder.embed_positions.weight has no gradient!
FAILED tests/models/seamless_m4t/test_modeling_seamless_m4t.py::SeamlessM4TModelWithSpeechInputTest::test_training_gradient_checkpointing - AssertionError: False is not true : speech_encoder.encoder.layers.1.ffn1_layer_norm.weight has no gradient!
FAILED tests/models/seamless_m4t/test_modeling_seamless_m4t.py::SeamlessM4TModelWithTextInputTest::test_training_gradient_checkpointing - AssertionError: False is not true : t2u_model.model.encoder.layers.0.self_attn.k_proj.weight has no gradient!
FAILED tests/models/time_series_transformer/test_modeling_time_series_transformer.py::TimeSeriesTransformerModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : model.encoder.embed_positions.weight has no gradient!
FAILED tests/models/umt5/test_modeling_umt5.py::UMT5ModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : encoder.block.0.layer.0.SelfAttention.q.weight has no gradient!
FAILED tests/models/visual_bert/test_modeling_visual_bert.py::VisualBertModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : visual_bert.pooler.dense.weight has no gradient!
==================================================================================================================================================== 25 failed, 243 passed, 19 skipped, 42539 deselected, 66 warnings in 59.45s ===============================================
```

I will test this against this commit 9286f0ac3939a7081773fc66480f651a7d6a8404 which is the commit before the first GC refactor commit and report here. If we get the same number of failed tests, this might indicate a silent bug in PyTorch.

Comments:
On the commit shared above, this new test led to the same failures:

```bash
FAILED tests/models/autoformer/test_modeling_autoformer.py::AutoformerModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : model.encoder.value_embedding.value_projection.weight in AutoformerForPrediction has no gradient!
FAILED tests/models/beit/test_modeling_beit.py::BeitModelTest::test_training_gradient_checkpointing - RuntimeError: GET was unable to find an engine to execute this computation
FAILED tests/models/big_bird/test_modeling_big_bird.py::BigBirdModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : bert.pooler.weight in BigBirdForMaskedLM has no gradient!
FAILED tests/models/canine/test_modeling_canine.py::CanineModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : canine.projection.conv.weight in CanineForMultipleChoice has no gradient!
FAILED tests/models/clipseg/test_modeling_clipseg.py::CLIPSegModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : clip.logit_scale in CLIPSegForImageSegmentation has no gradient!
FAILED tests/models/data2vec/test_modeling_data2vec_vision.py::Data2VecVisionModelTest::test_training_gradient_checkpointing - RuntimeError: GET was unable to find an engine to execute this computation
FAILED tests/models/dinov2/test_modeling_dinov2.py::Dinov2ModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : dinov2.embeddings.mask_token in Dinov2ForImageClassification has no gradient!
FAILED tests/models/dpt/test_modeling_dpt_hybrid.py::DPTModelTest::test_training_gradient_checkpointing - RuntimeError: GET was unable to find an engine to execute this computation
FAILED tests/models/flava/test_modeling_flava.py::FlavaForPreTrainingTest::test_training_gradient_checkpointing - AssertionError: False is not true : flava.text_model.pooler.dense.weight in FlavaForPreTraining has no gradient!
FAILED tests/models/fnet/test_modeling_fnet.py::FNetModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : fnet.pooler.dense.weight in FNetForMaskedLM has no gradient!
FAILED tests/models/gpt2/test_modeling_gpt2.py::GPT2ModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : multiple_choice_head.summary.weight in GPT2DoubleHeadsModel has no gradient!
FAILED tests/models/graphormer/test_modeling_graphormer.py::GraphormerModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : encoder.graph_encoder.graph_attn_bias.edge_encoder.weight in GraphormerForGraphClassification has no gradient!
FAILED tests/models/imagegpt/test_modeling_imagegpt.py::ImageGPTModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : transformer.h.0.ln_1.weight in ImageGPTForCausalImageModeling has no gradient!
FAILED tests/models/informer/test_modeling_informer.py::InformerModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : model.encoder.embed_positions.weight in InformerForPrediction has no gradient!
FAILED tests/models/layoutlm/test_modeling_layoutlm.py::LayoutLMModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : layoutlm.pooler.dense.weight in LayoutLMForMaskedLM has no gradient!
FAILED tests/models/lilt/test_modeling_lilt.py::LiltModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : lilt.embeddings.word_embeddings.weight in LiltForSequenceClassification has no gradient!
FAILED tests/models/luke/test_modeling_luke.py::LukeModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : luke.pooler.dense.weight in LukeForMaskedLM has no gradient!
FAILED tests/models/marian/test_modeling_marian.py::MarianModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : model.encoder.embed_positions.weight in MarianMTModel has no gradient!
FAILED tests/models/mra/test_modeling_mra.py::MraModelTest::test_training_gradient_checkpointing - ValueError: sequence length must be divisible by the block_size.
FAILED tests/models/pegasus/test_modeling_pegasus.py::PegasusModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : model.encoder.embed_positions.weight in PegasusForConditionalGeneration has no gradient!
FAILED tests/models/roformer/test_modeling_roformer.py::RoFormerModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : roformer.encoder.embed_positions.weight in RoFormerForMaskedLM has no gradient!
FAILED tests/models/seamless_m4t/test_modeling_seamless_m4t.py::SeamlessM4TModelWithSpeechInputTest::test_training_gradient_checkpointing - AssertionError: False is not true : speech_encoder.encoder.layers.0.ffn1_layer_norm.weight in SeamlessM4TForSpeechToSpeech has no gradient!
FAILED tests/models/seamless_m4t/test_modeling_seamless_m4t.py::SeamlessM4TModelWithTextInputTest::test_training_gradient_checkpointing - AssertionError: False is not true : t2u_model.model.encoder.layers.0.self_attn.k_proj.weight in SeamlessM4TForTextToSpeech has no gradient!
FAILED tests/models/time_series_transformer/test_modeling_time_series_transformer.py::TimeSeriesTransformerModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : model.encoder.embed_positions.weight in TimeSeriesTransformerForPrediction has no gradient!
FAILED tests/models/umt5/test_modeling_umt5.py::UMT5ModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : encoder.block.0.layer.0.SelfAttention.q.weight in UMT5ForConditionalGeneration has no gradient!
FAILED tests/models/visual_bert/test_modeling_visual_bert.py::VisualBertModelTest::test_training_gradient_checkpointing - AssertionError: False is not true : visual_bert.pooler.dense.weight in VisualBertForRegionToPhraseAlignment has no gradient!
=================================================================================================================================================== 26 failed, 242 passed, 19 skipped, 42543 deselected, 160 warnings in 51.80s =========================================================================
```

I propose to skip these tests for the architectures where the test fails, and file an issue on PyTorch with a simple reproducer.

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

Done! I will try to come up with a simple reproducible snippet and file an issue on PyTorch, the failing CI seems unrelated to this PR as it also fails on main. This PR is ready for another round of review!

Thanks for your suggestions @ydshieh , @amyeroberts , this PR is ready for another round of review!

Wow, thank you for the update, I wasn't aware we have to change all files 😅 .

For `test_training_gradient_checkpointing_use_reentrant` and the other 2, maybe we can define

`check_training_gradient_checkpointing(self, gradient_checkpointing_kwargs)`, and the 3 test methods just simply call that method with the desired argument value for `gradient_checkpointing_kwargs` - if that is the only difference between them. Otherwise, LGTM.

Sure @ydshieh , I just refactored the tests as suggested!