Cloud TPUの使い⽅ 〜BigBirdの⽇本語学習済みモデルを作る〜 板垣正敏＠Python機械学習勉強会in新潟 2021/7/10 

アジェンダ š Cloud TPUとは︖ š 今回のモチベーション š Cloud TPUが使える深層学習環境 š GCE上でのモデル訓練 š 予算枯渇と救世主の出現 š BigBirdの学習結果は︖ š まとめ 

Cloud TPUとは︖ 

Cloud TPUとは︖ š Googleが作っている Tensor計算専⽤プロ セッサ（ASIC） š 専⽤仮想マシン （VM）に接続され ておりgRPC経由で 利⽤する分散環境 š 2021年5⽉には直接 VMを利⽤できる サービスも発表 š Cloud TPU v2 š180 TFLOPS š64 GB HBM š Cloud TPU v3 š420 TFLOPS š128 GB HBM 

データ並列とモデル並列 š 並列化には2種類ある š メモリに乗り切らない⼤きなモデルを分割して処 理する「モデル並列」（上図右） š たくさんの学習データを分散して処理する「デー タ並列」（上図左） https://tech.preferred.jp/ja/blog/model-parallelism-in-dnn/ 

TPUを使う⽅法の⼀般論 フレーム ワーク データ並列 モデル並列 TensorFlow l tf.distribute.TPUStrategy() l tf.compat.v1.estimator.tpu.T PUEstimator l Mesh TensorFlow l ⼿動配置 PyTorch l torch_xla l torch.distributed.pipeline l ⼿動配置 l HuggingFace/Transformers では、⼀部のモデルで Attention単位の並列化 （Parallelize）が可能 

GPUとCloud TPUのコスト⽐較 タイプ vCPU GPU CPU メモリ GPU/TPU メモリ 料⾦／ 1時間 AWS p2.xlarge 4 K80/12GB 61GB 12GB $0.90 p3.2xlarge 8 V100/16GB 61GB 16GB $3.06 p3.8xlarge 32 V100/16GB×4 244GB 64GB $12.24 p3.16xlarge 64 V100/16GB×8 488GB 128GB $24.48 p4d.24xlarge 96 A100/40GB×8/NVSwitch 1152GB 320GB $32.7726 GCP n1-highmem-4 4 K80/12GB 52GB 12GB $0.65 n1-highmem-8 8 V100/16GB 52GB 16GB $2.95 n1-standard-96 96 V100/16GB×8 320GB 128GB $24.40 n1-standard-32/V2-8 32 Cloud TPU v2×8 120GB 64GB $6.02 n1-standard-64/V3-8 64 Cloud TPU v3×8 240GB 128GB $11.04 colab 無料版 1 Cloud TPU v2×8 12GB 64GB $0 

今回のモチベーション 

⽇本語要約処理で⻑⽂を扱いたい š 会社在籍時に某⼤学と⽇本語要約の共同研究をしていた š 現在の主流はTransformerを使ったモデル š 巨⼤な多⾔語コーパスを使った学習済みモデルmT5も試したが、多 くのモデルが最⼤トークン数512程度 š ⻑⽂を扱える⽇本語学習済みモデルがあれば… š ⻑⽂を扱えるNLPモデルの例 š Transformer-XL š Longformer š Reformer š BigBird š 公開されている⽇本語学習済みモデルがない…⾃分で作るしかない 

BigBirdとは︖ š Transformerを使った⻑い⽂章を 扱うモデルが作りにくいのは、 中核であるAttentionが⽂章の⻑ さ（トークン数）の2乗に⽐例 して⼤きくなるから š BigBird は Full Attention を Locality/Random/Global Connections の組み合わせに変 更 š さらにアテンションをブロック 化して計算⼿順を⼯夫し GPU/TPUでの計算効率を改善 https://github.com/google-research/bigbird 

訓練に使⽤したデータ š TensorFlow Datasetsに含まれるWikipediaのデータで、テキスト抽出な どの前処理が終わっているWiki40bを使⽤した š トークナイザはSentencepieceを使い、wiki40b/jaから抽出し⽂単位に 分割したテキストデータ（約1.9GB）で学習を⾏なった š wiki40bで使われている_START_ARTICLE_, _START_SECTION_, _START_PARAGRAPH_, _NEWLINE_をuser_defined_symbolsに指定 š BERTで使われる[CLS],{SEP],[MASK]をcontrol_symbolsに指定 š BigBirdの事前学習⽤データはBERTと同様なので、 https://github.com/yoheikikuta/bert-japanese の学習データ作成プログ ラムを利⽤した（サイズ約37GB） š max_token_length=4096, masked_lm_prob=0.15, max_predictions_per_seq=600, dupe_factor=10 š 属性は、input_ids, segment_ids, masked_lm_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels 

BigBirdのモデルサイズとメモリ š 通常のTransformerに⽐べて8分の1のメモリ使⽤量 š 標準的なTransformerで512だった最⻑トークン数を4096にして もモデルサイズはメモリ16GBのGPUで訓練できるほどに š でも⾃宅のGeForce RTX2080 Ti（メモリ11GB）では無理 š メモリと計算時間のトレードオフもあるので、GPU1基では訓 練時間が膨⼤に š そこでCloud TPUを使ったモデル並列化の出番︕ 

Cloud TPUが使える 深層学習環境 

GCP上の深層学習環境 サービス 説明 TF PyT orch GPU TPU Pre- empti -ble GCE 仮想マシンサービスで環境を構築 ○ ○ ○ ○ ○ AI Platform Training ソースコードをtarでパッケージ化、 ⼜は、学習環境をカスタムコンテナ 化してAI Platformで実⾏ ○ ○ ○ ○ × AI Platform Pipelines GKEクラスタを作成し、前処理から 訓練・評価、本番環境への実装など をコンテナベースで⼀貫して管理 ○ ○ ○ ○ ○ GKE ⾃分でKubernetesクラスタを作成し コンテナベースのジョブを起動 ○ ○ ○ ○ ○ Vertex AI Pipelines マネージドクラスタでトレーニング ジョブを実⾏ ○ ○ ○ × × Google Colab 無償あるいは定額のNotebook環境か らGPUおよびCloud TPUを使⽤可能 ○ ○ ○ ○ v2 - 

Google Colabではなく GCEを使う理由 š Google Colabで使える無料のTPU（v2-8）は、1基あたりメモリ 8GBで、トークン⻑4096のBigBirdは乗らない š 有償だがメモリ・性能が上位のv3-8を使ってみたい š GCPでアカウントを作成すると有効期間90⽇で$300分の無料 クレジットが割り当てられる š アカウントの作成にはクレジットカード登録が必要だが、課 ⾦設定をしない限り課⾦はされない š 前職で社内トレーニングを⾏った際に作成したアカウントの クレジットが27,000円分以上残っていたので、これを使って 試すことにした 

コスト削減の⽅法 š Cloud TPUもGCEのVMもプリエンプティブルインスタンスを使 えばコストはさらに安くなる š プリエンプティブルインスタンスはAWSのスポットインスタ ンスと同様に、空いているリソースを利⽤するもので、GCP の都合でいつ停⽌されても⽂句は⾔えないが、数分の1の料⾦ で利⽤可能（v3-8: $2.40/時間） š Cloud TPUのプリエンプティブルインスタンスは最⼤24時間で 必ず停⽌ š モデルの訓練ではチェックポイントを保存しておいて、次回 はそこから再開することで、活⽤が容易 

GCEでのモデル訓練 

データとモデルはGCSに保存 š Cloud TPUはユーザーがアクセスできるGCE VMとは異なるノー ドで稼働するため、GCE VMのローカルストレージにはアクセ スできない š このため、使⽤するデータやモデルのチェックポイントの保 存にはGCSバケットを使⽤する必要がある 

Preemptibleインスタンスを使う š TensorFlowの場合には上記でVMとTPUが同時に作成される š PyTorchの場合には複数バージョンがAnacondaでインストールされている ため、TPUは別に作成する š 上記のコマンドでsshセッションが開始されるので、あとはソースコード をgit cloneしてスクリプトを実⾏ š Jupyterを使いたい場合には下記のコマンドでsshトンネリング gcloud compute ssh train-bigbird –- L8888:localhost:8888 gcloud compute tpus execution-groups create ¥ --name=train-bigbird ¥ --zone=$GCP_ZONE ¥ --tf-version=2.5.0 ¥ --machine-type=n1-standard-8 ¥ --accelerator-type=v3-8 ¥ --preemptible ¥ --preemptible-vm 

BigBird事前学習⽤スクリプトの修正 š BigBirdの事前学習⽤スクリプトbigbird/pretrain/run_pretraining.pyには、 テキストデータを⼊⼒としてマスクされた訓練データを⽣成する関 数（numpy_masking）が含まれる š この関数はTensorFlowの計算グラフに組み込むために tf.compat.v1.py_func()を通じて呼ばれているが、Cloud TPU上ではエ ラーとなり実⾏ができない š このため、この機能を削除し、前処理済みのTFRecordファイルのみ を扱うように修正した š なお、この関数（numpy_masking）では、[CLS], [SEP], [MASK]という 制御トークンのIDが65, 66, 67としてハードコーディングしてあるた め、GPU環境などで⾃前のSentencePieceモデルを使⽤する際には修 正が必要である 

訓練の実⾏ š リポジトリをクローンした 後、環境変数を設定、 BigBirdおよび依存関係モ ジュールをインストールし てから、右のように訓練を 実⾏する š 正常終了（訓練完了）ある いは異常終了（Cloud TPUが PREEMPTED）の場合に関わ らず、SSHセッションを終 了する š SSHセッションが終了した らクリーンアップを実⾏ python3 run_pretraining.py ¥ --data_dir ${DATA_DIR} ¥ --output_dir ${MODEL_DIR} ¥ --preprocessed_data ¥ --max_encoder_length 4096 ¥ --max_predictions_per_seq 600 ¥ --do_train ¥ --train_batch_size 2 ¥ --optimizer AdamWeightDecay ¥ --learning_rate 1e-4 ¥ --num_train_steps 1400000 ¥ --num_warmup_steps 10000 ¥ --save_checkpoints_steps 5000 ¥ --vocab_model_file ${SPM_MODEL} ¥ --use_nsp ¥ --use_tpu ¥ --tpu_name ${TPU_NAME} ¥ --num_tpu_cores=8 

結構⽌まる プリエンプティブルインスタンス š 実際にやってみると結構短い時間で停⽌されることがわかっ た→Checkpointの頻度を上げた š GCE VMとCloud TPUの両⽅をプリエンプティブルにしていたの で、どちらが停⽌されても再開の必要が出てくる š GCE VMが先に落ちると、Cloud TPUが課⾦されたままアイドル になるので注意︕（VMはPreemptibleにしない選択もあり） š 訓練のプロセスが終了したらVMをシャットダウンすることに した š gcloud compute tpus create execution-group はインスタンス起動 後にsshセッションを起動し、終了まで待つ仕様なので、この コマンド後にcleanupを⾏うスクリプトを使うことにした 

予算枯渇と 救世主の出現 

無償枠を使い果たし 早くも予算オーバー š 実は試⾏錯誤をしていたら、無償枠がゼロに（⾃動的にス トップ） š 1ヶ⽉以内に課⾦設定を有効にすれば継続利⽤可能 š 私の場合には課⾦アカウントがあったので、ダウンロードし てあったデータをアップロードして継続へ š 100,000ステップ学習したところで、設定してた予算枠10,000 円をオーバー 

性能はまだまだ š 試しに100,000ステップ学習したモデルでファインチューニン グを⾏い、ライブドアニュースの分類タスクをやってみたが、 正解率93％程度でまだまだ学習不⾜ š ⽇本語版BERTの学習曲線を⾒ても100,000ステップはほんの序 の⼝ š ⽇本語版BERTの140万ステップを学習させたら、14万円が 吹っ⾶ぶ︖ š とても個⼈ができるレベルではなくなる… š どうする︖ 

救世主の出現 š bigbird/core/flags.pyに”use_gradient_checkpointing”というフラ グを発⾒ š ニューラルネットワークの訓練時に⼤量のメモリを消費する のは、バックプロパゲーション（誤差逆伝播法）のために、 計算した勾配を保持しておくため š バックプロパケーションの際に勾配を再計算すれば保持する 必要なないが、その代わりに計算時間が⼤幅に増える š このトレードオフをバランスするため、⼀部のレイヤのみ保 持しそれ以外を破棄する Gradient Checkpointing が考えられた š GPUでこのフラグをつけてみるとメモリ使⽤量が半減 

Google Colabでも訓練可能に š use_gradient_checkpointingを有効にすることにより、コアあた り8GBしかメモリがないCloud TPU v2でも訓練が可能に š GCEを使わずとも、Google Colaboratoryで訓練ができるように なった（ただし、GCSのコストはかかる） š ⽉額1,072円のGoogle Colab Proを使っているので、セッション の中断が少なくなった š CAPTCHaが出てくるのは下記のスクリプトで回避 š https://www.rockyourcode.com/script-to-stop-google-colab-from- disconnecting/ 

訓練が中断する環境で気をつけ るべきこと š 今回のように巨⼤なデータセットを使い、訓練が中断される 環境で訓練を⾏う際、データセットのシャッフルを⾏わない と、「いつもデータセットの最初の⽅しか訓練に使われな い」という落とし⽳が…（途中で気がついた） š 当初は1個の巨⼤な（34.5GB）のTFRecordファイルを使ってい たが、複数のファイルに分割することに š bigbird/pretrain/run_pretraining.pyには複数のTFRecordをシャッ フルし、tf.data.Dataset.interleave()を使って混合するコードが 含まれている š 今回は、TFRecordにする前のテキストファイルをドキュメン ト数で均等分割し使⽤している（つまり、各ファイルのサン プル数には差異がある） 

BigBirdの学習結果は︖ 

現在も訓練続⾏中 š 約18日間訓練中 š 807,000ステップ時点の損失1.876 š 1,400,000ステップでの学習結果と、応用について は次の機会をお楽しみに 

まとめ 

まとめ š ⻑⽂を扱える⾃然⾔語処理モデルBigBirdを⽇本語Wikipediaのデータ セットで事前学習させています š TPUが使えるモデルであれば、GPUよりも低コストでデータ並列を 使ったモデルの訓練が可能です š GCPでは、さまざまな⽅法でCloud TPUを使⽤することができますが、 コストを抑えるにはPreemptible TPUを使⽤するのが効果的です š Cloud TPUを使った訓練では、GCE VMの能⼒はさほど必要ありません š それでも、⼤量のデータを学習させるには時間とお⾦がかかります š モデル訓練時の使⽤メモリを削減させる⼿法として、Gradient Checkpointingは有効です（ただし、訓練時間は⻑くなります） š Cloud TPU v2で⼗分なモデルでは、Google Colaboratoryが最強で、Pro にすれば中断も抑えられます 

参考資料 š Cloud TPU š https://cloud.google.com/tpu?hl=ja š BigBird š https://github.com/google-research/bigbird š 板垣作成スクリプト š https://github.com/masa-ita/train_bigbird š Gradient Checkpointing論⽂ š Training Deep Nets with Sublinear Memory Cost š https://arxiv.org/abs/1604.06174v2 š Google Colab Pro š https://colab.research.google.com/signup