Remote process received SIGTERM on 96 core  tpu-vm during group_text map on datasets
## 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.9.0.dev0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (tpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: yes

### Who can help

t5: @patrickvonplaten
I am using the t5 model with sentencepiece tokenizer trained from scratch :
https://huggingface.co/flax-community/t5-small-dutch/blob/main/tokenizer.json

## Information

Model I am using (T5):

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

Following training the tokenizer for t5 and running t5_mlm from:
https://github.com/huggingface/transformers/blob/master/examples/flax/language-modeling/README.md

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

Pre-training T5 on Dutch oscar deduplicated nl

## To reproduce

Steps to reproduce the behavior:

(unfortunately this script will download and process the complete oscar deduplicated nl corpus which will take ~1 hour, apologies!)

1. setup transformer env on 96 core machine like in the example flax language-modeling example README linked above on the latest master
2. git clone https://huggingface.co/flax-community/t5-small-dutch
3. cd t5-small-dutch
4. ln -s ~/transformers/examples/flax/language-modeling/run_t5_mlm_flax.py run_t5_mlm_flax.py
5. ln -s ~/transformers/examples/flax/language-modeling/t5_tokenizer_model.py t5_tokenizer_model.py
6. ./run_t5_oscar.sh and capture output.
7. In the output, look for SIGTERM

During preprocessing, this text looks like a child process has crashed.

```
https://symbolize.stripped_domain/r/?trace=526cb0,7f98cb33820f,9222bf&map= 
*** SIGTERM received by PID 326216 (TID 326216) on cpu 51 from PID 323401; stack trace: ***
PC: @           0x526cb0  (unknown)  (unknown)
    @     0x7f969b419800        976  (unknown)
    @     0x7f98cb338210  (unknown)  (unknown)
    @           0x9222c0  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=526cb0,7f969b4197ff,7f98cb33820f,9222bf&map=2a762cd764e70bc90ae4c7f9747c08d7:7f968e4d7000-7f969b758280 
E0709 10:28:22.655698  326216 coredump_hook.cc:250] RAW: Remote crash gathering disabled for SIGTERM.
E0709 10:28:22.689142  326216 process_state.cc:771] RAW: Raising signal 15 with default behavior
```

## Expected behavior

No processes are expected to crash.

I added a notebook that shows that the error occurs during the map of the group_texts function on https://huggingface.co/flax-community/t5-base-dutch/blob/main/Load_token_group_dataset.ipynb<|||||>When executing the above notebook with 8 processing threads on my local machine, there are no errors.<|||||>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.