Colab TPUでHuggingface accelerateを用いて文章分類を試してみる

はじめに

以前のポストColab TPU で transformers Flax+JAX で文章分類を試してみるで、 TPU を使って Flax/JAX モデルで文章分類を試しました。 今回は、Pytorch 版のモデルで TPU(v2-8)を用いて文章分類を試してみます。Pytorch 版モデルで TPU 学習を簡単に行うための accelerate パッケージを利用します。

accelerate

accelerateは、pytorch 版のモデルを、マルチ GPU/TPU/fp16/deepspeed での学習を簡略化するためのパッケージで Huggingface によって公開されています。

accelerate はいろいろな実行方法があるのですが、今回は既存の学習スクリプトに対して CLI のラッパーのような形で実行する方法を試しました。 transformers の examples のなかには accelerate に対応したものがあり、これを使うと容易に TPU で学習できます。

データセット

前回と同様にlivedoor ニュースコーパスを用いて文章の分類をします。 学習用のデータとしては、ここに示した方法により生成した train.csv,dev.csv を使います。

基本方針

基本的にはaccelerate configで TPU 向けの設定をしたあとに、学習用スクリプトをaccelerate launchで実行します。

accelerate launch ./transformers/examples/pytorch/text-classification/run_glue_no_trainer.py

Colab notebook

Google Colab で実行するための手順を以下に示します。

Colab の全体はここに置いてあります。

python パッケージなど

!pip install transformers accelerate datasets==1.11.0 fugashi ipadic

いつもの transformers や datasets に加え、accelerate をインストールします。 datasets 1.12.0 や 1.12.1 を使うとデータセットのコンテンツが 2 度読み込まれてしまうという問題があったため、1.11.0 を使用しています。

!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl

pytorch を TPU で使うためのパッケージ torch_xla もインストールします。

transformers

!git clone -b v4.10.0 https://github.com/huggingface/transformers.git

example スクリプトを使うので transformers をクローニングします。

データセット

!ln -s /content/drive/MyDrive/colab/ldn .

今回は Google Drive に置いてあるデータセット(train.csv,dev.csv)を Drive をマウントして利用します。

accelerate の設定

!accelerate config
WARNING:root:TPU has started up successfully with version pytorch-1.9
In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): 3
What is the name of the function in your script that should be launched in all parallel scripts? [main]:
How many processes in total will you use? [1]: 8

accelerate で TPU を使うように設定します。最後のプロセス数は今回使用する TPUv2-8 のコア数(8)としています。 なお、accelerate config を一度実行したあとは以下のようにファイルに保存されて次回からは実行不要となります。

# cat ~/.cache/huggingface/accelerate/default_config.yaml
compute_environment: LOCAL_MACHINE
deepspeed_config: null
distributed_type: TPU
fp16: false
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
num_machines: 1
num_processes: 8

学習実行

MODEL="cl-tohoku/bert-base-japanese-whole-word-masking"
OUTDIR="./output"
BS=8
SEQLEN=512
EPOCH=5
!accelerate launch ./transformers/examples/pytorch/text-classification/run_glue_no_trainer.py \
       --model_name_or_path {MODEL} \
       --output_dir {OUTDIR} \
       --per_device_train_batch_size {BS} \
       --per_device_eval_batch_size {BS} \
       --max_length {SEQLEN} \
       --train_file ldn/data/train.csv \
       --validation_file ldn/data/dev.csv \
       --num_train_epochs {EPOCH}

バッチサイズは 8 としました。16 では OOM エラーとなりました。

環境など

今回の確認環境は以下のようになりました。

• Memory: 13GB (Colab Pro)
• TPU: v2-8
• python: 3.7
• pytorch: 1.9.0+cu102
• transformers: 4.10.2
• accelerate: 0.4.0
• datasets: 1.11.0
• torch-xla: 1.9

結果

ランタイムのタイプを TPU にして、すべてのセルを実行します。 結果は以下のようになりました。

09/16/2021 00:18:28 - INFO - run_glue_no_trainer - ***** Running training *****
09/16/2021 00:18:28 - INFO - run_glue_no_trainer -   Num examples = 5893
09/16/2021 00:18:28 - INFO - run_glue_no_trainer -   Num Epochs = 5
09/16/2021 00:18:28 - INFO - run_glue_no_trainer -   Instantaneous batch size per device = 8
09/16/2021 00:18:28 - INFO - run_glue_no_trainer -   Total train batch size (w. parallel, distributed & accumulation) = 64
09/16/2021 00:18:28 - INFO - run_glue_no_trainer -   Gradient Accumulation steps = 1
09/16/2021 00:18:28 - INFO - run_glue_no_trainer -   Total optimization steps = 465
 20% 93/465 [04:45<36:47,  5.93s/it]09/16/2021 00:23:34 - INFO - run_glue_no_trainer - epoch 0: {'accuracy': 0.85546875}
 40% 186/465 [06:21<03:29,  1.33it/s]09/16/2021 00:25:03 - INFO - run_glue_no_trainer - epoch 1: {'accuracy': 0.9192708333333334}
 60% 279/465 [07:50<02:17,  1.35it/s]09/16/2021 00:26:32 - INFO - run_glue_no_trainer - epoch 2: {'accuracy': 0.9381510416666666}
 80% 372/465 [10:42<01:13,  1.27it/s]09/16/2021 00:29:24 - INFO - run_glue_no_trainer - epoch 3: {'accuracy': 0.9446614583333334}
100% 465/465 [12:10<00:00,  1.39it/s]09/16/2021 00:30:52 - INFO - run_glue_no_trainer - epoch 4: {'accuracy': 0.9453125}
Configuration saved in ./output/config.json
Model weights saved in ./output/pytorch_model.bin
100% 465/465 [12:29<00:00,  1.61s/it]

5epoch の学習は 12.5 分程度で終了となりました。Flax/JAX よりも少しはやくなりました。 Flax/JAX では batch_size_per_device=8 にすると OOM となったため 4 としていたのですが、その違いかもしれません。 Accuracy はちょっと下がっていますが誤差の範囲でしょうか。

以前のローカル GPU(RTX2080Ti)、前回の TPU(Flax/JAX)で試したときとの比較が以下の表です。

おわりに

transformers + pytorch + accelerate で TPU を使って文章分類の学習を試してみました。 accelerate を使うと既存のスクリプトで容易に TPU 対応できることを確認しました。