Repository: transformers
Issue: FLAX core dump error on CloudTPU when running run_clm_flax.py
Labels: WIP

Hi, I'm having a weird problem trying to train a gpt-neo model from scratch on a v3-8 cloud TPU. Something similar to the [closed issue here](https://github.com/huggingface/transformers/issues/12404). Getting: 
```
https://symbolize.stripped_domain/r/?trace=7fb5dbf8a3f4,7fb5dbfe020f,7f&map= 
*** SIGTERM received by PID 64823 (TID 64823) on cpu 26 from PID 63364; stack trace: ***                                                                                                                                     | 0/1 [00:00<?, ?ba/s]
PC: @     0x7fb5dbf8a3f4  (unknown)  do_futex_wait.constprop.0
    @     0x7fb52fa377ed        976  (unknown)
    @     0x7fb5dbfe0210  440138896  (unknown)                                                                                                                                                                               | 0/1 [00:00<?, ?ba/s]
    @               0x80  (unknown)  (unknown)                                                                                                                                                                               | 0/1 [00:00<?, ?ba/s]
https://symbolize.stripped_domain/r/?trace=7fb5dbf8a3f4,7fb52fa377ec,7fb5dbfe020f,7f&map=44c8b163be936ec2996e56972aa94d48:7fb521e7d000-7fb52fd90330 
E1122 14:13:36.933620   64823 coredump_hook.cc:255] RAW: Remote crash gathering disabled for SIGTERM.                                                                                                                        | 0/1 [00:00<?, ?ba/s]
E1122 14:13:36.960024   64823 process_state.cc:776] RAW: Raising signal 15 with default behavior
```
randomly during preprocessing/loading the dataset. 

The env is clean, setup according to the Quickstart Flax guide from google's [help page](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm), and as well from [here](https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects#how-to-install-relevant-libraries). Jax is installed okay, sees 8 TPUs. I tried the standard pip install as well as the local install as some people suggested in the [issue](https://github.com/huggingface/transformers/issues/12404) above, still getting the same behavior. 

This error does **not** kill the training. 
So, question number 1 would be **how to get rid of this error ?**

Something else happens that _might_ be related: Running a dummy 300MB Wiki dataset for training only produces the error above, but training progresses. However, when running the full 40GB dataset, at a point during the first epoch I get:

``list([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, .... (many 1s) .. 1, 1, 1])]]' of type <class 'numpy.ndarray'> is not a valid JAX type.``

This error kills the training. I've found this related [issue](https://github.com/huggingface/transformers/issues/12502), but the last suggestion of  increasing ``max_seq_len`` does not apply here, as the preprocessing should automatically concatenate and cut the model len (and it is set in the config file). The dataset itself is clean, does not contain long words or chars or anything weird.

Thus, question 2: **Any pointers on how to solve this second error?**

Unfortunately I cannot share the dataset as it's private :disappointed: so I don't know how to help reproduce this error. There are 2 questions in this single issue as maybe there's a chance they are related (?). 

Thanks a bunch!

Update: [here is the output of the run_clm_flax.py](https://wtools.io/paste-code/b7US). Because there's a limit on how much you can paste online, I've deleted a few chunks of repeating lines in the output.