Add `accelerate` support for ViLT
Author: younesbelkada
Pull request: yes

# Motivation

Add `bnb` support for ViLT model as it has been asked by a user in https://github.com/TimDettmers/bitsandbytes/issues/14 
. This involved adding `accelerate` support for this model.

# What does this PR do?

Adds `_no_split_modules` attribute at `ViltModel` class to support  loading the model with `device_map=auto`. This also implied adding a `.to` operation inside `ViltLayer`. 
I also redefined `accelerate` tests since for this model the hidden states are not deterministic. However, it is possible to check the correctness of the operation by checking some output attributes such as `logits` or `pooler_output`. 

# Questions

The test `ViltModelIntegrationTest::test_inference_natural_language_visual_reasoning` seem to never pass on my machine (aka even without `_no_split_modules`), is it related to something I am missing? Also it seems that those tests were failing too on the nightly run: https://github.com/huggingface/transformers/runs/7882898294?check_suite_focus=true 

cc @NielsRogge @ArthurZucker @ydshieh 


Comments:
> I also redefined accelerate tests since for this model the hidden states are not deterministic.

Could you elaborate this part in more depth, please?

Sure, 
The accelerate tests defined in `test_modeling_common.py` such as [`test_cpu_offload`](https://github.com/huggingface/transformers/blob/d243112b651f64b87d5d2509ff75042794484c20/tests/test_modeling_common.py#L2321) or [`test_model_parallism`](https://github.com/huggingface/transformers/blob/d243112b651f64b87d5d2509ff75042794484c20/tests/test_modeling_common.py#L2352) compares the first element of the output dictionary returned by the model (aka [`base_output[0]` and `new_output[0]`](https://github.com/huggingface/transformers/blob/d243112b651f64b87d5d2509ff75042794484c20/tests/test_modeling_common.py#L2379) ). This is not correct in the case of `ViLT` since in most of the cases, the first element of the output list is a [`hidden_state`](https://github.com/huggingface/transformers/blob/d243112b651f64b87d5d2509ff75042794484c20/src/transformers/models/vilt/modeling_vilt.py#L868) which appears to be stochastic according to [this line](https://github.com/huggingface/transformers/blob/d243112b651f64b87d5d2509ff75042794484c20/tests/models/vilt/test_modeling_vilt.py#L317). This can be also confirmed also by running several inferences and see that the hidden states are not the same. This is because `VilT samples image tokens from a multinomial distribution` according to the linked line above. 
However, you can still check logits correctness using `logits` and `pooler_output` instead of the hidden states, therefore I added those in this PR. 

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

Thank you, @younesbelkada 

Let me bother you a bit more: why `logits` and `pooler_output` are deterministic while `hidden_states` is not. Strange to me, but I must miss some details in this model.



No worries! 
After deeply looking at the code it appears that the `pooler` and the other heads that are used by the model take as an input the [first element](https://github.com/huggingface/transformers/blob/358fc18613a737f3fbcebc5b2abed43386ff9cbc/src/transformers/models/vilt/modeling_vilt.py#L885) of the variable `hidden_states` (with the latest having the shape `batch_size`x `nb_tokens` x `hidden_dim`). With the hidden states containing an embedding for each image patch ([stochastic](https://github.com/huggingface/transformers/blob/358fc18613a737f3fbcebc5b2abed43386ff9cbc/src/transformers/models/vilt/modeling_vilt.py#L174)) [concatenated together with the text embeddings (non stochastic)](https://github.com/huggingface/transformers/blob/358fc18613a737f3fbcebc5b2abed43386ff9cbc/src/transformers/models/vilt/modeling_vilt.py#L236). Therefore since the heads uses the first text embedding (or CLS token) to perform its downstream task, there is no stochasticity involved there. The stochasticity is only involved on the image patch embedding side.
This can be confirmed as well from the model achitecture (taken from [https://arxiv.org/pdf/2102.03334.pdf](arxiv)):
![Screenshot 2022-08-18 at 17 27 05](https://user-images.githubusercontent.com/49240599/185434093-1b6ca99e-8fb2-4353-9861-a382f675b9ac.png)
Although one can always take the first hidden states to perform the `accelerate` tests, I do think the proposed fix is slightly better since it checks the output of modules that would not be checked if we consider only the first hidden states (e.g. `Pooler`).


Should still be deterministic from my intuition, let me have a look

After looking into it with @ArthurZucker fixing a manual seed on each accelerate test before each forward pass seems to fix the non passing tests. 
We can either keep the changes as they are, or re-define the tests with a seed that is set before each forward pass

I am open to the fix of using seed. It makes the change smaller.
But I am confused by the fact that the daily scheduled CI never report this issue (I should check on slack channels too).

https://github.com/huggingface/transformers/runs/7891383348?check_suite_focus=true

The nightly CI you mentioned above uses nightly torch version (i.e. daily built torch) for which more test failures are expected.

I think it's better for us to figure out why the relevant tests don't fail on scheduled CI for a long period, but here it fails. This seems contradictory.

Could you mention on which hardware you tested, @younesbelkada ?

For the `accelerate` tests it is normal that they were passing since the `_no_split_modules` attribute was never defined in the model class, therefore these tests were never run. 
Regarding the second test maybe it's a hardware issue yes! Let me print you the details:
```
Python version: 3.10.4 (main, Mar 31 2022, 08:41:55) [GCC 7.5.0]
transformers version: 4.22.0.dev0
Torch version: 1.12.1
Cuda available: True
Cuda version: 11.3
CuDNN version: 8302
Number of GPUs available: 2
NCCL version: (2, 10, 3)
DeepSpeed version: None
TensorFlow version: 2.9.1
TF GPUs available: True
Number of TF GPUs available: 2
```

The non passing test pass on my VM with `atol=4e-2` by the way

Thanks a lot for your comments @ydshieh & @NielsRogge 

@NielsRogge : I can confirm this decorator works fine for ViTMAE too, I replaced the function you pointed me by:
```
@set_reproducible(2)
def test_save_load(self):
    super().test_save_load()
```
And the test was passing so I guess we can safely replace stochastic tests with this kind of trick

I would love to have a quick review from @sgugger or @stas00 if possible 🙏 As I think that this decorator can be useful for future models

Thanks 💪  

Just a small comment, in terms of performances I think the decorator can be a little bit improve to only run on the model's forward and not on every single forward pass (if I understand correctly here all the functions named `forward` are affected, instead of just the model's main forward pass) 

I think that @ArthurZucker wanted to enhance the util to make it more memory efficient, will let him decide whenever he feels it ready to merge 💪 

I can confirm that it now fixes the tests for `vit_mae` as suggested by @NielsRogge. Just added a quick fix to only set the seed on the model's forward. LGTM 

Thank you all for your comments!
Reverted the changes related to the context manager and kept the tests to be as simple as possible!
Can confirm all the slow tests pass now (expect for [this test](https://github.com/huggingface/transformers/pull/18683#discussion_r949422087) but this is expected as stated in the comment)
I propose to address the stochasticity issue potentially in a follow-up PR, and keep this PR only for its main goal: add `ViLT` support for `accelerate` 

@sgugger thanks for your comment! And sorry for my late reply
I should have addressed the proposed suggestion in the commit 43b087a and can confirm the slow tests for this model are passing! I hesitated between this solution and having an attribute on `ModelTesterMixin` but I thought that this solution is simpler to understand for future users (the `ModelTesterMixin` has already a lot of attributes)

I can also take care of opening a follow-up PR to update the tests of ViTMAE to make sure these changes are consistent for stochastic models tests inside `transformers`

Perfect thanks! @sgugger 
Should have addressed the changes now! Let me know if you still need some modifications