No attribute '_mp_fn' when fine-tuning mbart for en-ro translation task using TPU
Author: abedkhooli
Pull request: no

I followed the TPU example in the [examples folder](9https://github.com/huggingface/transformers/tree/master/examples) and found xla_spawn.py calls 
`xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)`  
and [fine-tune.py](https://github.com/huggingface/transformers/blob/master/examples/seq2seq/finetune.py) does not have the "_mp_fn" found in some training scripts.  
I get 
```
Traceback (most recent call last):
  File "examples/xla_spawn.py", line 72, in <module>
    main()
  File "examples/xla_spawn.py", line 68, in main
    xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
AttributeError: module 'finetune' has no attribute '_mp_fn'
``` 
Tried to fix it by adding the _mp_fn:
```
#def _mp_fn(index):
    # For xla_spawn (TPUs)
#    pass
    #main()
``` 
with and without args `main(args)` but neither worked.

Comments:
@abedkhooli  Could I have the command you ran + environment details so that I can try to replicate this?
Thanks!


Thanks @sshleifer for looking into this. 
TPU type: TPU v2 which is 8 cores, 64 GB (using Google Colab)
```
%%bash
export ENRO_DIR='/content/wmt_en_ro' # Download instructions above
#export WANDB_PROJECT="MT" # optional
export MAX_LEN=32
export BS=8
cd /content/transformers
./mbart_enro.sh
```
mbart_enro.sh:
```
#!/usr/bin/env bash
export PYTHONPATH="../":"${PYTHONPATH}"

python examples/xla_spawn.py --num_cores 8 \
	  examples/seq2seq/finetune.py \
    --learning_rate=3e-5 \
    --fp16 \
    --do_train \
    --val_check_interval=0.25 \
    --adam_eps 1e-06 \
    --num_train_epochs 1 --src_lang en_XX --tgt_lang ro_RO \
    --data_dir $ENRO_DIR \
    --max_source_length $MAX_LEN --max_target_length $MAX_LEN --val_max_target_length $MAX_LEN --test_max_target_length $MAX_LEN \
    --train_batch_size=$BS --eval_batch_size=$BS \
    --task translation \
    --warmup_steps 500 \
    --freeze_embeds \
    --model_name_or_path=facebook/mbart-large-cc25 \
    --output_dir enro_finetune_baseline \
    --label_smoothing 0.1 \
    --fp16_opt_level=O1 --sortish_sampler --n_train 5000 --n_val 500 \
    "$@"
```
I believe the issue is adding the correct _mp_fn to examples/seq2seq/finetune.py that matches the main() call (I am not an experienced coder :-)).


I see a related [PR#5960](https://github.com/huggingface/transformers/pull/5960) - does that mean moving away from [xla_spawn](https://github.com/huggingface/transformers/blob/master/examples/xla_spawn.py) ?

That PR is stalled, I am open to using any tpu implementation that works!


If using [xla_spawn](https://github.com/huggingface/transformers/blob/master/examples/xla_spawn.py), and adding  _mp_fn(..) to [finetune.py](https://github.com/huggingface/transformers/blob/master/examples/seq2seq/finetune.py), how should it (_mp_fn) be defined?

I don't know, great question. Maybe @LysandreJik would know the answer.

`_mp_fn(index)` should simply be an entry point to your script that leverages `transformers.Trainer`. You can see examples of it [here](https://github.com/huggingface/transformers/blob/master/examples/text-classification/run_glue.py).

Please note that we implemented this to mimic torch's `torch.distributed.launch`. I have no idea how this would work with a `pytorch-lightning` implementation. Doesn't pytorch-lightning have its own way of managing TPU training?

The main() function in [finetune.py](https://github.com/huggingface/transformers/blob/master/examples/seq2seq/finetune.py#L350) takes arguments, so _mp_fn(index) signature won't work. 
```
def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()
```
`Exception in device=TPU:0: main() missing 1 required positional argument: 'args'`

Right, but even if you manage to make it work with the args, `finetune.py` is using pytorch-lightning so it won't work with `xla_spawn.py`. You can check the [pytorch-lightning docs](https://pytorch-lightning.readthedocs.io/en/latest/tpu.html) to see how to run on TPU.

So, [lightning_base.py](https://github.com/huggingface/transformers/blob/master/examples/lightning_base.py#L165) is not ready for TPU yet.

This is now supported by `Seq2SeqTrainer` which doesn't use PL.
See https://github.com/huggingface/transformers/blob/master/examples/seq2seq/builtin_trainer/finetune_tpu.sh