Running BERT on Google Cloud Platform With TPU
By: Jason Zhang
Tags: Machine Learning, Technology, Google Cloud Platform, Data Science, Deep Learning

Running BERT on Google Cloud Platform With TPU

Use Google Cloud and TPU to Build a Deep Learning Environment

Photo by Rajeshwar Bachu on Unsplash

By 2020, various deep learning models such as Transformer, BERT, and GPT-series make AI become smarter and smarter, using them can achieve many incredible tasks. On the hardware, Google launched Cloud TPU allows programmers to train these “Big” models at a reasonable cost.

In this article, I share how I built a Deep learning environment using Google Cloud Platform and Cloud TPU. And I also demonstrate running a BERT model task with this environment.

There are four parts in this article:

Part 1. Create a Deep Learning VM using the Google Cloud Console

Part 2. Create a Deep Learning VM using the command-line tools

Part 3. Run the BERT Task

Part 4. Clean up

Part 1. Create a Deep Learning VM using the Google Cloud Console

First, we use Google Cloud Console to create a deep learning environment. There are three steps:

Create a VM(virtual machine) instance Create a Cloud TPU node Create a Bucket

Before You Begin, you need:

In the Google Cloud Console, on the project selector page, select or create a Google Cloud project. Make sure that billing is enabled for your Cloud project. Enable the Compute Engine API

Step 1: Create a VM instance

In the Google Cloud Platform Console, go to the VM instances page Click the Create button. Specify a Name for your instance. Change the Zone for this instance. This option is related to the TPU type. Please refer to the next step: Zone section in Create a Cloud TPU. I choose us-central1-b . Select a Machine configuration for your instance.

IMPORTANT: TPU is not in the same physical space as your VM, and its use depends on network speed. And Google’s network speed is directly proportional to the number of CPUs. So in order to better use TPU, I recommend choosing the number of vCPU at least 16 . I select Machine type with n1-standard-16 (16vCPU, 60GB memory) In the Boot disk section, click Change to configure your boot disk.

For the Deep learning project, I choose Deep Learning on Linux in the Operating System section.

In Version section, we need to choose Debian image which contains Tensorflow or PyTorch pre-installed and optimized for Cloud TPUs. Depending on your project, you select the corresponding Tensorflow or PyTorch Version. I choose Debian GNU/Linux 10 Buster + TF 2-1-1 that is A Debain Linux image with Tensorflow Version 2–1–1 pre-installed and optimized for cloud TPUs or Debian GNU/Linux 9 Streth + PyTorch/XLA for PyTorch deep learning framework.

See also: Supported Tensorflow, Cloud TPU versions, and official models list.

IMPORTANT: Change Size value. Since the Deep learning image requires a dozen gigabytes, and larger models like BERT require 1–2 gigabytes for each checkpoint, the Size needs to be larger. I chose 500 GB. In Access scopes section, we check Allow full access to all Cloud APIs , because we need access to TPU and Google Bucket. In Preemptibility section, we recommend turning off this option. While a preemptible VM costs much less, it has the potential to cause loss of training results and state. Click Create, after waiting a few minutes, your VM has been built.

VM Information

Boot disk Information

Access scopes section

Turn off Preemptibility

VM instances lists

STEP 2: Create the Cloud TPU Node

In the Google Cloud Platform Console, go to the TPUs page Click CREATE TPU NODE In the Zone section, select the same zone as in the VM setting page. I chose us-central1-b here. Different zones support a different TPUs, please refer to TPU types and zones. In TPU type section, I choose v3-8 to get better performance. To make more efficient use of the TPU, you can use one TPU V3 at the same time when you’re training multiple models. In Preemptibility section, we can check out Turn on preemptibility for this TPU node .

Google provides the checkpoint mechanism for TPU so that even if your TPU is stopped by other users, you won’t lose your data. Using this option can save you a lot of money.

You can see below, when I checked this option, the estimated Hourly rate of this TPU dropped from US$8.00 to $2.40. In TPU Software version, select the same version as Operating system version in VM Settings. I choose 2.1.1 . Other sections can use the default value. Click CREATE, after waiting a few minutes, your TPU is available.

TPU Information

Compare Preemptibility Option On/Off

STEP 3: Create a Bucket

We do not save training data, results, and models on the VM’s drive because its price is more expensive. We use Google Cloud Storage to save them.

Open the Cloud Storage browser. Click Create bucket Enter your bucket information and click Continue to complete each step

CREATE BUCKET

The Bucket Information Part1

The Bucket Information Part2

Bucket lists

IMPORTANT: Set up your Compute Engine VM, your Cloud TPU node, and your Cloud Storage bucket in the same region/zone to reduce network latency and network costs.

After the above 3 steps, we have created a virtual machine optimized for deep learning with one TPU node and one Bucket.