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

Transcript

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

   View Slide

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

   View Slide

3. Cloud TPUとは︖
   

   View Slide

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
   

   View Slide

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

   View Slide

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）が可能
   

   View Slide

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
   

   View Slide

8. 今回のモチベーション
   

   View Slide

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

   View Slide

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

    View Slide

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
    

    View Slide

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

    View Slide

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

    View Slide

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
    -
    

    View Slide

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

    View Slide

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

    View Slide

17. GCEでのモデル訓練
    

    View Slide

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

    View Slide

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
    

    View Slide

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モデルを使⽤する際には修
    正が必要である
    

    View Slide

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
    

    View Slide

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

    View Slide

23. 予算枯渇と
    救世主の出現
    

    View Slide

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

    View Slide

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

    View Slide

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

    View Slide

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/
    

    View Slide

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

    View Slide

29. BigBirdの学習結果は︖
    

    View Slide

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

    View Slide

31. まとめ
    

    View Slide

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

    View Slide

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
    

    View Slide