# Trainer with TPUs

**URL:** <https://discuss.huggingface.co/t/trainer-with-tpus/16718>\
**Category:** Beginners\
**Created:** [April 12, 2022, 8:59am UTC](https://discuss.huggingface.co/t/trainer-with-tpus/16718 "2022-04-12T08:59:49Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![Kioto97](https://avatars.discourse-cdn.com/v4/letter/k/c77e96/32.png) [@Kioto97](https://discuss.huggingface.co/u/Kioto97)\
**Post date:** [April 12, 2022, 8:59am UTC](https://discuss.huggingface.co/t/trainer-with-tpus/16718/1 "2022-04-12T08:59:49Z")

</div>

Hello, I’ve been trying to train a BERT-based model for SQuAD with the Trainer API, and I would like to use TPUs from COLAB to speed-up the training.  
I have followed the official notebook for training a SQuAD model, and then, I have followed the suggestions on this older topic:  
[Tpu Trainer](https://discuss.huggingface.co/t/when-can-we-expect-tpu-trainer/10353)  
Which says the only thing to edit is to set a max\_length parameter for the padding, which I already did.

So what I did is basically:

- Install the xla library
- Set max\_length for padding
- Set TPU environment in COLAB  
But I don’t get any speedup for training, am I missing something? This is my code:  
[My code](https://colab.research.google.com/drive/17B9KK_D-3sUeobRb9Bfds6BEPcGEnGko?usp=sharing)

Edit: I made some progress! I have to restart the runtime as soon as I install the XLA library, otherwise there are compatibility issues with older versions.  
But still, I didn’t get the speed-up I hoped for, if I do 0.8 iterations/second with a GPU, I only get 3.2 iterations/second with TPU, is it normal?  
Also I get the following data:

> Num examples = 66107  
> Num Epochs = 3  
> Instantaneous batch size per device = 16  
> Total train batch size (w. parallel, distributed & accumulation) = 16

Is it normal? To my understanding, each TPU should have 8 cores, so if I have a batch size of 16 per each devices, the total train batch size should be 128? How do I use all the cores? Thanks

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [April 12, 2022, 4:21pm UTC](https://discuss.huggingface.co/t/trainer-with-tpus/16718/2 "2022-04-12T16:21:02Z")

</div>

Looking at the source code, it might be that you need to specify the `tpu_num_cores`, assuming that the notebook/script in your case does not set this automatically.

> <https://github.com/huggingface/transformers/blob/cc034f72eb6137f4c550e911fba67f8a0e1e98fa/src/transformers/training_args.py#L258-L259>

---

<div class="post-metadata">

**Author:** ![Kioto97](https://avatars.discourse-cdn.com/v4/letter/k/c77e96/32.png) [@Kioto97](https://discuss.huggingface.co/u/Kioto97)\
**Post date:** [April 12, 2022, 8:42pm UTC](https://discuss.huggingface.co/t/trainer-with-tpus/16718/3 "2022-04-12T20:42:00Z")

</div>

Thank you for the answer Bram! Unfortunately that’s not the case, I think by default the parameter is set to 8, since even tuning it manually doesn’t change anything.

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [April 13, 2022, 8:34am UTC](https://discuss.huggingface.co/t/trainer-with-tpus/16718/4 "2022-04-13T08:34:36Z")

</div>

In that case, I am not sure if there are any other parameters that may improve your speed more. Perhaps someone else can help!