# Tutorials for using Colab TPUs with Huggingface Transformers?

**URL:** <https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970>\
**Category:** 🤗Transformers\
**Created:** [November 10, 2020, 8:01pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970 "2020-11-10T20:01:00Z")\
**Posts on this page:** 17\
**Page:** 1

<div class="post-metadata">

**Author:** ![facehugger2020](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/facehugger2020/32/3052_2.png) [@facehugger2020](https://discuss.huggingface.co/u/facehugger2020)\
**Post date:** [November 10, 2020, 8:01pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/1 "2020-11-10T20:01:00Z")

</div>

I looking for an easy-to-follow tutorial for using Huggingface Transformer models (e.g. BERT) in PyTorch on Google Colab with TPUs. I found guides about XLA, but they are largely centered around TensorFlow.

Any help would be appreciated.

---

<div class="post-metadata">

**Author:** ![thomwolf](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/thomwolf/32/2198_2.png) [@thomwolf](https://discuss.huggingface.co/u/thomwolf)\
**Post date:** [November 12, 2020, 12:55pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/2 "2020-11-12T12:55:17Z")

</div>

There are a few contributed notebooks which might help you here: [https://github.com/huggingface/transformers/tree/master/notebooks](https://github.com/huggingface/transformers/tree/master/notebooks)

For instance this one by @valhalla is about training T5 on TPU in PyTorch: [https://colab.research.google.com/github/patil-suraj/exploring-T5/blob/master/T5\_on\_TPU.ipynb#scrollTo=QLGiFCDqvuil](https://colab.research.google.com/github/patil-suraj/exploring-T5/blob/master/T5_on_TPU.ipynb#scrollTo=QLGiFCDqvuil)

---

<div class="post-metadata">

**Author:** ![facehugger2020](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/facehugger2020/32/3052_2.png) [@facehugger2020](https://discuss.huggingface.co/u/facehugger2020)\
**Post date:** [December 6, 2020, 11:21pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/3 "2020-12-06T23:21:14Z")

</div>

I still cannot get any HuggingFace Tranformer model to train with a Google Colab TPU.

I tried out the notebook mentioned above illustrating [T5 training on TPU](https://colab.research.google.com/github/patil-suraj/exploring-T5/blob/master/T5_on_TPU.ipynb#scrollTo=QLGiFCDqvuil), but it uses the `Trainer` API and the XLA code is very ad hoc.

I also tried a more principled approach based on [an article by a PyTorch engineer](https://medium.com/pytorch/get-started-with-pytorch-cloud-tpus-and-colab-a24757b8f7fc).

My understanding is that using the **GPU** is simply a matter of creating a variable `device` and assigning it `cuda`, like this:

```python
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

```

And then you would move your model and your tensors to the device.

```python
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
model = model.to(device)

# Training loop:
    for i, batch in enumerate(dl):

        optimizer.zero_grad()

        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        y = batch['y'].to(device)

        result = model(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)

```

To use a **TPU** , the article above mentioned creating an analogous `device` variable but setting it to use XLA.

```python
import torch_xla
import torch_xla.core.xla_model as xm

device = xm.xla_device()

model = model.to(device)

```

I tried this ostensibly straight-forward approach but when I run training, it’s running extremely slowly, practically at the same speed as CPU-only training.

Here is my [Google Colab notebook](https://colab.research.google.com/drive/1up7f3CvOLRhXYXEBlcA_oZxcaz9MxipW?usp=sharing) with my attempt. It runs well with GPU, but exceedingly slowly with TPU.

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](https://discuss.huggingface.co/u/sgugger)\
**Post date:** [December 7, 2020, 2:31pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/4 "2020-12-07T14:31:18Z")

</div>

At a first glance, you have `loss.item()` in your training loop, which you should absolutely avoid on TPUs (it’s a big slowdown). You should use `loss.detach()` to accumulate your losses on the TPU then only do the `.item()` at the very end of your epoch.

---

<div class="post-metadata">

**Author:** ![facehugger2020](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/facehugger2020/32/3052_2.png) [@facehugger2020](https://discuss.huggingface.co/u/facehugger2020)\
**Post date:** [December 7, 2020, 8:58pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/5 "2020-12-07T20:58:20Z")

</div>

Thanks. I did as you suggested, but the training loop is still making very slow progress.

OLD:

```python
    epoch_loss = 0.0
    for i, batch in enumerate(dl):

        loss = loss_fn(yhat, y)
        loss.backward()

        epoch_loss += loss.item()

    return epoch_loss/len(dl)

```

NEW:

```python
    epoch_loss = 0.0
    for i, batch in enumerate(dl):

        loss = loss_fn(yhat, y)
        loss.backward()

        epoch_loss += loss.detach() # <-- NEW

    return epoch_loss.item()/len(dl) # <-- NEW

```

One batch is still taking a long time to complete. I suspect it’s running on the CPU rather than the TPU. However, I think I followed all XLA setup correctly. If this issue is out of Transformers’ domain, I’ll go bug the XLA folks.

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](https://discuss.huggingface.co/u/sgugger)\
**Post date:** [December 7, 2020, 9:59pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/6 "2020-12-07T21:59:17Z")

</div>

Mmmm, I don’t see the call to the spawn function, so yes, you’re probably training on CPU. Normally, you are supposed to called `train` through

```auto
import torch_xla.distributed.xla_multiprocessing as xmp
xmp.spawn(train, args=potential_args, nprocs=num_tpu_cores)

```

---

<div class="post-metadata">

**Author:** ![facehugger2020](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/facehugger2020/32/3052_2.png) [@facehugger2020](https://discuss.huggingface.co/u/facehugger2020)\
**Post date:** [December 7, 2020, 10:53pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/7 "2020-12-07T22:53:43Z")

</div>

My understanding from reading the [PyTorch XLA documentation](https://pytorch.org/xla/release/1.7/index.html#running-models-on-xla-devices) is that `xmp.spawn()` is used for multi-TPU processing. For single-TPU training, you only need to define `device` correctly. The difference is also shown in the PyTorch example code for [single-core AlexNet](https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/single-core-alexnet-fashion-mnist.ipynb) and [multi-core AlexNet](https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/multi-core-alexnet-fashion-mnist.ipynb) training.

At this time, I’m interested in just single-TPU execution.

---

<div class="post-metadata">

**Author:** ![facehugger2020](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/facehugger2020/32/3052_2.png) [@facehugger2020](https://discuss.huggingface.co/u/facehugger2020)\
**Post date:** [December 7, 2020, 11:15pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/8 "2020-12-07T23:15:33Z")

</div>

I filed a Github issue with the XLA team.

---

<div class="post-metadata">

**Author:** ![tillfurger](https://avatars.discourse-cdn.com/v4/letter/t/6bbea6/32.png) [@tillfurger](https://discuss.huggingface.co/u/tillfurger)\
**Post date:** [February 15, 2021, 8:30pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/9 "2021-02-15T20:30:25Z")

</div>

Any update on this?

---

<div class="post-metadata">

**Author:** ![finiteautomata](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/finiteautomata/32/47620_2.png) [@finiteautomata](https://discuss.huggingface.co/u/finiteautomata)\
**Post date:** [April 12, 2021, 7:09pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/10 "2021-04-12T19:09:26Z")

</div>

I’ve succesfully trained a NLI model using colab’s TPUs. I had to struggle a little bit with configuring the environment but thanks to the help of the fellows of `huggingface` I was able to do it.

You can check it here:

> **[Google Colaboratory](https://colab.research.google.com/drive/1dVEfoxGvMAKd0GLnrUJSHZycGtyKt9mr#scrollTo=declared-pioneer)**

---

<div class="post-metadata">

**Author:** ![PremalMatalia](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/premalmatalia/32/1457_2.png) [@PremalMatalia](https://discuss.huggingface.co/u/PremalMatalia)\
**Post date:** [June 14, 2021, 5:38pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/11 "2021-06-14T17:38:00Z")

</div>

I can see an error while training in your colab notebook…was it really successful?

---

<div class="post-metadata">

**Author:** ![finiteautomata](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/finiteautomata/32/47620_2.png) [@finiteautomata](https://discuss.huggingface.co/u/finiteautomata)\
**Post date:** [June 17, 2021, 2:53pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/12 "2021-06-17T14:53:10Z")

</div>

@PremalMatalia: you’re right. I couldn’t identify the source of error (perhaps a memory limit on Colab?).

BTW, I reworked the notebook because it had some sub-optimal stuff (for example, preprocessing the whole dataset for each fork). Also, I added some explanations. Now it works ok and it’s faster than before, and I hope it’s more readable.

Check it back:

> **[Google Colaboratory](https://colab.research.google.com/drive/1dVEfoxGvMAKd0GLnrUJSHZycGtyKt9mr#scrollTo=FyTR9V5jJWcS)**

---

<div class="post-metadata">

**Author:** ![GenV](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/genv/32/12132_2.png) [@GenV](https://discuss.huggingface.co/u/GenV)\
**Post date:** [February 15, 2022, 4:05pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/13 "2022-02-15T16:05:38Z")

</div>

@finiteautomata I’m trying this code now but I don’t see any prints, I don’t know if it’s working or not. Is there a way to see a percentage of the training? Thanks

---

<div class="post-metadata">

**Author:** ![finiteautomata](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/finiteautomata/32/47620_2.png) [@finiteautomata](https://discuss.huggingface.co/u/finiteautomata)\
**Post date:** [February 15, 2022, 5:35pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/14 "2022-02-15T17:35:05Z")

</div>

Hum, there should be `tqdm` bars all around the training. Perhaps it didn’t even start in case you see nothing.

By the way, a couple of months after this, I suggest going for a different task to learn how TPUs work. NLI is by no means a good example of this. Pre-training or fine-tuning a language model is a good use case for this hardware. Check [run\_mlm.py](https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_mlm.py) example, it can be easily adapted using the same ideas of my former notebook.

---

<div class="post-metadata">

**Author:** ![GenV](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/genv/32/12132_2.png) [@GenV](https://discuss.huggingface.co/u/GenV)\
**Post date:** [February 16, 2022, 1:17pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/15 "2022-02-16T13:17:57Z")

</div>

@finiteautomata I have opened a thread for this and I’m running an official code → [link](https://discuss.huggingface.co/t/tpu-slow-finetuning-t5-base/14629/2), my problem is that TPU is very very slow on google colab and I don’t know why.

---

<div class="post-metadata">

**Author:** ![JKMO](https://avatars.discourse-cdn.com/v4/letter/j/c0e974/32.png) [@JKMO](https://discuss.huggingface.co/u/JKMO)\
**Post date:** [December 8, 2022, 2:12pm UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/16 "2022-12-08T14:12:23Z")

</div>

how does one learn how exactly one can use and train from notebook?

---

<div class="post-metadata">

**Author:** ![KraoESPfan1n](https://avatars.discourse-cdn.com/v4/letter/k/6a8cbe/32.png) [@KraoESPfan1n](https://discuss.huggingface.co/u/KraoESPfan1n)\
**Post date:** [June 3, 2024, 1:57am UTC](https://discuss.huggingface.co/t/tutorials-for-using-colab-tpus-with-huggingface-transformers/1970/17 "2024-06-03T01:57:15Z")

</div>

but how can I run a current model, not train, use a model with TPU v2