Cloud TPU での PyTorch を使用した FairSeq RoBERTa の事前トレーニング

このチュートリアルでは、Cloud TPU で FairSeq の RoBERTa を事前トレーニングする方法を説明します。具体的には、公開されている wikitext-103 データセットでモデルを事前トレーニングする FairSeq のチュートリアルに従います。

目標

PyTorch 環境を作成および構成する
データセットを準備する
トレーニングジョブを実行する
出力結果を表示できることを確認する

費用

このチュートリアルでは、課金対象である次の Google Cloud コンポーネントを使用します。

Compute Engine
Cloud TPU

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。新しい Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

始める前に

このチュートリアルを開始する前に、Google Cloud プロジェクトが正しく設定されていることを確認します。

Google Cloud アカウントにログインします。Google Cloud を初めて使用する場合は、アカウントを作成して、実際のシナリオでの Google プロダクトのパフォーマンスを評価してください。新規のお客様には、ワークロードの実行、テスト、デプロイができる無料クレジット $300 分を差し上げます。

Google Cloud Console の [プロジェクトセレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

プロジェクトセレクタに移動

Cloud プロジェクトに対して課金が有効になっていることを確認します。詳しくは、プロジェクトで課金が有効になっているかどうかを確認する方法をご覧ください。

Google Cloud Console の [プロジェクトセレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

プロジェクトセレクタに移動

このチュートリアルでは、Google Cloud の課金対象となるコンポーネントを使用します。費用を見積もるには、Cloud TPU の料金ページを確認してください。不要な課金を回避するために、このチュートリアルを完了したら、作成したリソースを必ずクリーンアップしてください。

Compute Engine インスタンスを設定する

Cloud Shell ウィンドウを開きます。

Cloud Shell を開く
プロジェクト ID の変数を作成します。
```
export PROJECT_ID=project-id
```
Cloud TPU を作成するプロジェクトを使用するように Google Cloud CLI を構成します。
```
gcloud config set project ${PROJECT_ID}
```
このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した GCP API の呼び出しを許可します。

Cloud Shell から、このチュートリアルで必要となる Compute Engine リソースを起動します。

gcloud compute instances create roberta-tutorial \
--zone=us-central1-a \
--machine-type=n1-standard-16  \
--image-family=torch-xla \
--image-project=ml-images  \
--boot-disk-size=200GB \
--scopes=https://www.googleapis.com/auth/cloud-platform

新しい Compute Engine インスタンスに接続します。
```
gcloud compute ssh roberta-tutorial --zone=us-central1-a
```
重要なポイント: これ以降、接頭辞 (vm) $ は Compute Engine VM インスタンスでコマンドを実行する必要があることを意味します。

Cloud TPU リソースを起動する

Compute Engine 仮想マシンから、次のコマンドを使用して Cloud TPU リソースを起動します。

(vm) $ gcloud compute tpus create roberta-tutorial \
--zone=us-central1-a \
--network=default \
--version=pytorch-1.11  \
--accelerator-type=v3-8

Cloud TPU リソースの IP アドレスを識別します。
```
(vm) $ gcloud compute tpus describe --zone=us-central1-a roberta-tutorial
```
重要: IP アドレスは NETWORK_ENDPOINTS 列の下に表示されます。この IP アドレスは、ポート番号なしで PyTorch 環境を作成して構成するときに必要になります。

PyTorch 環境を作成および構成する

conda 環境を開始します。
```
(vm) $ conda activate torch-xla-1.11
```
Cloud TPU リソースの環境変数を構成します。

注: TPU_IP_ADDRESS 変数は、Cloud TPU リソースを起動する際に特定される Cloud TPU の IP アドレスと同じである必要があります。これは、Google Cloud Console の [Compute Engine] > [TPU] に表示される内部 IP アドレスです。
```
(vm) $ export TPU_IP_ADDRESS=ip-address
```
```
(vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
```

データを設定する

以下を実行して FairSeq をインストールします。

(vm) $ pip install --editable /usr/share/torch-xla-1.11/tpu-examples/deps/fairseq

モデルデータを格納する pytorch-tutorial-data ディレクトリを作成します。
```
(vm) $ mkdir $HOME/pytorch-tutorial-data
(vm) $ cd $HOME/pytorch-tutorial-data
```
FairSeq RoBERTa の README「データの前処理」セクションの手順に従います。データセットの準備には約 10 分かかります。

モデルのトレーニング

モデルをトレーニングするには、まず環境変数を設定します。

(vm) $ export TOTAL_UPDATES=125000    # Total number of training steps
(vm) $ export WARMUP_UPDATES=10000    # Warmup the learning rate over this many updates
(vm) $ export PEAK_LR=0.0005          # Peak learning rate, adjust as needed
(vm) $ export TOKENS_PER_SAMPLE=512   # Max sequence length
(vm) $ export UPDATE_FREQ=16          # Increase the batch size 16x
(vm) $ export DATA_DIR=${HOME}/pytorch-tutorial-data/data-bin/wikitext-103

次のスクリプトを実行します。

(vm) $ python3 \
      /usr/share/torch-xla-1.10/tpu-examples/deps/fairseq/train.py $DATA_DIR \
      --task=masked_lm --criterion=masked_lm \
      --arch=roberta_base --sample-break-mode=complete \
      --tokens-per-sample=512 \
      --optimizer=adam \
      --adam-betas='(0.9,0.98)' \
      --adam-eps=1e-6 \
      --clip-norm=0.0 \
      --lr-scheduler=polynomial_decay \
      --lr=0.0005 \
      --warmup-updates=10000 \
      --dropout=0.1 \
      --attention-dropout=0.1 \
      --weight-decay=0.01 \
      --update-freq=16 \
      --train-subset=train \
      --valid-subset=valid \
      --num_cores=8 \
      --metrics_debug \
      --save-dir=checkpoints \
      --log_steps=30 \
      --log-format=simple \
      --skip-invalid-size-inputs-valid-test \
      --suppress_loss_report \
      --input_shapes 16x512 18x480 21x384 \
      --max-epoch=1

トレーニングスクリプトは約 15 分間実行され、終了すると次のようなメッセージが表示されます。

saved checkpoint /home/user/checkpoints/checkpoint1.pt
(epoch 1 @ 119 updates) (writing took 25.19265842437744 seconds)
| done training in 923.8 seconds

出力結果を確認する

トレーニングジョブが完了の後、モデルチェックポイントが次のディレクトリに置かれます。

$HOME/checkpoints

クリーンアップ

作成したリソースを使用した後、アカウントに不要な請求が発生しないようにクリーンアップを行います。

Compute Engine インスタンスとの接続を切断していない場合は切断します。
```
(vm) $ exit
```
プロンプトが user@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。
Cloud Shell で、Google Cloud CLI を使用して Compute Engine インスタンスを削除します。
```
$ gcloud compute instances delete roberta-tutorial --zone=us-central1-a
```
Google Cloud CLI を使用して Cloud TPU リソースを削除します。
```
$ gcloud compute tpus delete roberta-tutorial --zone=us-central1-a
```

次のステップ

次のように PyTorch colabs を試す