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

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

目標

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

費用

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

  • Compute Engine
  • Cloud TPU

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

始める前に

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

  1. Google アカウントにログインします。

    Google アカウントをまだお持ちでない場合は、新しいアカウントを登録します。

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

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

  3. Google Cloud プロジェクトに対して課金が有効になっていることを確認します。プロジェクトに対して課金が有効になっていることを確認する方法を学習する

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

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

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

    Cloud Shell を開く

  2. プロジェクト ID の変数を作成します。

    export PROJECT_ID=project-id
    
  3. Cloud TPU を作成するプロジェクトを使用するように gcloud コマンドライン ツールを構成します。

    gcloud config set project ${PROJECT_ID}
    
  4. 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
    
  5. 新しい Compute Engine インスタンスに接続します。

    gcloud compute ssh roberta-tutorial --zone=us-central1-a
    

Cloud TPU リソースを起動する

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

    (vm) $ gcloud compute tpus create roberta-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-1.6  \
    --accelerator-type=v3-8
    
  2. Cloud TPU リソースの IP アドレスを識別します。

    (vm) $ gcloud compute tpus list --zone=us-central1-a
    

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

  1. conda 環境を開始します。

    (vm) $ conda activate torch-xla-1.6
    
  2. Cloud TPU リソースの環境変数を構成します。

    (vm) $ export TPU_IP_ADDRESS=ip-address
    
    (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    

データを設定する

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

    (vm) $ pip install --editable /usr/share/torch-xla-1.6/tpu-examples/deps/fairseq
    
  2. モデルデータを格納する pytorch-tutorial-data ディレクトリを作成します。

    (vm) $ mkdir $HOME/pytorch-tutorial-data
    (vm) $ cd $HOME/pytorch-tutorial-data
    
  3. 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) $ python /usr/share/torch-xla-1.6/tpu-examples/deps/fairseq/train.py $DATA_DIR \
    --task masked_lm \
    --criterion masked_lm \
    --arch roberta_base \
    --sample-break-mode complete \
    --tokens-per-sample $TOKENS_PER_SAMPLE \
    --optimizer adam \
    --adam-betas '(0.9,0.98)' \
    --adam-eps 1e-6 \
    --clip-norm 0.0 \
    --lr-scheduler polynomial_decay \
    --lr $PEAK_LR \
    --warmup-updates $WARMUP_UPDATES \
    --total-num-update $TOTAL_UPDATES \
    --dropout 0.1 \
    --attention-dropout 0.1 \
    --weight-decay 0.01 \
    --update-freq $UPDATE_FREQ \
    --max-update $TOTAL_UPDATES \
    --log-format simple \
    --valid-subset=valid \
    --train-subset=train \
    --num_cores=8 \
    --metrics_debug \
    --input_shapes 16x512 18x480 21x384 \
    --save-dir=${HOME}/checkpoints \
    --log_steps=30 \
    --max-epoch=1 \
    --skip-invalid-size-inputs-valid-test
   

トレーニング スクリプトは約 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

クリーンアップ

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

  1. Compute Engine インスタンスとの接続を切断していない場合は切断します。

    (vm) $ exit
    

    プロンプトが user@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

  2. Cloud Shell で、gcloud コマンドライン ツールを使用して、Compute Engine インスタンスを削除します。

    $ gcloud compute instances delete roberta-tutorial --zone=us-central1-a
    
  3. gcloud コマンドライン ツールを使用して、Cloud TPU リソースを削除します。

    $ gcloud compute tpus delete roberta-tutorial --zone=us-central1-a
    

次のステップ

次のように PyTorch colabs を試す