How to Colab with TPU. Training a Huggingface BERT on Google... | by Rohan Jagtap | Towards Data Science

TPUs (Tensor Processing Units) are application-specific integrated circuits (ASICs) that are optimized specifically for processing matrices.
Cloud TPU resources accelerate the performance of linear algebra computation, which is used heavily in machine learning applications
— Cloud TPU Documentation
Google Colab provides experimental support for TPUs for free! In this article, we’ll be discussing how to train a model using TPU on Colab. Specifically, we’ll be training BERT for text classification using the transformers package by huggingface on a TPU.
First things first. Since the TPU is optimized for some specific operations, we need to check if our model actually uses them; i.e. we need to check if the TPU actually helps our model to train faster. Following are some use cases where we might want to use a TPU as mentioned in the Cloud TPU Documentation:
Models dominated by matrix computations
Models with no custom TensorFlow operations inside the main training loop
Models that train for weeks or months
Larger and very large models with very large effective batch sizes
If your model uses a custom TensorFlow operation that is not present in Cloud TPU supported TensorFlow ops, you may rather use a GPU accelerator instead.
TPUs work on the cloud, unlike GPUs or CPUs that work locally. Hence we need some initialization before we start:
If you observe the output from the snippet above, our TPU cluster has 8 logical TPU devices (0–7) that are capable of parallel processing. Hence, we define a distribution strategy for distributed training over these 8 devices:
strategy = tf.distribute.TPUStrategy(resolver)
For more on distributed training refer: https://www.tensorflow.org/guide/distributed_training
In this section, we’ll be actually seeing how to train a BERT on TPU. We will do this in 2 ways:
Using model.fit()Using Custom Training Loop.
Using model.fit()
Using Custom Training Loop.
Since we are using a distribution strategy, the model must be created on each device for parameter sharing. Hence, the model needs to be created and built within the strategy’s scope:
Then, we simply fit the model on the data as one would do in a regular training setting:
To save the model weights:
model.save_weights("checkpoint/tpu-model.h5")
In the next subsection, we will discuss how to do the same thing using a custom training loop.
Here, we just need to tweak a few things manually that TensorFlow did for us in the backend in the previous method.
First, we create a data pipeline using tf.data API:
Now, the reason we didn’t have to worry about this in the previous part is that TensorFlow took care of these things itself; i.e. when we called model.fit(). Also this time, we need to manually distribute the dataset among the TPU devices:
Next, we create and build the model in the exact same way as in the previous method. Optionally, we can add some metrics in the strategy scope for manual loss and accuracy monitoring:
Now comes the most important part, i.e. the training step function. But first, let’s create an iterator for the distributed dataset:
train_iterator = iter(train_dataset)
Then we write the train_step function and decorate it using @tf.function as one would normally do. We use strategy.run() to execute the training step:
Finally, we run this train_step function for a few epochs in the training loop:
This time, let’s try saving the model using Checkpoints. Now there’s a issue here. Well, we can’t save the model just like that.
The error is pretty clear, it says you cannot access the local file system while executing eagerly, as the execution is taken to the cloud for the TPU to perform its ops.
So, to overcome this, we need to save the checkpoint in a GCS bucket. You can create a free tier GCP account here. First, we need to create a cloud storage bucket. Here is a tutorial from the official docs for creating a GCS bucket.
Next, we need to sign in using our GCP credentials and set the GCP project to the active configuration:
from google.colab import authauth.authenticate_user()!gcloud config set project <project-id>
gcloud config set sets the specified property in your active configuration only.
— Google Cloud Docs
Once this is done, we can access our bucket by simply using:
gs://<bucket-name>/<file-path>
Now the checkpoint looks something like this:
And this time, it will successfully save the model checkpoint to your bucket!
In this article, we saw why and how to tweak the original code that one would write to train a model, to make it TPU compatible. We also discussed when to and when not to use a TPU for training.