latest TPU VM dies on import of TrainingArguments?
Hi! I'm wondering if  TPU VMs are supported by transformers (I think I saw a tweet that they are, among other bits of documentation :-) ). I'm having some trouble getting some basic scripts running. The error message is pretty odd --- it implies that somewhere within huggingface, a large malloc is being attempted.
```
jack@t1v-n-f3eee39e-w-0:~/transformers/examples/tensorflow/language-modeling$ python3 run_clm.py  --help

tcmalloc: large alloc 113702387720192 bytes == (nil) @  0x7f6a2b8f1680 0x7f6a2b911ff4 0x7f6a2b408309 0x7f6a2b409fb9 0x7f6a2b40a056 0x7f66f4db3659 0x7f66ea7e9954 0x7f6a2bae5b8a 0x7f6a2bae5c91 0x7f6a2b844915 0x7f6a2baea0bf 0x7f6a2b8448b8 0x7f6a2bae95fa 0x7f6a2b6b934c 0x7f6a2b8448b8 0x7f6a2b844983 0x7f6a2b6b9b59 0x7f6a2b6b93da 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=7f6a2b72718b,7f6a2b72720f&map= 
*** SIGABRT received by PID 15450 (TID 15450) on cpu 95 from PID 15450; stack trace: ***
PC: @     0x7f6a2b72718b  (unknown)  raise
    @     0x7f67f8afb1e0        976  (unknown)
    @     0x7f6a2b727210  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=7f6a2b72718b,7f67f8afb1df,7f6a2b72720f&map=ca1b7ab241ee28147b3d590cadb5dc1b:7f67ebdfc000-7f67f8e2eb20 
E0923 20:01:20.325205   15450 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0923 20:01:20.325261   15450 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0923 20:01:20.325271   15450 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0923 20:01:20.325278   15450 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0923 20:01:20.325298   15450 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0923 20:01:20.325316   15450 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0923 20:01:20.325320   15450 coredump_hook.cc:525] RAW: Discarding core.
E0923 20:01:20.329836   15450 process_state.cc:771] RAW: Raising signal 6 with default behavior
```

When I run the tensorflow `tpu-test.py` script from [here](https://cloud.google.com/tpu/docs/tensorflow-quickstart-tpu-vm) all is well.
```
...
PerReplica:{
  0: tf.Tensor(2.0, shape=(), dtype=float32),
  1: tf.Tensor(2.0, shape=(), dtype=float32),
  2: tf.Tensor(2.0, shape=(), dtype=float32),
  3: tf.Tensor(2.0, shape=(), dtype=float32),
  4: tf.Tensor(2.0, shape=(), dtype=float32),
  5: tf.Tensor(2.0, shape=(), dtype=float32),
  6: tf.Tensor(2.0, shape=(), dtype=float32),
  7: tf.Tensor(2.0, shape=(), dtype=float32)
}
```

Interestingly, if I try other backends like `jax`, it also doesn't work
```
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip install optax
git clone https://github.com/google/flax.git ; pip install --user -e flax
```
then

```
python3 run_clm_flax.py --help
tcmalloc: large alloc 33883982246649856 bytes == (nil) @  0x7f641c085680 0x7f641c0a5ff4 0x7f641bb9c309 0x7f641bb9c370 0x7f641bb9c406 0x7f60c5dc55ca 0x7f60c5bac6e4 0x7f60bb5e2954 0x7f641c279b8a 0x7f641c279c91 0x7f641bfd8915 0x7f641c27e0bf 0x7f641bfd88b8 0x7f641c27d5fa 0x7f641be4d34c 0x7f641bfd88b8 0x7f641bfd8983 0x7f641be4db59 0x7f641be4d3da 0x67299f 0x682dcb 0x684321 0x5c3cb0 0x5f257d 0x56fcb6 0x56822a 0x5f6033 0x56ef97 0x5f5e56 0x56a136 0x5f5e56
terminate called after throwing an instance of 'std::bad_alloc'
  what():  std::bad_alloc
https://symbolize.stripped_domain/r/?trace=7f641bebb18b,7f641bebb20f&map= 
*** SIGABRT received by PID 21617 (TID 21617) on cpu 20 from PID 21617; stack trace: ***
PC: @     0x7f641bebb18b  (unknown)  raise
    @     0x7f619eba5c75        976  (unknown)
    @     0x7f641bebb210  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=7f641bebb18b,7f619eba5c74,7f641bebb20f&map=03ff3e3b5ed284dc8d852c5156c3a04c:7f6191131000-7f619eeeaf80 
E0923 20:24:28.285293   21617 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0923 20:24:28.285317   21617 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0923 20:24:28.285325   21617 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0923 20:24:28.285333   21617 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0923 20:24:28.285344   21617 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0923 20:24:28.285352   21617 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0923 20:24:28.285359   21617 coredump_hook.cc:525] RAW: Discarding core.
E0923 20:24:28.289724   21617 process_state.cc:772] RAW: Raising signal 6 with default behavior
Aborted (core dumped)
```

Given that this issue is happening during the imports, I did a binary search of all imports within the flax clm script, and found that this import causes the large malloc.

```
from transformers import TrainingArguments
```
and confirmed that that line by itself is sufficient to cause the issue.

_Notably, though, not /all/ of transformers is subject to this malloc._ I have a different script (which I haven't cleaned up/wasn't intending to share) that seems to run on this same vm with tpu+flax/jax without issue ¯\\_(ツ)_/¯

## Environment info

running `transformers-cli env` actually causes the error I'm describing, so I can't give the exact diagnostics. the issue occurs on the latest, vanilla google TPU VM with a v3-8.

```
>>> import transformers
>>> transformers.__version__
'4.11.0.dev0'
>>> import tensorflow as tf
>>> tf.__version__
'2.6.0'
```

### Who can help

@patrickvonplaten @sgugger

## Information

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

## To reproduce

Steps to reproduce the behavior:

1. Boot up a new TPU VM with a v3-8 TPU
2. install transformers
3. `python3 -c "from transformers import TrainingArguments"`

## Expected behavior

Not a giant malloc/core dump on the TPU VM

The `TrainingArguments` are trying to set the XLA device when imported, so a failure there implies something is going wrong on your setup and that the script can't reach the TPU. Did you follow the steps in the [TPU VM guide](https://cloud.google.com/tpu/docs/pytorch-xla-ug-tpu-vm)? In particular
```
export XRT_TPU_CONFIG="localservice;0;localhost:51011"
```
is necessary for PyTorch.<|||||>Ahh --- thanks for the info! That makes sense. I guess because I wasn't trying to use pytorch (I had only tried jax and tensorflow), I didn't follow that part of the TPU VM guide. Let me see if that works.<|||||>I tried different configurations of `XRT_TPU_CONFIG`, `LD_PRELOAD`, etc. None seemed to support the import.
```
jack@t1v-n-f3eee39e-w-0:~$ echo $XRT_TPU_CONFIG; echo $LD_PRELOAD; cat malloc_import.py ; python3 malloc_import.py ; LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 python3 malloc_import.py
localservice;0;localhost:51011

from transformers import TrainingArguments
free(): invalid pointer
https://symbolize.stripped_domain/r/?trace=7f3a74e8b18b,7f3a74e8b20f,7f3a74edd47b,7f370ee8fad2,7f3a75052b89&map= 
*** SIGABRT received by PID 6132 (TID 6132) on cpu 77 from PID 6132; stack trace: ***
PC: @     0x7f3a74e8b18b  (unknown)  raise
    @     0x7f37ece091e0        976  (unknown)
    @     0x7f3a74e8b210  (unknown)  (unknown)
    @     0x7f3a74edd47c        288  (unknown)
    @     0x7f370ee8fad3         64  _GLOBAL__sub_I_xla_cpu_device.cc
    @     0x7f3a75052b8a  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=7f3a74e8b18b,7f37ece091df,7f3a74e8b20f,7f3a74edd47b,7f370ee8fad2,7f3a75052b89&map=ca1b7ab241ee28147b3d590cadb5dc1b:7f37e010a000-7f37ed13cb20 
E0923 22:13:24.683057    6132 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0923 22:13:24.683077    6132 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0923 22:13:24.683085    6132 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0923 22:13:24.683107    6132 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0923 22:13:24.683115    6132 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0923 22:13:24.683127    6132 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0923 22:13:24.683131    6132 coredump_hook.cc:525] RAW: Discarding core.
E0923 22:13:24.686808    6132 process_state.cc:771] RAW: Raising signal 6 with default behavior
Aborted (core dumped)
tcmalloc: large alloc 500236124160 bytes == (nil) @  0x7f12ec8e0680 0x7f12ec900ff4 0x7f12ec3f7309 0x7f12ec3f8fb9 0x7f12ec3f9056 0x7f0f9cf65659 0x7f0f9299b954 0x7f12ecad4b8a 0x7f12ecad4c91 0x7f12ec833915 0x7f12ecad90bf 0x7f12ec8338b8 0x7f12ecad85fa 0x7f12ec6a834c 0x7f12ec8338b8 0x7f12ec833983 0x7f12ec6a8b59 0x7f12ec6a83da 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=7f12ec71618b,7f12ec71620f&map= 
*** SIGABRT received by PID 6270 (TID 6270) on cpu 95 from PID 6270; stack trace: ***
PC: @     0x7f12ec71618b  (unknown)  raise
    @     0x7f10647b71e0        976  (unknown)
    @     0x7f12ec716210  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=7f12ec71618b,7f10647b71df,7f12ec71620f&map=ca1b7ab241ee28147b3d590cadb5dc1b:7f1057ab8000-7f1064aeab20 
E0923 22:13:26.997615    6270 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0923 22:13:26.997671    6270 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0923 22:13:26.997678    6270 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0923 22:13:26.997690    6270 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0923 22:13:26.997701    6270 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0923 22:13:26.997711    6270 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0923 22:13:26.997717    6270 coredump_hook.cc:525] RAW: Discarding core.
E0923 22:13:27.000837    6270 process_state.cc:771] RAW: Raising signal 6 with default behavior
Aborted (core dumped)
```<|||||>(similar to https://github.com/huggingface/transformers/issues/4668 )<|||||>I also tried running the pytorch xla demo script as described in the google tpu vm tutorial. seems to work fine

```
git clone --recursive https://github.com/pytorch/xla.git
python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
...
| Training Device=xla:0/3 Epoch=1 Step=80 Loss=0.04206 Rate=416.50 GlobalRate=149.87 Time=22:24:32
| Training Device=xla:0/1 Epoch=1 Step=80 Loss=0.04206 Rate=416.24 GlobalRate=150.04 Time=22:24:32
| Training Device=xla:0/7 Epoch=1 Step=80 Loss=0.04206 Rate=416.63 GlobalRate=149.92 Time=22:24:32
| Training Device=xla:0/6 Epoch=1 Step=80 Loss=0.04206 Rate=416.46 GlobalRate=152.41 Time=22:24:32
| Training Device=xla:1/0 Epoch=1 Step=80 Loss=0.04206 Rate=416.20 GlobalRate=136.09 Time=22:24:32
| Training Device=xla:0/5 Epoch=1 Step=80 Loss=0.04206 Rate=416.28 GlobalRate=149.83 Time=22:24:32
| Training Device=xla:0/4 Epoch=1 Step=80 Loss=0.04206 Rate=416.25 GlobalRate=150.01 Time=22:24:32
| Training Device=xla:0/2 Epoch=1 Step=80 Loss=0.04206 Rate=416.04 GlobalRate=149.87 Time=22:24:32
```
and with tmalloc...
```
LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
...      
| Training Device=xla:1/0 Epoch=1 Step=20 Loss=6.10956 Rate=55.42 GlobalRate=50.91 Time=22:27:10
| Training Device=xla:0/4 Epoch=1 Step=20 Loss=6.10956 Rate=56.12 GlobalRate=57.35 Time=22:27:10
| Training Device=xla:0/7 Epoch=1 Step=20 Loss=6.10956 Rate=56.13 GlobalRate=57.42 Time=22:27:10
| Training Device=xla:0/6 Epoch=1 Step=20 Loss=6.10956 Rate=56.07 GlobalRate=56.98 Time=22:27:10
| Training Device=xla:0/2 Epoch=1 Step=20 Loss=6.10956 Rate=56.04 GlobalRate=56.72 Time=22:27:10    
| Training Device=xla:0/3 Epoch=1 Step=20 Loss=6.10956 Rate=55.82 GlobalRate=54.83 Time=22:27:10
| Training Device=xla:0/5 Epoch=1 Step=20 Loss=6.10956 Rate=55.92 GlobalRate=55.79 Time=22:27:10
| Training Device=xla:0/1 Epoch=1 Step=20 Loss=6.10956 Rate=55.94 GlobalRate=55.95 Time=22:27:10
```<|||||>I just tried this on a new TPU VM and can reproduce it. Surprisingly installing `tensorflow-cpu` (`pip install tensorflow-cpu`) magically solves this issue. Gently pinging @skye , since it looks similar to what you reported here https://github.com/huggingface/transformers/issues/12761#issuecomment-915620910

The same is with the `run_clm_flax.py` script, if you install `tensorflow-cpu` and run the script with `USE_TORCH=0`, it should run fine. Here's what I did

```bash
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip install tensorflow-cpu
git clone https://github.com/huggingface/transformers
cd transformers
pip install .
pip install flax optax datasets
cd examples/flax/language-modeling
USE_TORCH=0 python3 run_clm_flax.py     --model_name_or_path gpt2     --dataset_name wikitext     --dataset_config_name wikitext-2-raw-v1     --do_train --do_eval     --block_size 1024     --num_train_epochs 1     --learning_rate 1e-5     --per_device_train_batch_size 8  --per_device_eval_batch_size 8     --output_dir ~/tmp/flax-clm     --overwrite_output_dir
```
`USE_TORCH=0` is not needed if you are using CPU version of torch or if you set `export XRT_TPU_CONFIG="localservice;0;localhost:51011"`

I usually create different `venv` to avoid such framework clashes, so didn't see this issue before.<|||||>BTW: I was also able to solve this without installing `tensorflow-cpu`.

these give the malloc error:
`python3 -c "from transformers import TrainingArguments"`

`XRT_TPU_CONFIG="localservice;0;localhost:51011" python3 -c "from transformers import TrainingArguments"`

but simply doing
`USE_TORCH=0 python3 -c "from transformers import TrainingArguments"`
works.

But, weirdly, you can still run pytorch...
here is `torch_test.py`
```python
import torch
import torch_xla.core.xla_model as xm
dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)
```

`XRT_TPU_CONFIG="localservice;0;localhost:51011" python3 torch_test.py`

returns the expected
```
tensor([[ 0.9678, -1.2321, -1.1490],
        [ 3.3503,  1.5816,  0.9326],
        [-1.1846, -0.7140, -0.4168]], device='xla:1')
```

This all leads me to believe there may be some small piece of transformers, probably unrelated to the models/optimizers/etc., that doesn't play nice with TPUs on torch<|||||>I am getting a similar error with TPU VM, except it is caused by **merely having transformers installed.** Error message below. Error still occurs when I remove transformers from my code. It disappears (and model trains fine) when I uninstall transformers.

**Environment info**

```
transformers: 4.11.3
pytorch: 1.9.1
pytorch-lightning: 1.4.9
```

**Error message**
```
tcmalloc: large alloc 28254628860207104 bytes == (nil) @  0x7f3fc05e3680 0x7f3fc0603ff4 0x7f3fc00fa309 0x7f3fc00fbfb9 0x7f3fc00fc056 0x7f3c700a5659 0x7f3c65adb954 0x7f3fc07d7b8a 0x7f3fc07d7c91 0x7f3fc0536915 0x7f3fc07dc0bf 0x7f3fc05368b8 0x7f3fc07db5fa 0x7f3fc03ab34c 0x7f3fc05368b8 0x7f3fc0536983 0x7f3fc03abb59 0x7f3fc03ab3da 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=7f3fc041918b,7f3fc041920f&map= 
*** SIGABRT received by PID 15664 (TID 15664) on cpu 95 from PID 15664; stack trace: ***
PC: @     0x7f3fc041918b  (unknown)  raise
    @     0x7f3d380271e0        976  (unknown)
    @     0x7f3fc0419210  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=7f3fc041918b,7f3d380271df,7f3fc041920f&map=ca1b7ab241ee28147b3d590cadb5dc1b:7f3d2b328000-7f3d3835ab20 
E1011 23:49:12.957856   15664 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E1011 23:49:12.957878   15664 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E1011 23:49:12.957886   15664 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E1011 23:49:12.957894   15664 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E1011 23:49:12.957905   15664 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E1011 23:49:12.957916   15664 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E1011 23:49:12.957923   15664 coredump_hook.cc:525] RAW: Discarding core.
E1011 23:49:12.960721   15664 process_state.cc:771] RAW: Raising signal 6 with default behavior
Aborted (core dumped)
```
<|||||>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.<|||||>Commenting since this issue is still relevant.<|||||>We still haven't been able to reproduce it and don't have any problem with TPU VMs on our side (the CI even runs tests on it).<|||||>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.