Training of BART slow on TPU - aten ops investigation
Author: phtephanx
Pull request: no
Labels: name: wontfix

Referencing: https://github.com/huggingface/transformers/issues/8339

Versions: 
* `transformers==4.0.1`
* `pytorch==1.7.0`
* `pytorch_xla==1.7.0`

### problem

I've been trying to find out why training of `BartForConditionalGeneration` with `Trainer` is so **slow on TPU**. With slow I mean >30 min/batch (8000 samples) on 8 cores of TPUv3. I can very likely exclude slowdowns on behalf of the host machine.

Following the [xla troubleshooting guide](https://github.com/pytorch/xla/blob/master/TROUBLESHOOTING.md), I made sure that the **input tensors have fixed shape**. Furthermore, I created a `metrics` report which reveals that there occur many context switches between the XLA device and CPU due to:
* aten::_local_scalar_dense
* aten::isnan

I tried to localize the culprits via **debugging** on 1 TPU core and printing the metrics report at each breakpoint.

`aten::isnan` is obviously caused by `torch.isnan` in [L362 of BertEncoderLayer](https://github.com/huggingface/transformers/blob/master/src/transformers/models/bart/modeling_bart.py#L362). I don't know any fix but I'd just turn the condition off for my training.

### questions
1. According to this [issue](https://github.com/pytorch/xla/issues/909), the counter for `aten::_local_scalar_dense` increases "every time the Python code switches from tensor context" to Python scalar context". I couldn't really pinpoint, however, which line causes this. I'm aware that the printing of the metrics report at each breakpoint does so but there are others. Any idea?
2. Just to be on the safe side: `pytorch_xla`'s wrapper `ParallelLoader` for the `DataLoader` doesn't have to do anything with TPU-related slowdown or context-switch in e.g. the `Dataset` instance because the loading still takes place on the host machine and the tensors are put into TPU-queues _afterwards_ as far as I read the code, correct?
3. Is training of `BART` (and probably similar models) running fast on TPU for anyone?

AFAIU @patil-suraj knows more.

Comments:
Hi @phtephanx 

Thank you for troubleshooting the problem. I haven't really dived into this problem yet.

Maybe @LysandreJik (our TPU expert ) will be able to help here once he's back from vacation ;)

Also cc @patrickvonplaten 

Hey @phtephanx,

Thanks for opening this issue! I'm actually very interested in solving this problem as well...did you try executing Bart on TPU with just a simple training loop to see if the slow-down persists? 

Maybe we can work together in a colab to solve this problem. Feel free to dive a bit more into the problem if you feel like it (I'd suggest on a google colab with TPU) and I'm happy to guide you along the way. Else I hope to find some time in mid-January to tackle the problem :-) 

@patrickvonplaten Sorry for the late reply. 

> Thanks for opening this issue! I'm actually very interested in solving this problem as well...did you try executing Bart on TPU with just a simple training loop to see if the slow-down persists?
Maybe we can work together in a colab to solve this problem. Feel free to dive a bit more into the problem if you feel like it (I'd suggest on a google colab with TPU) and I'm happy to guide you along the way. Else I hope to find some time in mid-January to tackle the problem :-)

Would be great if we could cooperate on this because I'm stuck with this for a while! ;-)

I created two minimalistic training notebooks:
* **(I)** [BART on TPU w/o Trainer](https://colab.research.google.com/drive/10crQewhWImt9vHD1UJo-HnzzSgbwwFyA?usp=sharing)
* **(II)** [BART on TPU w/ Trainer](https://colab.research.google.com/drive/1C_8EmDmnisYPLfIkwL7tu-_iVWoPuu1I?usp=sharingv)

**Settings:**
* `BART-base`
* `batch_size=64`
* `gradient_accumulation_steps=1`

**Observations:**
* **(I)** and **(II)** run almost equally fast. Their graph seems to stabilize after approx. 4 steps and subsequently runs at constant throughput of 0.43 [it/s] (~ 2.34 [s/it])
* **(I)** and **(II)** exhibit (almost) the same count of  `aten::isnan` and `aten::_local_scalar_dense`. Thus, no significant number of ops (only 1) w/o XLA lowering is introduced by the `Trainer`

**Conclusions:**
* The throughput of both is really decent which indicates IMO that everything is probably ok with the TPU adaptation of `BART` and `Trainer` even if these two ops w/o XLA lowering occur
* (I couldn't, however, reproduce this throughput on a private GCE VM, so far at all. If you're still interested, I'll take the exact same script and report!)

BTW: I also tried out the training loop of (I) **without** wrapping it into a function and calling `xmp.spawn` on it. The throughput is very low and the graph never stabilizes which was suggested by a `CompileTime` of more than 20 [min].

**Reproducing (II) on GCE VM:**
I conducted a run for **(II)** on **1** core of TPUv3. Apart from the `CompileTime` being considerably larger, which is expected because Colab somehow works instantaneously, the execution time is similar ([metrics_report.txt](https://github.com/huggingface/transformers/files/5762134/tpu_bart_w_trainer_1_core.txt)). It was actually a bit faster: 6 [s] on GCE VM vs. 12 [s] on Colab. Furthermore, I observed that the `CompileTime` is even noticeably larger when using **8** cores.

(It might be that during my actual targeted training of `BART-large` on 8 cores, the stabilization phase of the graph simply takes much longer than for `BART-base` and I never arrived at a stable graph).

Hmm, so it works as expected on a google colab, but not on a private machine? The behavior you described for (I) and (II) seems reasonable to me. It's normal that compilation time is quite high for PyTorch/XLA IMO

I did some runs for `BART-large` on GCE TPU v3-8 with different settings using `Trainer`:

| batch-size | lengths | num-cores | grad-acc | initial-speed [s/it] | final-speed [s/it] | final speed at step |
|------------|---------|-----------|----------|----------------------|--------------------|---------------------|
| 1          | 128     | 1         | 1        | 120                  | 1                  | ~7                   |
| 32         | 128     | 1         | 1        | 318                  | 2.3                | ~6                   |
| 32         | 128     | 8         | 1        | 300                  | 3                  | ~12                  |
| 32         | 128     | 8         | 4        | 440                  | 14.2               | ~20                  |

Extrapolating these numbers to the hparams used by the authors (batch size of 8000) results in the "slow" throughput due to which I opened the issue. I think, everything is fine - one just needs a bigger device like TPU v3-128. Out of scope for me. 

@patrickvonplaten Feel free to close unless there's something else we can discuss / tune.

Hey @phtephanx,

Thanks a lot for posting this, it's very useful! Yeah, I think for now I don't see a big issue either

> Hey @phtephanx,
> 
> Thanks a lot for posting this, it's very useful! Yeah, I think for now I don't see a big issue either

You're welcome ;)

This issue has been automatically marked as stale and been closed because it has not had recent activity. Thank you for your contributions.

If you think this still needs to be addressed please comment on this thread.