

TITLE
Training of BART slow on TPU - aten ops investigation

COMMENTS
9

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 1
rocket: 0
eyes: 0

BODY
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.