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

このチュートリアルでは、FairSeq バージョンの Transformer と英語からドイツ語に翻訳する WMT 18 翻訳タスクを中心に取り上げます。

目標

データセットを準備します。
トレーニングジョブを実行します。
出力結果を確認します。

費用

このチュートリアルでは、課金対象である次の 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 に認証情報を使用した API の呼び出しを許可します。

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

gcloud compute --project=${PROJECT_ID} instances create transformer-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 transformer-tutorial --zone=us-central1-a
```
重要なポイント: これ以降、接頭辞 (vm) $ は Compute Engine VM インスタンスでコマンドを実行する必要があることを意味します。

Cloud TPU リソースを起動する

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

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

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

データのダウンロード

モデルデータを格納する pytorch-tutorial-data ディレクトリを作成します。
```
(vm) $ mkdir $HOME/pytorch-tutorial-data
```
pytorch-tutorial-data ディレクトリに移動します。
```
(vm) $ cd $HOME/pytorch-tutorial-data
```

モデルデータをダウンロードします。

(vm) $ wget https://dl.fbaipublicfiles.com/fairseq/data/wmt18_en_de_bpej32k.zip

データを抽出します。

(vm) $ sudo apt-get install unzip && \
unzip wmt18_en_de_bpej32k.zip

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

conda 環境を開始します。
```
(vm) $ conda activate torch-xla-1.13
```
Cloud TPU リソースの環境変数を構成します。
```
(vm) $ export TPU_IP_ADDRESS=ip-address; \
export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
```
注: TPU_IP_ADDRESS 変数は、Cloud TPU リソースを起動する際に特定される Cloud TPU の IP アドレスと同じである必要があります。

モデルのトレーニング

モデルをトレーニングするには、次のスクリプトを実行します。

(vm) $ python /usr/share/torch-xla-1.13/tpu-examples/deps/fairseq/train.py \
  $HOME/pytorch-tutorial-data/wmt18_en_de_bpej32k \
  --save-interval=1 \
  --arch=transformer_vaswani_wmt_en_de_big \
  --max-target-positions=64 \
  --attention-dropout=0.1 \
  --no-progress-bar \
  --criterion=label_smoothed_cross_entropy \
  --source-lang=en \
  --lr-scheduler=inverse_sqrt \
  --min-lr 1e-09 \
  --skip-invalid-size-inputs-valid-test \
  --target-lang=de \
  --label-smoothing=0.1 \
  --update-freq=1 \
  --optimizer adam \
  --adam-betas '(0.9, 0.98)' \
  --warmup-init-lr 1e-07 \
  --lr 0.0005 \
  --warmup-updates 4000 \
  --share-all-embeddings \
  --dropout 0.3 \
  --weight-decay 0.0 \
  --valid-subset=valid \
  --max-epoch=25 \
  --input_shapes 128x64 \
  --num_cores=8 \
  --metrics_debug \
  --log_steps=100

注: `input_shapes` ハイパーパラメータの値を変更する（たとえば、以下のとおりに）と、パフォーマンスが向上します。

--input_shapes 256x64 512x32

これらの変更により、初期コンパイルは大幅に遅くなりますが、安定化後はエポック時間が短くなります。

これらの入力シェイプを使用するには、bfloat16 の使用を有効にする必要があります。具体的には、次のコマンドを実行します。

(vm) $ export XLA_USE_BF16=1

クリーンアップ

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

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

次のステップ

次のように PyTorch colabs を試す