❓ [TPU] [Trainer] Moving model to device before setting optimizer slow the training 
# ❓ Questions & Help

On master, after applying the fix of #4450, the training on 8 TPU cores is much slower.

* Before the fix : **20 min / epoch (8 iterations / s)**
* After the fix : **1h30 / epoch (2~3 iterations / s)**

Of course the training before the fix was not working (loss was not decreasing).

But this slow is not expected : for the TF2 equivalent with the same dataset, it takes **20 min / epoch**.

---

Anyone meeting the same problem ?

I have the same problem, but I'm not sure if it's really a problem. I thought the extreme speed before the fix was because it wasn't training properly, and that the slower speed now is supposed to be this way, but that's just my guess.<|||||>Hi, you're right @LeonieWeissweiler. Previous to that fix, the optimizer wasn't actually adjusting weights, resulting in a major speed-up (but the script in itself wasn't working).

@Colanim, do you mind specifying what exactly you're training on? When training on TPU there's a lot you should take into account: batch size, sequence length, number of cores being the most important. Do you mind giving a bit of context as to what you're trying to run?

 We've also merged this PR https://github.com/huggingface/transformers/pull/4467 which solves quite a few issues with the TPU training. Please make sure to install from source to benefit from that commit.

From my tests, on TPU with 8 cores (v3-8), on MNLI I reach 22 minutes/epoch with a batch size of 8, but **6 minutes/epoch** (with a 2 minute tracing, that isn't necessary for the following epochs) with a batch size of 128 (which does train with a final accuracy of 81% using `bert-base-cased`, single epoch). <|||||>I'm training ELECTRA for Extractive Text Summarization.

I will try to increase the batch size and see the results, thanks for the pointer.

What bother me is that I used TFElectra and I could train the model at ~8 iterations per sec.

Same model and same hyper-parameters on Pytorch and it's slower.

But I realized that in my model, since it's extractive summarization, I'm extracting the [CLS] representation of each sentence. These CLS position varies from sample to sample. Maybe that's why it's slower on pytorch-xla ?<|||||>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.
