Dwango Media Village(DMV)の佐々木です。大量のデータで大規模なニューラルネットワークのパラメータを求める深層学習は多くの計算コストを必要とします。一般的にはGPUを使った高速化などが行われますが、それでも実験には数日〜数週間の長い時間がかかってしまうことがあります。
この記事ではGoogleとの戦略的PoC(概念実証)として行ったTensor Processing Unit(TPU)による深層学習の速度向上の検証を紹介します。
TPUはGoogleが提供している機械学習用の計算デバイスです。GPUに比べて高い計算能力を備えており、大幅な学習時間の短縮が見込めます。DMVではGPUを使って学習実験を行なっていますが、モデルの規模やデータ量によっては数週間の期間がかかることがあります。もし実験の大幅な高速化が出来れば、より多くのハイパーパラメーターやモデル構造を試すことが可能になります。今回の検証では私たちが日頃扱っている深層学習モデルをTPU環境へ移行することを想定し、どのぐらいの速度向上が得られるのかを調べました。
TPUは Google Colaboratory 、または Google Cloud 上で提供されるVirtual Machine(VM)から利用できます。Google Colabは手軽にTPU環境を構築できますが、コードを管理しながら段階的に進めていくため今回はVMによる方法を採用しました。2022年春の現在、VMからTPUを使うには 2つの方法 が提供されています。「TPUノード」と「TPU VM」です。TPUノードはGoogle CloudのVMからTPUを管理する別のVM(TPUホストVM)へgRPCを使って接続する方法です。TPUホストVMへ接続するVMは一般のGoogle CloudのVMと同様にスペックを自由に変更できます。対してTPU VMはTPUホストVMへ直接接続するためgRPCによる接続が不要で高速な通信が期待できます。その代わりにTPU VMのスペック(CPUやメモリなど)は予め決められています。
TPUを使った計算では XLA(Accelerated Linear Algebra) と呼ばれる線形代数に特化したコンパイラを通して計算を実行します。Pythonのフレームワークとしては PyTorch , Jax , TensorFlow がXLAへのインターフェースを提供しています。今回は移行元の実装がPyTorchだったため、PyTorchのXLA対応パッケージ PyTorch XLA を使用しました。PyTorchからPyTorch XLAへの移行ではGPUデバイス(torch.device)をPyTorch XLAのTPUデバイスに置き換えるだけです。ただしXLAは非同期に計算を実行するため、Optimizer.stepの後などにデバイスとの同期を明示的に行う必要があります。また、1つのTPUには8つのコア（torchのデバイスに相当）が搭載されているため並列計算を行えます。この場合にはデバイス・プロセス間同期などの実装が追加で必要になります。
はじめにデータIOなしのスループット比較結果を紹介します。深層学習の学習計算はデータIO、前向き計算、後ろ向き計算、パラメータ更新の4つのステップで行われますが、この実験ではデータIO以外の3ステップについて単位時間あたりにどのぐらいのサンプルを学習できるか、すなわち学習スループットを調査します。
モデルはNAGAモデル、StyleGAN3の2種類を調べました。NAGAモデルはDMVの 麻雀AI NAGA の解析エンジンのコアとなるモデルで、 GPT-3 のようなFeed Forward層を主に使用しています。StyleGAN3はConvolutional Neural Network(CNN)ベースであり画像生成に用いられるモデルです。提案論文の 公式実装 をもとにしていますが、PyTorchのCustom Opが使えないためPyTorchの基本的なOpのみを使うよう修正を行いました。どちらもモデルもPyTorchで実装されており、これらをTPU環境へと移行しました。
比較元となるGPUはシングルGPU(RTX8000, CuDNN + CUDA)をPyTorchで学習した場合で、Torch JITなどの特別な最適化は行いませんでした。TPUではv3-8（8コア、GPUのデバイスに相当）を用いて、1GPUとの比較をよりフェアに行うため、1つのコアのみを使用しました。
上の図は様々なバッチサイズで学習を行った時のGPUとのスループット比を示したものです。横軸は底が2の対数スケールで示したバッチサイズ、縦軸は単位時間あたりに学習できたサンプル数がGPUに対してどのくらい多いかを示しています。スループットが1の場合はGPUと同じ速度、2は2倍高速化したという意味です。
NAGAモデルでは全てのバッチサイズにおいてTPU利用方式に関わらず速度向上がみられました。特にバッチサイズが大きいほど速度向上の恩恵が得られることがわかります。TPU利用方式に関してはTPU VMの方がより速度向上比が高いことがわかります。特にバッチサイズが\(2^{11}\)(=2048)の場合ではTPU-VMとTPUノードでは対GPUの速度向上比に5倍近い差があります。これは恐らくバッチサイズが大きくなるとTPUノードのgRPC通信がボトルネックとなることが原因でしょう。
対してStyleGAN3ではバッチサイズを増やしても大きな速度向上は得られませんでした。このモデルは活性化関数を適用する前に特徴量マップの解像度を2倍にするなどのNAGAモデルに比べてテンソル操作が多いモデル構造であることが要因であると考えられます。
ここまでの結果を踏まえて、TPU-VMを使ってNAGAモデルの大規模データIOを含めた学習実験を行いました。前の実験で分かったようにTPUはGPUに比べて高速に前向き計算〜パラメータ更新を行うことができます。Google Cloudの 公式ブログ にもあるように、このような場合にはデータIOがボトルネックとなります。そこで1つのTPUデバイスに搭載されている8つのコア全てを並列に利用することでデータIOを分散することにしました。具体的にはコアごとにデータIOを含む学習実行を行
うCPUプロセスとモデルを用意してモデル並列の学習を行いました。コアごとのバッチサイズは1024とし、一度の学習ステップで8192サンプルを学習できるようにしました。GPU側は前の実験と同じく1GPU、バッチサイズは512としました。
上の図は学習終了までにかかった時間から算出した1秒あたりの学習したサンプル数の比較です。GPU比で11.23倍の高速化を行うことができました。NAGAモデルの学習はGPU環境で20~30日程度かかっていたのですが、TPUへ移行することで数日へと大幅に短縮できました。
高速化に加えて嬉しい別の効果も得られました。バッチサイズを増やしたことサンプルの独立性が向上し、学習が安定したのです。学習済みモデルを使って麻雀の自己対戦などの評価を行ったところ、解析エンジンとしての性能が大幅に向上しました。このモデルはNAGAの新バージョンとして現在利用されています:
この記事ではTPUを用いた深層学習の高速化の検証結果を紹介しました。データIOを行わないスループットの比較調査ではバッチサイズに応じて速度向上効果が得られること、またモデルの種類によって高速化の比率が異なることがわかりました。またマルチTPUコアによる実験ではGPUに比べて約11倍学習速度が向上しました。
TPUを使うことによって高速化を見込める一方で、気をつけるべき点もあります。GPU用のPyTorch実装をTPU向けに移植する場合、ほとんどの場合はコード追加は不要ですがXLAが対応していない命令を避ける必要があります。またモデル構造とバッチサイズによっては速度向上がそこまで得られない場合もあります。そしてTPUの複数コアを使って高速化を望む場合にはさらに実装を修正する必要があります。とはいえ今回のように数十日もかかる計算を数日で終わらせられることもあるため、検討する価値は大いにあると思います。
Google Cloud、Tensor Processing Unit、Colaboratory および TensorFlow は Google LLC の商標です。
Kazuma Sasaki
