Tutorials for using Colab TPUs with Huggingface Transformers?
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.
2020-11-12T12:55:17Z
There are a few contributed notebooks which might help you here:https://github.com/huggingface/transformers/tree/master/notebooksFor instance this one by@valhallais 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
2020-12-06T23:21:14Z
I still cannot get any HuggingFace Tranformer model to train with a Google Colab TPU.I tried out the notebook mentioned above illustratingT5 training on TPU, but it uses theTrainerAPI and the XLA code is very ad hoc.I also tried a more principled approach based onan article by a PyTorch engineer.My understanding is that using theGPUis simply a matter of creating a variabledeviceand assigning itcuda, like this:device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')And then you would move your model and your tensors to the device.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 aTPU, the article above mentioned creating an analogousdevicevariable but setting it to use XLA.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 myGoogle Colab notebookwith my attempt. It runs well with GPU, but exceedingly slowly with TPU.
2020-12-07T14:31:18Z
At a first glance, you haveloss.item()in your training loop, which you should absolutely avoid on TPUs (it’s a big slowdown). You should useloss.detach()to accumulate your losses on the TPU then only do the.item()at the very end of your epoch.
2020-12-07T20:58:20Z
Thanks. I did as you suggested, but the training loop is still making very slow progress.OLD: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: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) # <-- NEWOne 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.
2020-12-07T21:59:17Z
Mmmm, I don’t see the call to the spawn function, so yes, you’re probably training on CPU. Normally, you are supposed to calledtrainthroughimport torch_xla.distributed.xla_multiprocessing as xmp
xmp.spawn(train, args=potential_args, nprocs=num_tpu_cores)
2020-12-07T22:53:43Z
My understanding from reading thePyTorch XLA documentationis thatxmp.spawn()is used for multi-TPU processing. For single-TPU training, you only need to definedevicecorrectly. The difference is also shown in the PyTorch example code forsingle-core AlexNetandmulti-core AlexNettraining.At this time, I’m interested in just single-TPU execution.
2020-12-07T23:15:33Z
I filed a Github issue with the XLA team.
2021-02-15T20:30:25Z
Any update on this?
2021-04-12T19:09:26Z
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 ofhuggingfaceI was able to do it.You can check it here:colab.research.google.comGoogle Colaboratory
2021-06-14T17:38:00Z
I can see an error while training in your colab notebook…was it really successful?
2021-06-17T14:53:10Z
@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:colab.research.google.comGoogle Colaboratory
2022-02-15T16:05:38Z
@finiteautomataI’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
2022-02-15T17:35:05Z
Hum, there should betqdmbars 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. Checkrun_mlm.pyexample, it can be easily adapted using the same ideas of my former notebook.
2022-02-16T13:17:57Z
@finiteautomataI have opened a thread for this and I’m running an official code →link, my problem is that TPU is very very slow on google colab and I don’t know why.
2022-12-08T14:12:23Z
how does one learn how exactly one can use and train from notebook?
2024-06-03T01:57:15Z
but how can I run a current model, not train, use a model with TPU v2