large alloc when running code flax speech recognition.
Author: pphuc25
Pull request: no

### System Info

When I run the code for training, the problem appears when the warning large alloc, but I do not understand why this happens, this make the training progress really slow, I am currently run on TPU v3-8

The warning be like
```
cmalloc: large alloc 13261799424 bytes == 0x53d90000 @  0x7efe22acb680 0x7efe22aec824 0x7efe22aecb8a 0x7efd7854d2dc 0x7efd7343b7e5 0x7efd734462b0 0x7efd734521f7 0x7efd75b16a37 0x7efd733dc26a 0x7efd733dd56e 0x7efd730a3e04 0x7efd73077cbf 0x5f6939 0x5f7506 0x50b8d3 0x570556 0x5697da 0x5f6ec3 0x5f60b2 0x56ccfc 0x5697da 0x5f6ec3 0x59d21e 0x5f62ae 0x56ccfc 0x5697da 0x5f6ec3 0x5f60b2 0x56ccfc 0x5697da 0x5f6ec3
```

### Reproduction

1. Use code [run_flax_speech_recognition_seq2seq.py](https://github.com/huggingface/transformers/blob/main/examples/flax/speech-recognition/run_flax_speech_recognition_seq2seq.py)

2. Install libraries
  ```
  datasets[audio]>=2.14.0
  jax>=0.3.6
  jaxlib>=0.3.6
  flax>=0.4.1
  optax>=0.0.8
  torch>=1.9.0
  jiwer
  evaluate
  git+https://github.com/huggingface/transformers.git
  tokenizers
  accelerate
  soundfile>=0.12.1
  audiomentations
  ```

3. Run command
  ```
  python3 src/jax/train.py \
	  --model_name_or_path="openai/whisper-tiny" \
	  --dataset_name="mozilla-foundation/common_voice_12_0" \
	  --dataset_config_name="vi" \
	  --language="vi" \
	  --train_split_name="train+validation" \
	  --eval_split_name="test" \
	  --output_dir="./whisper-tiny-vi-flax-example" \
	  --per_device_train_batch_size="16" \
	  --per_device_eval_batch_size="8" \
	  --num_train_epochs="10" \
	  --learning_rate="1e-4" \
	  --accumulate_gradient_steps 2 \
	  --warmup_steps="500" \
	  --logging_steps="25" \
	  --generation_max_length="30" \
	  --preprocessing_num_workers="10" \
	  --dataloader_num_workers="10" \
	  --max_duration_in_seconds="10" \
	  --text_column_name="sentence" \
	  --overwrite_output_dir \
	  --do_train \
	  --do_eval \
	  --predict_with_generate \
	  --push_to_hub \
	  --use_auth_token
  ```

### Expected behavior

It's would be nice if I can have a explain of why this happened, and the way to fix it

Comments:
cc @sanchit-gandhi 

Hey @pphuc25 - you can ignore this warning. It doesn't mean that training will be slower or that there's anything wrong with the script. _c.f._ https://stackoverflow.com/a/52399665. You can change the `tcmalloc` threshold if you want to hide these warnings, e.g. as per https://github.com/huggingface/distil-whisper/blob/914dcdf3919552d5a3826a9d5db99b059ddcc16e/training/flax/distillation_scripts/run_distillation_32_2_timestamped.sh#L3

with my expriments, this make the training really slow on TPU, so I think it's really affect
The way I fix is reinstall library with jax[tpu]

That's a different issue @pphuc25! You **definitely** want to be installing the TPU version of jax / jax-lib on TPU, otherwise you'll have no utilisation of the accelerator cores. This is unrelated to the `tcmalloc` allocation threshold. It simply requires you to install the correct version of JAX:
1. Official instructions: https://jax.readthedocs.io/en/latest/installation.html#pip-installation-google-cloud-tpu
2. More detailed instructions (possibly outdated): https://github.com/huggingface/transformers/tree/main/examples/research_projects/jax-projects#tpu-vm

Oh, I see, thank you for supported