Upgrade to Pro — share decks privately, control downloads, hide ads and more …

Cloud TPUの使い方〜BigBirdの日本語学習済みモデルを作る〜

Cloud TPUの使い方〜BigBirdの日本語学習済みモデルを作る〜

F8865f41777ef3caced0e4e6801ff83a?s=128

masa-ita

July 10, 2021
Tweet

Transcript

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

  2. アジェンダ š Cloud TPUとは︖ š 今回のモチベーション š Cloud TPUが使える深層学習環境 š

    GCE上でのモデル訓練 š 予算枯渇と救世主の出現 š BigBirdの学習結果は︖ š まとめ
  3. Cloud TPUとは︖

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

  6. 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)が可能
  7. 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
  8. 今回のモチベーション

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

    š ⻑⽂を扱えるNLPモデルの例 š Transformer-XL š Longformer š Reformer š BigBird š 公開されている⽇本語学習済みモデルがない…⾃分で作るしかない
  10. BigBirdとは︖ š Transformerを使った⻑い⽂章を 扱うモデルが作りにくいのは、 中核であるAttentionが⽂章の⻑ さ(トークン数)の2乗に⽐例 して⼤きくなるから š BigBird は

    Full Attention を Locality/Random/Global Connections の組み合わせに変 更 š さらにアテンションをブロック 化して計算⼿順を⼯夫し GPU/TPUでの計算効率を改善 https://github.com/google-research/bigbird
  11. 訓練に使⽤したデータ š 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
  12. BigBirdのモデルサイズとメモリ š 通常のTransformerに⽐べて8分の1のメモリ使⽤量 š 標準的なTransformerで512だった最⻑トークン数を4096にして もモデルサイズはメモリ16GBのGPUで訓練できるほどに š でも⾃宅のGeForce RTX2080 Ti(メモリ11GB)では無理

    š メモリと計算時間のトレードオフもあるので、GPU1基では訓 練時間が膨⼤に š そこでCloud TPUを使ったモデル並列化の出番︕
  13. Cloud TPUが使える 深層学習環境

  14. 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 -
  15. Google Colabではなく GCEを使う理由 š Google Colabで使える無料のTPU(v2-8)は、1基あたりメモリ 8GBで、トークン⻑4096のBigBirdは乗らない š 有償だがメモリ・性能が上位のv3-8を使ってみたい š

    GCPでアカウントを作成すると有効期間90⽇で$300分の無料 クレジットが割り当てられる š アカウントの作成にはクレジットカード登録が必要だが、課 ⾦設定をしない限り課⾦はされない š 前職で社内トレーニングを⾏った際に作成したアカウントの クレジットが27,000円分以上残っていたので、これを使って 試すことにした
  16. コスト削減の⽅法 š Cloud TPUもGCEのVMもプリエンプティブルインスタンスを使 えばコストはさらに安くなる š プリエンプティブルインスタンスはAWSのスポットインスタ ンスと同様に、空いているリソースを利⽤するもので、GCP の都合でいつ停⽌されても⽂句は⾔えないが、数分の1の料⾦ で利⽤可能(v3-8:

    $2.40/時間) š Cloud TPUのプリエンプティブルインスタンスは最⼤24時間で 必ず停⽌ š モデルの訓練ではチェックポイントを保存しておいて、次回 はそこから再開することで、活⽤が容易
  17. GCEでのモデル訓練

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

    存にはGCSバケットを使⽤する必要がある
  19. 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
  20. 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モデルを使⽤する際には修 正が必要である
  21. 訓練の実⾏ š リポジトリをクローンした 後、環境変数を設定、 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
  22. 結構⽌まる プリエンプティブルインスタンス š 実際にやってみると結構短い時間で停⽌されることがわかっ た→Checkpointの頻度を上げた š GCE VMとCloud TPUの両⽅をプリエンプティブルにしていたの で、どちらが停⽌されても再開の必要が出てくる

    š GCE VMが先に落ちると、Cloud TPUが課⾦されたままアイドル になるので注意︕(VMはPreemptibleにしない選択もあり) š 訓練のプロセスが終了したらVMをシャットダウンすることに した š gcloud compute tpus create execution-group はインスタンス起動 後にsshセッションを起動し、終了まで待つ仕様なので、この コマンド後にcleanupを⾏うスクリプトを使うことにした
  23. 予算枯渇と 救世主の出現

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

    š 100,000ステップ学習したところで、設定してた予算枠10,000 円をオーバー
  25. 性能はまだまだ š 試しに100,000ステップ学習したモデルでファインチューニン グを⾏い、ライブドアニュースの分類タスクをやってみたが、 正解率93%程度でまだまだ学習不⾜ š ⽇本語版BERTの学習曲線を⾒ても100,000ステップはほんの序 の⼝ š ⽇本語版BERTの140万ステップを学習させたら、14万円が

    吹っ⾶ぶ︖ š とても個⼈ができるレベルではなくなる… š どうする︖
  26. 救世主の出現 š bigbird/core/flags.pyに”use_gradient_checkpointing”というフラ グを発⾒ š ニューラルネットワークの訓練時に⼤量のメモリを消費する のは、バックプロパゲーション(誤差逆伝播法)のために、 計算した勾配を保持しておくため š バックプロパケーションの際に勾配を再計算すれば保持する

    必要なないが、その代わりに計算時間が⼤幅に増える š このトレードオフをバランスするため、⼀部のレイヤのみ保 持しそれ以外を破棄する Gradient Checkpointing が考えられた š GPUでこのフラグをつけてみるとメモリ使⽤量が半減
  27. 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/
  28. 訓練が中断する環境で気をつけ るべきこと š 今回のように巨⼤なデータセットを使い、訓練が中断される 環境で訓練を⾏う際、データセットのシャッフルを⾏わない と、「いつもデータセットの最初の⽅しか訓練に使われな い」という落とし⽳が…(途中で気がついた) š 当初は1個の巨⼤な(34.5GB)のTFRecordファイルを使ってい たが、複数のファイルに分割することに

    š bigbird/pretrain/run_pretraining.pyには複数のTFRecordをシャッ フルし、tf.data.Dataset.interleave()を使って混合するコードが 含まれている š 今回は、TFRecordにする前のテキストファイルをドキュメン ト数で均等分割し使⽤している(つまり、各ファイルのサン プル数には差異がある)
  29. BigBirdの学習結果は︖

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

  31. まとめ

  32. まとめ š ⻑⽂を扱える⾃然⾔語処理モデルBigBirdを⽇本語Wikipediaのデータ セットで事前学習させています š TPUが使えるモデルであれば、GPUよりも低コストでデータ並列を 使ったモデルの訓練が可能です š GCPでは、さまざまな⽅法でCloud TPUを使⽤することができますが、

    コストを抑えるにはPreemptible TPUを使⽤するのが効果的です š Cloud TPUを使った訓練では、GCE VMの能⼒はさほど必要ありません š それでも、⼤量のデータを学習させるには時間とお⾦がかかります š モデル訓練時の使⽤メモリを削減させる⼿法として、Gradient Checkpointingは有効です(ただし、訓練時間は⻑くなります) š Cloud TPU v2で⼗分なモデルでは、Google Colaboratoryが最強で、Pro にすれば中断も抑えられます
  33. 参考資料 š 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