日本語CTRLを1から学習する - 5

はじめに

前回まで学習の基本的な設定を行ってきた。
今回はCTRLの学習高速化について調査・検証する。

Deep Learning学習高速化について

Deep Learningの学習高速化については以下のようなアプローチが考えられる。

Neural Netの軽量化
プログラムの最適化
並列化

Neural Netの軽量化

パラメータの削減

層の数・潜在状態数を削減することで学習は速くなる。
ただ、パラメータ数を減らしすぎるとNeural Netの表現能力が下がり、
パフォーマンスにも影響が出るため不用意に下げることはできず、
どの程度まで下げて問題ないか検証する必要がある。

アーキテクチャの軽量化

BERTはTransformerを大量に重ねたアーキテクチャであり、
表現能力の高さと引き換えに学習・推論の速度が問題となる。

それを解決するためにBERTを高速化したアーキテクチャがいくつか
提案されている。

DistilBERT¹ : DistilationでBERTを軽量化
ALBERT² : embeddingと層間パラメータ共有で軽量化

DistilBERTは元の学習済みモデルがないとDistilationできないため今回は利用できない。
ALBERTの層間のパラメータ共有は良さそうに見えるが、文章生成系で有効かは不明。

プログラムの最適化

Deep Learningの学習プログラムの最適化は以下のような要素がある

ライブラリの選択
マシンに対する最適化

ライブラリの選択

以下ではtensorflowとpytorchの2つのライブラリについて速度面で比較する。
※筆者は以前tensorflow1系を使っていたが、ここ1年ほどはpytorchを利用してきたため
最新のtensorflowの事情には詳しくない。

pytorch と tensorflowにおけるtransformersライブラリの性能比較 (推論速度) https://medium.com/huggingface/benchmarking-transformers-pytorch-and-tensorflow-e2917fb891c2

上記記事いよると推論速度的にはtensroflow, pytorch間ではそれほど大きな違いはない。
ただ上記はあくまで推論速度の推定である。

Titan RT上での各ライブラリの比較
https://medium.com/syncedreview/tensorflow-pytorch-or-mxnet-a-comprehensive-evaluation-on-nlp-cv-tasks-with-titan-rtx-cdf816fc3935

タスクによって学習速度の上下は入れ替わっている。
ただ、この比較はgithubから得たソースコードによって行っているらしいので、
個々のソースコードがどの程度最適化しているかが均一ではない。
またtensorflowのバージョンも1.10~1.13とやや古いものとなっているため2系でどうなるかは不明。

どちらのライブラリも最終的にcudnnでGPUを利用しているので、
GPUで計算しているうちはそれほど大きな差は出ないと思われる。
(TPUを使うと変わってくるかもしれない。)
当面は使い慣れたpytorchを使っていこうと思う。

混合精度による高速化

Neural Net内の計算はfloat32のデータ型で計算されることが多いが、
これをfloat16など低精度の型に落とすことでメモリ使用量が減少する。
Neural Net内の全てのパラメータの型を変換してしまうと性能が低下して
しまうため、必要な部分だけ高精度で必要ない部分は低精度とする
混合精度(mixed precision)という手法が利用される。

Nvidia apex³ は自動でpytorchのmixed precisionを実現してくれるライブラリである。
Transformer系のモデルはメモリの使用量も多くなるためこれを用いることでより大きなバッチサイズを利用できる。

並列化

並列化は複数のCPUまたはGPUを利用して1つのモデルを学習することである。
Deep Learningの場合は複数のGPUを利用して学習を高速化することは珍しくない。
pytorchは複数GPUによる学習をサポートしている。

参考
https://qiita.com/arutema47/items/2b92f94c734b0a11609d
https://nomoto-eriko.hatenablog.com/entry/2019/06/03/174633
http://kaga100man.com/2019/12/08/post-110/

基本的にはほとんどコード変更なしで並列化してくれるようだが、
BatchNormalizationの挙動などには注意が必要なようである。
また、GPU間の負荷分散も考える必要がある。

検証

今回はGoogle Colaboratory上でapexの検証を行った。
学習したのは前回と同様の小さめのCTRLモデル。

変更点1 apexをインストール, import

!git clone https://github.com/NVIDIA/apex
!cd apex;pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .
from apex import amp

変更点2 モデルとオプティマイザーの設定後にコード追加

model, optimizer = amp.initialize(model, optimizer)

変更点3 lossの逆伝播部分を修正

#  loss.backward()
  with amp.scale_loss(loss, optimizer) as scaled_loss:
      scaled_loss.backward()

以上の修正を行った結果、1秒あたりのIter数が2.1→6.2に上昇した。
ただこの性能向上の原因はapexを導入したことだけではない可能性がある。
Google ColaboratoryでGPUを利用する場合性能の異なる2種類のGPU
がランダムに割り振られるらしい⁴。今回の実行では性能の良いTesla T4の
GPUが割り振らていたので、GPUの違いにより前回より性能がよかった可能性も残る。
(前回実行した時はGPUは確認していなかった。)

ただ、以下のページにあるようにT4とK80(ColaboのもうひとつのGPU)の間の性能差は
3倍もないようなのでapexによりある程度性能向上しているものと思われる。
https://electric-blue-industries.com/2019/10/30/hardware-gpu%E3%82%B9%E3%83%9A%E3%83%83%E3%82%AF%E6%AF%94%E8%BC%83%EF%BC%88nvidia%EF%BC%89/

学習パラメータとマシン構成について

学習するマシンの構成を決める上でハイパーパラメータを決める必要がある。

学習リソース例

NICT BERT 日本語 Pre-trained モデルではNVIDIA Tesla V100を32枚使い
1週間程度かけてBERTを学習している。バッチサイズは4096。
(※GCPで Tesla V100 8枚を1ヶ月借りる費用は1万ドルほどである...)

京都大学黒橋村脇研ではGTX-1080Tiで1ヶ月ほどかけてBERTを学習している。
これくらいならば現実的にできなくはないと思うが、 max_seq_lenが128とやや小さく感じる。

学習データの読み込み

今回学習するデータはWikipediaデータ(3GB程度)、青空文庫データ(1GB程度)、
その他となり、おそらく10GB未満になると思われる。
トークン化などの前処理を終えたデータならばメモリに乗りそうなので
データ読み込みのオーバーヘッドをなくすために32GB程度のメインメモリを持つ
CPUとする。 GPUへはbatch単位でデータを送るため転送量はそれほど問題ないと思われる。

想定パラメータ

モデルサイズに関連するパラメータは以下のようなものがある。

語彙数(30000程度)
Transformer層数(6~12層?)
潜在状態数(384~768?)
学習中の文章の最大長(256程度)
バッチサイズ (?)

これらのパラメータは実際に使用するマシンのGPUメモリに乗るかなどを検証しなければ分からない。
できれば2GPUマシン以内で2週間程度で終わる規模の学習にしたい。
バッチサイズは元のBERTが256(Cloud TPU 16コア) CTRLが1024 (Cloud TPU 256コア)で学習しているがそれに合わせると計算リソースが恐ろしくかかりそうなので利用可能なリソースに合わせて考えたい。

まとめと今後

今回は学習の高速化手法について考えた。
混合精度学習、並列化などの高速化手法と
マシンリソースを考慮して最適なパラメータを検討したい。
ただ、今回の検証の目的はとりあえずやってみることなので
あまりチューニングに時間を使うつもりはない。
今後は学習するデータを決めて、どの程度の計算ができるかをGCP上で検証したい。

ML Over the Horizon