[FLAX] Core dump using example code

## Environment info
- `transformers` version: 4.8.1
-  `flax` version: 0.3.4
- `python` version: 3.8.5

## Who can help
@patrickvonplaten

## Models:
FLAX - RoBERTa MLM

## Information
Following the official guides for creating VMs and TPUs:
https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm

Following this guide for training RoBERTa on the Norwegian OSCAR training set.
https://github.com/huggingface/transformers/tree/master/examples/flax/language-modeling

I am unable to run the run_mlm_flax.py without getting a core dump. The same happens on the run_clm_flax.py script.



## Error message
```
tcmalloc: large alloc 435677134848 bytes == (nil) @  0x7f61ae7be680 0x7f61ae7deff4 0x7f61ae2d5309 0x7f61ae2d6fb9 0x7f61ae2d7056 0x7f5e637fd659 0x7f5e59233a09 0x7f61ae9b2b8a 0x7f61ae9b2c91 0x7f61ae711915 0x7f61ae9b70bf 0x7f61ae7118b8 0x7f61ae9b65fa 0x7f61ae58634c 0x7f61ae7118b8 0x7f61ae711983 0x7f61ae586b59 0x7f61ae5863da 0x67299f 0x682dcb 0x684321 0x5c3cb0 0x5f257d 0x56fcb6 0x56822a 0x5f6033 0x56ef97 0x5f5e56 0x56a136 0x5f5e56 0x569f5e
terminate called after throwing an instance of 'std::bad_alloc'
  what():  std::bad_alloc
https://symbolize.stripped_domain/r/?trace=7f61ae5f418b,7f61ae5f420f&map= 
*** SIGABRT received by PID 8576 (TID 8576) on cpu 95 from PID 8576; stack trace: ***
PC: @     0x7f61ae5f418b  (unknown)  raise
    @     0x7f5f7fb581e0        976  (unknown)
    @     0x7f61ae5f4210  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=7f61ae5f418b,7f5f7fb581df,7f61ae5f420f&map=ca1b7ab241ee28147b3d590cadb5dc1b:7f5f72e59000-7f5f7fe8bb20 
E0628 20:40:48.745220    8576 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0628 20:40:48.745291    8576 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0628 20:40:48.745305    8576 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0628 20:40:48.745322    8576 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0628 20:40:48.745346    8576 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0628 20:40:48.745362    8576 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0628 20:40:48.745366    8576 coredump_hook.cc:525] RAW: Discarding core.
E0628 20:40:48.749975    8576 process_state.cc:771] RAW: Raising signal 6 with default behavior
Aborted (core dumped)
```

## To reproduce
Follow the guide.

I got the same error, I temporarily fixed it with this [patch](https://katb.in/meow2590).  
<|||||>Thanks a lot. I was actually also able solve this issue by using flax from the git.
git clone https://github.com/google/flax.git
pip install --user -e flax
<|||||>Thanks, I will try it too<|||||>Hey @Wikidepia,

Thanks a lot for your error report. I'm not really sure what is causing the error here. Are you using a TPU VMv3-8? It's important to make sure that jax/flax is installed correctly as explained here: https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm<|||||>Also could you maybe follow the guide here: https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects#how-to-install-relevant-libraries to make sure everything is correctly installed? The code should work as is - please let me know if you continue getting a core dump error.<|||||>Exactly. That turned out to be the issue. Was a bit confused because installing flax with pip install gives me flax version 0.3.4. Installing from git still gives version 0.3.4, but  now it works. Thanks a lot.<|||||>> Hey @Wikidepia,
> 
> Thanks a lot for your error report. I'm not really sure what is causing the error here. Are you using a TPU VMv3-8? It's important to make sure that jax/flax is installed correctly as explained here: https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm

Yes I'm using TPU VM, I might need to update transformers or just create new TPU
Thanks for your help :D <|||||>Also pinging @avital @marcvanzee here in case they have seen something similar before - I don't think it's necessary to clone the flax repo to make it work on TPU VM no? Think one can just `pip install ...` it<|||||>I see you have updated the guide with additional info. Getting the TPU VMs up and running was a bit back and forth. I will do a fresh install of this in a couple of days, and check if I can reproduce this. Thanks a lot @patrickvonplaten  for your response.<|||||>@patrickvonplaten. I have now reinstalled from scratch following the guides, including this https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm. The error I got was not caused by installing Flax from pip instead from source. My mistake. 

The core dump was caused by not running "pip install --upgrade clu". Sorry for the confusion.  <|||||>No worries, thanks for documenting everything here! I'm sure it'll be helpful for others :-)<|||||>Hey @patrickvonplaten , I'm in the TRC program so I could also test some of scripts for TPU VM (before the projects starts on July, 7).<|||||>@stefan-it this would be great, actually. Could you check whether you can set up the libraries correctly according to:
unshuffled_deduplicated_no
and then run these steps: https://github.com/huggingface/transformers/tree/master/examples/flax/language-modeling#masked-language-modeling where instead of running it on **norwegian (no)** (which would take to long) could you run it on **Alemannic (als)**? So simply replacing all occurrences of `unshuffled_deduplicated_no` with `unshuffled_deduplicated_als` ?<|||||>Hi @patrickvonplaten ,

I created a virtual environment (`venv') and followed the installation instructions from [here](https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects#how-to-install-relevant-libraries).

Here are my obserations:

After installing `jax` there's a strange wheel output:

```bash
Building wheels for collected packages: jax                                                                                                                                                 
  Building wheel for jax (setup.py) ... error                                                                                                                                               
  ERROR: Command errored out with exit status 1:                                                                                                                                            
   command: /home/stefan/dev/bin/python3 -u -c 'import sys, setuptools, tokenize; sys.argv[0] = '"'"'/tmp/pip-install-gl33kipc/jax/setup.py'"'"'; __file__='"'"'/tmp/pip-install-gl33kipc/ja
x/setup.py'"'"';f=getattr(tokenize, '"'"'open'"'"', open)(__file__);code=f.read().replace('"'"'\r\n'"'"', '"'"'\n'"'"');f.close();exec(compile(code, __file__, '"'"'exec'"'"'))' bdist_wheel
 -d /tmp/pip-wheel-itpzq66t                                                                                                                                                                        cwd: /tmp/pip-install-gl33kipc/jax/                                                                                                                                                  
  Complete output (6 lines):                                                                                                                                                                  usage: setup.py [global_opts] cmd1 [cmd1_opts] [cmd2 [cmd2_opts] ...]                                                                                                                     
     or: setup.py --help [cmd1 cmd2 ...]                                                                                                                                                    
     or: setup.py --help-commands                                                                                                                                                           
     or: setup.py cmd --help                                                                                                                                                                
                                                                                                                                                                                            
  error: invalid command 'bdist_wheel'                                                                                                                                                      
  ----------------------------------------                                                                                                                                                  
  ERROR: Failed building wheel for jax                                                                                                                                                        Running setup.py clean for jax                                                                                                                                                            
Failed to build jax                                                                                                                                                                         
Installing collected packages: six, absl-py, numpy, opt-einsum, scipy, flatbuffers, jaxlib, libtpu-nightly, jax                                                                             
    Running setup.py install for jax ... done 
```

However, `jax` is installed but:

```
In [1]: import jax
/home/stefan/dev/lib/python3.8/site-packages/jax/__init__.py:27: UserWarning: cloud_tpu_init failed: ModuleNotFoundError("No module named 'requests'")
 This a JAX bug; please report an issue at https://github.com/google/jax/issues
  _warn(f"cloud_tpu_init failed: {repr(exc)}\n This a JAX bug; please report "
```

So there's something wrong with dependency management from `jax`, I manually installed `requests` and it is working.

Then I could run the tokenizer script, which was perfectly working. For the `run_mlm_flax.py` I just found this message:

```bash
Traceback (most recent call last):                                                                                                                                                          
  File "./run_mlm_flax.py", line 319, in <module>                                                                                                                                               f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"                                                                               
  File "/home/stefan/transformers/src/transformers/file_utils.py", line 1641, in wrapper                                                                                                    
    raise ImportError(f"Method `{func.__name__}` requires PyTorch.")
ImportError: Method `device` requires PyTorch.

```

Ok, I did not install PyTorch, and this method is only used in a `logger.info` command, maybe we can write a small logic around it to not use a PyTorch-specific function. I did comment it out, so training could start.

For Alemannic, the following error is thrown after first epoch:

```bash
[18:40:31] - INFO - absl - Starting the local TPU driver.                                                                                                                                   
[18:40:31] - INFO - absl - Unable to initialize backend 'tpu_driver': Not found: Unable to find driver in registry given worker: local://                                                   [18:40:31] - INFO - absl - Unable to initialize backend 'gpu': Not found: Could not find registered platform with name: "cuda". Available platform names are: TPU Interpreter Host          
                                                                                                                                                                                            
[18:40:38] - INFO - absl - A polynomial schedule was set with a non-positive `transition_steps` value; this results in a constant schedule with value `init_value`.                         
/home/stefan/dev/lib/python3.8/site-packages/jax/lib/xla_bridge.py:382: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please upd
ate your code.                                                                                                                                                                              
  warnings.warn(                                                                                                                                                                            
/home/stefan/dev/lib/python3.8/site-packages/jax/lib/xla_bridge.py:369: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update
 your code.                                                                                                                                                                                   warnings.warn(                                                                              
Epoch ... (1/18):   0%|                                                                                                                                              | 0/18 [00:00<?, ?it/s]
Training...:   0%|                                                                                                                                                    | 0/4 [00:00<?, ?it/s]
Training...: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [01:13<00:00, 18.38s/it]
Epoch... (1/18 | Loss: [11.003486 11.003486 11.003486 11.003486 11.003486 11.003486 11.003486                                                                                               
 11.003486], Learning Rate: [9.0000685e-07 9.0000685e-07 9.0000685e-07 9.0000685e-07 9.0000685e-07                                                                                          
 9.0000685e-07 9.0000685e-07 9.0000685e-07])                                                                                                                                                
Epoch ... (1/18):   0%|                                                                                                                                              | 0/18 [01:15<?, ?it/s]
Traceback (most recent call last):
  File "/home/stefan/dev/lib/python3.8/site-packages/numpy/lib/shape_base.py", line 867, in split                                                                                               len(indices_or_sections)                                                                                                                                                                
TypeError: object of type 'int' has no len()                                                                                                                                                
                                                                                                                                                                                            During handling of the above exception, another exception occurred:                                                                                                                         
                                                                                                                                                                                            
Traceback (most recent call last):                                                                                                                                                          
  File "./run_mlm_flax.py", line 622, in <module>                                                                                                                                           
    eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)                                                                                                               
  File "./run_mlm_flax.py", line 268, in generate_batch_splits                                                                                                                              
    batch_idx = np.split(samples_idx, sections_split)                                                                                                                                       
  File "<__array_function__ internals>", line 5, in split                                                                                                                                     File "/home/stefan/dev/lib/python3.8/site-packages/numpy/lib/shape_base.py", line 871, in split                                                                                           
    if N % sections:                                                                                                                                                                        
ZeroDivisionError: integer division or modulo by zero
```

Maybe the training corpus is just too small.

I'm currently training a model for Amharic, and training is running (2 epochs and 1 evaluation phase) :)<|||||>Just a question - maybe it is documented already, but how should we deal with the limited hard disk space? 

I've seen discussions on the Alpha VM TPU Google channel, where it was suggested to use [gcsfuse](https://github.com/GoogleCloudPlatform/gcsfuse), but I've just seen your thread in the Google channel right now, so let's wait :)

<|||||>@stefan-it did you install JAX as follows: 

```
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```

or just via installing `transformers` from source? <|||||>Regarding limited disk space - we are currenty working on a solution :-)<|||||>@patrickvonplaten yeah, I was using the `pip install` command as above (so before installing transformers library).

It seems that disk can be attached during creating of the VM, but not after creation, unfortunately, example [here](https://cloud.google.com/sdk/gcloud/reference/alpha/compute/tpus/tpu-vm/create#--data-disk).<|||||>I also got this core dump on TPUv3-8 and TPUv2-8 VMs. I'll try some of the proposed fixes tomorrow and post an update. 
@patil-suraj <|||||>Alemanic is also a super small dataset so if your batch size is too large it might actually be bigger than the number of examples in the eval set<|||||>The alemanic `als` script really is just a dummy dataset and should be run with a small batch size (2 per device) for testing<|||||>Hi @stefan-it, for limited disk space I found a work around which don't need using any gsutil.
Since TPU-VM has huge RAM (335gb), I mount part of it as a partition, and set HF_HOME to this mount partition **before** running any tokenizers (which cache the preprocessed dataset into $HF_HOME). For more specific:

```bash
mkdir $HOME/hfcache
sudo mount -t tmpfs -o size=128000m tmpfs $HOME/hfcache  # mount 125Gb RAM as disk
export HF_HOME=/home/lethanh/hfcache
```<|||||>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.