examples/seq2seq/finetune.py and BART supports TPU
Author: sshleifer
Pull request: no
Labels: name: Help wanted

name: wontfix

- [ ] test the code on tpu
- [ ] if it doesn't work well: change code as little as possible to get it working.
- [ ] add a script/command that works

Comments:
Do you want to wait for a stable release for the torch-xla? Also shouldn't the end user pass in "(num_tpu_cores=8)" when they are creating the lightning's trainer. It should automatically handle the rest for us i think. Also we have "bloat16" for TPUs as well.

It totally could work out of the box. In which case this issue could be as simple as running a shell command on a tpu machine, seeing that it works well, and then checking in the working command or commenting here :).


Alright cool! Let's see if I can make an attempt to do it! Very new to this examples section of hugging face. 

Hey! So i tried running it on the [colab](https://colab.research.google.com/drive/16q2GWrnZ0Tjg1OxJQUcaWKCWwn3Jh5z0?usp=sharing) first, but it seems there's some error wrt bart model. The traceback it threw, (another issue open as well https://github.com/huggingface/transformers/issues/5915).

To reproduce the same, I executed the finetune.sh with
```bash
!sh finetune.sh \
--data_dir /content/xsum \
--model_name_or_path facebook/bart-base \
--output_dir=xsum_results \
--train_batch_size=2 \
--eval_batch_size=2 \
--num_train_epochs 1 \
--n_tpu_cores 8 \
--tpu_cores 8
```
And modified the args in the shell snip when we invoke the python finetune.py (removed fp16 and gpus to 0)


```
Attempted to call `variable.set_data(tensor)`, but `variable` and `tensor` have incompatible tensor type.
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 119, in _start_fn
    fn(gindex, *args)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 222, in tpu_train
    self.run_pretrain_routine(model)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 1196, in run_pretrain_routine
    False)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/evaluation_loop.py", line 293, in _evaluate
    output = self.evaluation_forward(model, batch, batch_idx, dataloader_idx, test_mode)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/evaluation_loop.py", line 470, in evaluation_forward
    output = model.validation_step(*args)
  File "finetune.py", line 145, in validation_step
    return self._generative_step(batch)
  File "finetune.py", line 176, in _generative_step
    decoder_start_token_id=self.decoder_start_token_id,
  File "/usr/local/lib/python3.6/dist-packages/torch/autograd/grad_mode.py", line 15, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/transformers/generation_utils.py", line 248, in generate
    if self.get_output_embeddings() is None:
  File "/usr/local/lib/python3.6/dist-packages/transformers/modeling_bart.py", line 1113, in get_output_embeddings
    return _make_linear_from_emb(self.model.shared)  # make it on the fly
  File "/usr/local/lib/python3.6/dist-packages/transformers/modeling_bart.py", line 190, in _make_linear_from_emb
    lin_layer.weight.data = emb.weight.data
```

Makes sense. You could try to instantiate
```python
self.lm_head =  _make_linear_from_emb(self.model.shared) 
```
in BartForConditionalGeneration.__init__ 
and then have `get_output_embeddings` return `self.lm_head`.


Alright! Thanks; Will keep you posted 🎉;

Edit -1:

So things seems to work after i make the changes you said. The trouble is that training seems to be frozen on TPUs.

`emb.weight.data` is...
```
tensor([[-0.0370,  0.1117,  0.1829,  ...,  0.2054,  0.0578, -0.0750],
        [ 0.0055, -0.0049, -0.0069,  ..., -0.0030,  0.0038,  0.0087],
        [-0.0448,  0.4604, -0.0604,  ...,  0.1073,  0.0310,  0.0477],
        ...,
        [-0.0138,  0.0278, -0.0467,  ...,  0.0455, -0.0265,  0.0125],
        [-0.0043,  0.0153, -0.0567,  ...,  0.0496,  0.0108, -0.0099],
        [ 0.0053,  0.0324, -0.0179,  ..., -0.0085,  0.0223, -0.0020]],
       device='xla:1')
```
...and `lin_layer.weight.data` is...
```
tensor([[-1.0449e-03,  4.0973e-03, -9.7727e-04,  ...,  8.2363e-04,
         -3.2153e-03,  3.5317e-03],
        [ 2.3644e-03,  3.5527e-03, -1.2428e-03,  ..., -1.0983e-04,
         -2.1916e-03,  5.3099e-05],
        [-4.2492e-03,  3.8183e-04,  3.2527e-03,  ..., -4.4359e-03,
          7.6555e-04, -4.1728e-03],
        ...,
        [-4.3412e-03,  2.8537e-03,  7.9720e-04,  ...,  2.9499e-03,
          2.6357e-03, -3.5283e-03],
        [ 3.7042e-03, -3.0546e-03,  3.9206e-03,  ..., -2.3771e-03,
          4.3551e-03,  1.1703e-04],
        [ 3.5616e-03, -3.1224e-03,  1.3898e-03,  ..., -2.1096e-05,
          5.4077e-04,  1.6183e-03]])
```
... and you try to do...
```
lin_layer.weight.data = emb.weight.data
```

Isn't the problem that they are on different devices?

We can remove the thing which does it on fly(refer to sshleifer comment above). That won't work in case of TPUs.

Good to know. Once you get everything working it would be great to have all the required changes consolidated into one PR.

I am not quite sure that lightning does this optimization when we use multiple TPU cores (only available at the nightly-xla's).  Refer [here](https://github.com/pytorch/xla/issues/1870#issuecomment-623603323).

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.