GPT2 distributed TPU pre-training using run_clm.py 
Author: mukhtar-algezoli
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 4.2.0dev0
- Platform: Linux-4.9.0-14-amd64-x86_64-with-debian-9.13
- Python version: 3.6.10
- PyTorch version (GPU?): 1.8.0a0+5c3788d (False)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: yes using V3-8 TPUs

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
 
 albert, bert, GPT2, XLM: @LysandreJik 
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Speed and Memory Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 TextGeneration: @TevenLeScao 
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @patrickvonplaten @TevenLeScao
 Blenderbot: @patrickvonplaten
 Bart: @patrickvonplaten
 Marian: @patrickvonplaten
 Pegasus: @patrickvonplaten
 mBART: @patrickvonplaten
 T5: @patrickvonplaten
 Longformer/Reformer: @patrickvonplaten
 TransfoXL/XLNet: @TevenLeScao
 RAG: @patrickvonplaten, @lhoestq
 FSMT: @stas00
 examples/seq2seq: @patil-suraj
 examples/bert-loses-patience: @JetRunner
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 -->
 albert, bert, GPT2, XLM: @LysandreJik 
Trainer: @sgugger


## Information

Model I am using (Bert, XLNet ...): GPT2

The problem arises when using:
* [ 1] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ 1] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1.Create a GCP with pytorch/xla support
2.create a V3-8 TPU
3.install transformers and datasets libs and then run code:

```ruby
python3 transformers/examples/xla_spawn.py --num_cores=8 \
   transformers/examples/language-modeling/run_clm.py \
    --model_name_or_path gpt2 \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
    --do_eval \
    --output_dir /tmp/test-clm
```
then this error comes up:
```ruby
Traceback (most recent call last):
  File "transformers/examples/xla_spawn.py", line 85, in <module>
    main()
  File "transformers/examples/xla_spawn.py", line 81, in main
    xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 394, in spawn
    start_method=start_method)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 205, in start_processes
    while not context.join():
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 160, in join
    exit_code=exitcode
torch.multiprocessing.spawn.ProcessExitedException: process 0 terminated with exit code 17
```
<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

and those exceptions come up before the error:
```ruby
     [[{{node XRTCompile}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

  (1) Resource exhausted: Ran out of memory in memory space hbm. Used 16.61G of 15.98G hbm. Exceeded hbm capacity by 645.48M.

Total hbm usage >= 16.63G:
    reserved         18.00M 
    program          13.93G 
    arguments         2.68G (100.0% utilization)

Output size 192.01M (100.0% utilization); shares 0B with arguments.

Program hbm requirement 13.93G:
    global             4.0K
    HLO temp         13.93G (91.5% utilization: Unpadded (12.75G) Padded (13.93G), 0.0% fragmentation (2.60M))

  Largest program allocations in hbm:

  1. Size: 1.53G
     Shape: pred[8,1023,50257]{1,2,0:T(8,128)E(32)}
     Unpadded size: 392.25M
     Extra memory due to padding: 1.15G (4.0x expansion)
     XLA label: %broadcast.4850.remat3 = pred[8,1023,50257]{1,2,0:T(8,128)E(32)} broadcast(pred[]{:T(256)E(32)} %constant.4065), dimensions={}
     Allocation type: HLO temp
     ==========================

  2. Size: 785.38M
     Shape: bf16[8,1023,50257]{1,2,0:T(8,128
```

## Expected behavior
I don't think this is an OOM problem since I am using 8 cores TPU, so it must be an XLA multiprocessing problem.
<!-- A clear and concise description of what you would expect to happen. -->


Comments:
It clearly states you're out of hbm memory, which is the TPU memory from what Google tells me. I think you have to specify a lower batch size or a lower `block_size` (GPT-2 uses a very big one by default).

@sgugger yup, this was exactly what's wrong, my batch size (per device) is 8 so didn't consider that I am overloading the TPU but totally forgot about block_size (which defaults to 1024). 
Thank you very much.