weird large memory usage of mbert model 
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.4.2
- Platform:
- Python version: 3.7
- PyTorch version (GPU?): 1.8 
- Tensorflow version (GPU?):
- Using GPU in script?: - 
- Using distributed or parallel set-up in script?:

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- nlp datasets: [different repo](https://github.com/huggingface/nlp)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->
albert, bert, xlm: @LysandreJik

## Information

I am using mbert model, this is 110 M params, I am testing the pretrianing codes with mt5-small and compare this with mbert, mbert weirdly use a lot of memory, I need to reduce 1/2 of batch_size on the same machine I train my mt5-small model which is 3x larger than mbert model. This is weird that mbert while being 1/3 of mt5-small require large memory.
* I am using run_mlm command
* I run the codes on V100  GPU

## To reproduce

Steps to reproduce the behavior:
python  run_mlm.py     --model_name_or_path bert-base-multilingual-uncased     --dataset_name wikitext     --dataset_config_name wikitext-2-raw-v1     --do_train     --do_eval     --output_dir output --per_device_train_batch_size 88  --fp16 --max_seq_length 128




## Expected behavior

larger batch size should be possible with mbert 
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.<|||||>One possible reason for that is bigger tokenizer cardinality and the fact that HF code predicts all tokens, not only one that are missing, so this operation might dominate overall runtime/memory while doing no useful output.

In my case it was the case and I did a small fix here https://github.com/yurymalkov/transformers/commit/0fe0725c0f7fcc13df698bba1bd01847c1494e43 that ended up in 6X larger batch at my memory usage (1060 GTX) and 3-4X times faster training of a small mbert model, but it causes some tests to fail (I think mainly due to pytorch jit failure which sees named tensor in slice, secondly to missing output of non-masked tokens - though who cares about them?).