はじめに
みなさんはじめまして、リーガルテックベンチャーMNTSQの取締役の堅山です。この度弊社でテックブログを開設することになり、その第一号として記事を書いています。弊社の取り組む「法務」の世界はエンジニア・リサーチャーの方々から見ると縁遠いことも多いかなと思いますが、そういった方に向けてリーガルテック企業が実際何をやっているのか発信していけたらなと思っています。特に、
といったことについて弊社メンバーで書いていければと思っています。
目次
今日は、言語モデルのpretrainについて書こうと思います。
言語モデルの進歩は目まぐるしく、BERT以来様々なモデルが発表されています。また、huggingface/transformersなどのライブラリを使えば、すでに学習済みの言語モデルの利用については非常に手軽です。一方でそのような新しめの言語モデルをスクラッチで学習する方法について、私の探した範囲では実はあまりウェブに情報がなく、いろいろ試行錯誤が必要でした。その結果をハマりそうなポイントなどを含めて記事にまとめてみました。
以下の内容について記しています
- fairseqを使って言語モデルをスクラッチで学習する
- pytorch/xlaを使って、fairseqをTPU上で動かす
- Preemptive Instancesを使ってコストを抑える
- Tips: Tensorflow Research Cloud (TFRC)に応募してみよう
こころがまえ
そもそも、なんでこんなにpretrainが大変かと言うと、
- 言語モデルを扱うライブラリはいままさに活発に開発されており、APIや実行手順が目まぐるしく変わる
- 計算量が多いため、分散学習をしたり、TPUのようなデバイスを使う必要がある。すると、インフラ周りでさらに扱わないといけないソフトウェアの範囲が増える
という状況がありそうです。おそらくこの記事はすぐ古くなってしまいます。お試しいただく際は、ある特定のブランチを使わないといけないとか、古いバージョンを使わないといけないとか、環境変数を新たに設定しないといけないなど、いろいろなワークアラウンドが必要になる可能性があります。ドキュメントにのっていない困ったことがあったらrepositoryのissueを目を皿のようにしてみましょう。
fairseqを使って言語モデルをスクラッチで学習する
fairseq はpytorchベースの言語モデルを扱うためのライブラリです。 github.com このあたりのライブラリとしては、前述のtransformersが有名だと思いますが、fairseqにはなんとCLI Interfaceがついてきます。カスタマイズ性はtransformersのほうが高いのですが、fairseqの場合は前処理を間違えていきなり変なことになる、みたいな不安はなさそうです。例えば、transformersで頑張るとすると、tokenizerの処理の方法がこのバージョンでは変わっていてtutorialのままでは正しく動かないであるとか、transformersのDatasetクラスやTrainerクラスをきちんと理解しないといけないなど、柔軟な分いろいろ大変だったりします。
fairseqの使い方はシンプルです。たくさんテキストを用意しましょう。今回私の場合は、train.txtとtest.txtというファイルに纏めておきました。今回使ったデータセットは実は前後の文に相関がないため、BERTのようなNSP (Next Sentence Prediction)タスクが入っているモデルを学習していないのですが（Robertaを使いました）、NSPを学習する場合は文のペアを用意しないといけないため、たぶん何かしないといけないと思います（未確認です）。短いテキストファイルがたくさんある場合、何も考えずにファイルを結合してしまうと、NSPの部分でだめになってしまうかもしれないので、気をつけてください。fairseqは単語分割については面倒を見てくれません。bpeをする場合はsubword_nmtなどのbpeツールを使いましょう。 github.com 今回のデータセットではわたしはbpeを使ってないので、その部分は省略します。bpeによる単語分割が終わったら、fairseq-pretrainをかけます。fairseq-preprocessはテキストデータをバイナリ化して保存します。今回はすでにspaceでtokenizeされているのでspaceを使いました。
fairseq-preprocess \ --trainpref train.txt --validpref test.txt \ --workers 8 --tokenizer space
ここで気をつけてほしいのが、fairseqのCLIのドキュメントにあるオプションは、すべて有効なわけではないということです。fairseq-preprocessにも--tpuとか指定できるのですが、前処理にtpuを使えるはずもなく、ドキュメントにかかれているオプションが実際にその段階で有効なのかは確認してみてください。
さて、preprocsessが終わったら、pretrainができます。TPU上での学習は、以下のようなコマンドを走らせればOKです。（わたしが実際に使ったコマンドです）TPUでない場合は--tpuオプションを外して、--distributed-world-sizeをGPUの台数にすればそのまま動くはずです。（nvidiaのapexやncclなどのインストールが必要な場合があります）
fairseq-train ./data-bin/ \ --tpu --task masked_lm --criterion masked_lm \ --arch roberta_base --sample-break-mode complete \ --optimizer adam --adam-betas '(0.9,0.98)' --adam-eps 1e-6 \ --clip-norm 0.0 --lr-scheduler polynomial_decay \ --lr 1e-5 --warmup-updates 3000 --dropout 0.1 \ --attention-dropout 0.1 --weight-decay 0.01 \ --update-freq 1 --log-format simple --log-interval 10 \ --tokenizer space --train-subset=train --save-dir ./checkpoints \ --keep-interval-updates 10 --tokens-per-sample 32 \ --max-sentences 1792 --valid-subset test --max-epoch=25 \ --num-workers=2 --tensorboard-logdir ./runs --keep-last-epochs=25 \ --keep-best-checkpoints=1 --seed 0 --distributed-world-size 8 \ --max-positions 32 > /path/to/log/`date +\%Y\%m\%d\%H\%M\%S`.log 2>&1
今回、メインのパラメータは、GCPのチュートリアルからとってきています。一方で、いくつかのパラメータは実際のデータに合わせる必要があります、重要なパラメータとしては以下があります。
./data-bin/: 第一引数はモデルの保存先を指定します。これがデフォルトのようです--lr: learning rateはデータ、バッチサイズなどで大きく適切な値が変わります。1e-5がよく使われているらしく、わたしのケースでもうまくいきました。1e-4, 1e-6ともに学習が進まなかったです--warmup-updates: このパラメータの回数分のminibatchを処理した段階でlrが最大値になります。1epochのminibatchの数は後述のパラメータで大きく変わるので、変更し忘れないように注意して下さい--update-freq: 何回かのminibatchの勾配をまとめます。バッチサイズを擬似的に大きくしたいときに使います
以下がbatchsizeに関するパラメータです
--tokens-per-sample: 1文の最大長を指定します--max-sentences: 1つのバッチに詰め込まれる文の数を指定します。tokens-per-sample x max-sentencesがだいたいTPU上で専有するメモリサイズなので、ここを最大にできるようにしていくのが重要です。紛らわしいのですが、--max-positionsはpositional embeddingの最大値を決めるものなので、--tokens-per-sampleと同じ値を設定しておけばよさそうです。
pytorch/xlaを使って、fairseqをTPU上で動かす
さて、そもそもTPUを使うには、GCP上でVMを立てる必要があります。TPU(-v3)には8つの独立したプロセッサがあるので、fairseqも少なくとも8プロセスをデータ送信に使います。1つのTPUのプロセッサにnum-workers x 8分のプロセスでデータを送っているようです. いまのバージョンのfairseqは、Pythonのmultiprocessingで全データをメモリに読み込んでいるようで、データサイズ x num-workers x 8ぐらいのメモリを専有していきます。なので、めちゃくちゃVMのメモリサイズが必要です。わたしの場合は、num-workers=1で、10 CPU、256GBメモリのマシンを使いました。SSDなどもろもろで$2.5/hourくらいかかってしまいます。結構高いですね。後で説明するpreemptive instancesを使うとこれが$0.6くらいになります。
TPU上で計算するにはTPUにコンパイルしたモデルを送信しないといけません。これに、pytorch/xlaを使います。
pytorch/xlaはpipなどではなく、VMのImageでインストールされているものを使います。VMを作成するときに、OSイメージを選択する画面でDeep Learning > pytorch/xlaを選ぶ必要があります。このVMにはanacondaがインストールされており、その中の仮想環境にtorch-xlaが何バージョンかインストールされています。conda env listをうってみると、一覧が出てきます。私の場合はtorch-xla-1.6のvenvを使いました。
conda activate torch-xla-1.6
前述のGCPのチュートリアルのとおり、環境変数を設定します。
export TPU_IP_ADDRESS=XXX.XXX.XXX.XXX export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
ご存知の通り、anacondaはIntel C++ Compilerでコンパイルされたnumpyなどがついてくるのですが、なぜかそのバグを踏み抜いている場合があり、fairseqを実行する際に環境変数を渡してあげないとエラーになりました。さらに、
export MKL_THREADING_LAYER=GNUTPUとしましょう。
TPUを起動しないと、IPアドレスが決まりません。GCPの画面で新しいノードを起動しましょう。ここで注意するべきなのが、TPUのソフトウェアの指定です。使うpytorchのバージョンに対応したものを選びましょう。普通のバージョンを選ぶときどうしません。わたしの場合はpytorch-nightlyでOKでした。
あとは前述のコマンドを打てば、fairseqが動くはずです。なお、試行錯誤の過程でたくさんエラーが出ると思いますが、エラーの多くはおそらくメモリ関連のものです。わたしの場合は、multiprocessingのエラーはホストマシンのメモリが、TPUのエラーはTPUのメモリが足りないエラーでした。
Preemptive Instancesを使ってコストを抑える
Preemptive instanceは最大24時間しか連続で走らせることができず、途中で突然シャットダウンされる可能性があるインスタンスです。fairseqは定期的にモデルをcheckpointとして保存してくれるので（--save-dirで保存場所を指定します。保存頻度・残しておく数も、--keep-last-epochsや--keep-best-checkpointsといったオプションで指定ができます）、突然シャットダウンされても、適切なところから再開できます。再開を勝手にしてもらうために、設定をしていきましょう。
まず、VMがシャットダウンされたら再起動されるようにします。gcloudコマンドからVMは起動できるので、適当な手元のパソコンからwatchコマンドで定期的にstartコマンドを送りましょう。すでに起動している場合は、何も起きないので安心です。
watch -n 600 gcloud compute start [your-vm-name]
VMが起動したら、fairseqを自動で起動してもらう必要があります。いろいろ方法がありますが、一番手軽なのはcronを使うことでしょう。単にcronを使うと定期的に新しいfairseqを立ち上げてしまうので、flockコマンドで抑制します。10分おきに再起動するなら以下の感じでしょう。新しいログファイルも作成してくれます。VMの中でcrontab -eで以下の様なコマンドを設定してください。
* 10 * * * * flock -n /home/xxx/lock \ source /anaconda3/envs/torch-xla-1.6/etc/conda/activate.d/env_vars.sh; \ /anaconda3/envs/torch-xla-1.6/bin/fairseq-train ... 2>&1 > `date`-cron.log
なお、condaのvenvをcron内でactivateするには上記のようにenvファイルを読み込む必要があります。torch-xla-1.6は使っているcondaのvenvで適宜調整してください。
パフォーマンス
今回の実験はTPU v3–8を1台使って行いました。わたしの場合、RTX-2080ti 1台に対して、TPU v3-8 1台でだいたい5倍速ぐらいの学習時間でした。Cloud TPU profilerを使うと、いつTPUのプロセッサが空いているかなどがわかるので、もっと最適化ができるかもしれません。
学習済みモデルを読み込んでみよう
先程のコマンドで指定した./data-binフォルダにモデルファイルが保存されていきます。もっともvalidation lossが低いモデルがcheckpoint_best.ptで保存されています。
from fairseq.models.roberta import RobertaModel model = RobertaModel.from_pretrained( './checkpoints/', checkpoint_file="checkpoint_best.pt", data_name_or_path="./data-bin/xxxx/", ) model.eval() model.fill_mask("今日はいい<mask>ですね")
最後の行はタスクに合わせて変更して頂く必要がありますが、要はmasked language modelを試すには、<mask>というトークンで単語を置き換えればOKです。こんな感じの返答が返ってくるはずです。
[('今日はいい天気ですね', 0.6828255653381348, '天気'),...]
Tips: Tensorflow Research Cloud (TFRC)に応募してみよう
pretrainでつらいところは、コストです。TPU-v3の利用料金はなんと$8/hourです！ですが安心してください。まず、TPUにもPreemptiveオプションがあり$2.4/hourまでコストを下げれれます。TPUもgcloudコマンドで開始できるので、上記のスクリプトをいじれば使えるようになるはずです。
もう一つ、TPUのコストを節約する方法があります。「Tensorflow Research Cloud (TFRC) 」というGoogleが公開しているプログラムに応募することです。これに応募してプロジェクトが認められると、私が以前申し込んだときは以下の資源へのアクセスが30日間認められました。
- TPU v2-8 x5 @ europe-west-4
- TPU v3-8 x5 @ europe-west-4
- Preemptive TPU v2-8 x100 @ europe-west-4
Googleの寛大なプログラム、すごいですね…。TFRC プログラムの参加者は、TFRC を利用した研究結果を、査読を受けた論文、オープンソース コード、ブログ投稿などの形で全世界に公表することが求められます。条件を満たせそうであれば、応募してみるのもよいのではないでしょうか。
なお、これだけTPUが使えるならなんでもできそうですが、残念ながら複数のTPUを使ってfairseqを学習することは現時点ではできないようです。おそらくtensorflowベースで頑張ると、preemptive TPU x100の真価を発揮できるのでしょうが、なかなか大変そうです。
