Flax LM scripts doesn't scale-up on TPU Pod.
## Environment info
- `transformers` version: 4.18.0.dev0
- Platform: Linux-5.8.0-1035-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.4.0
- PyTorch version (GPU?): 1.11.0+cu102 (False)
- Tensorflow version (GPU?): 2.8.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.4.1 (cpu)
- Jax version: 0.3.4
- JaxLib version: 0.3.2
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help

Models:

- T5, Pegasus, EncoderDecoder: @patrickvonplaten
- JAX/Flax: @patil-suraj


## Information

Model I am using (Bert, XLNet ...): T5 - Flax

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

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

The T5 MLM as well as any Masking models for flax is working fine in a single TPU node. However, when running the script on a TPU pod, it doesn't scale up correctly.
Assuming the batch size on a single tpu v4-8 node per device is 32 and the global batch size is 128 when we scale it up to a tpu v4-64, then the per device batch size of 32 doesn't work and gives OOM, and it has to be reduced to 4, which means the global batch size will be kept 128.
Furthermore, the training speed remains the same between a single node and a pod.


## To reproduce

Steps to reproduce the behavior:

1. Run any example on https://github.com/huggingface/transformers/tree/main/examples/flax/language-modeling using a single tpu node.
2. Re-run it again using a tpu pod.


## Expected behavior

We are expecting the script scale-up correctly on a tpu pod. 

So, I think we have figured out the problem with flax team :
https://github.com/google/flax/discussions/2017

The main problem that you send to the shard function the global batch for all hosts and nodes.
In the examples, it should send only local batch size per host.<|||||>I believe something like that should work:

replace:
`model_inputs = shard(model_inputs.data)`

with:
```
loca_host_model_inputs = {}
for key, value in model_inputs.data.items():
    loca_host_model_inputs[key] = np.split(model_inputs.data[key], jax.process_count(), axis=0)[jax.process_index()]

model_inputs = shard(loca_host_model_inputs)
```
<|||||>A more simplified  version:

loca_host_model_inputs = {key:np.split(model_inputs.data[key], jax.process_count(), axis=0)[jax.process_index()]  for key, value in model_inputs.data.items()}<|||||>If you think it makes sense, I will create a pull request to fix this issue.<|||||>I think @patil-suraj @borisdayma are working quite a bit with TPU Pods sharding at the moment<|||||>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/main/CONTRIBUTING.md) are likely to be ignored.