ブログで紹介したようにStockmarkではMLOpsの考えを取り入れたシステム開発を行なっており、今回の開発でもMLOpsの概念に基づき、以下のようにCI/CDをできる限り自動化しました。
全件処理の際にも、上をベースとしたシステムを用いますが、さらなる処理の高速化を図りました。
数千万件のニュース記事の処理する際には、データの前処理や後処理にも、かなりの時間がかかります。そこでデータ処理のワークフローを(i)形態素解析などのデータの前処理、(ii)BERTによる処理、(iii)後処理の3段階に明確に分離しました。そして、CPUで計算を行うデータの前処理と後処理はCloud Runを用いて最大で100程度のインスタンスを用いて並列しました。Cloud Runはリクエストに応じて自動的にインスタンスが増えるので(自動スケーリング)、多数のインスタンスでの効率的な並列処理が可能になりました。さらにBERTの処理では2つのTPUのノードを用いました。このようにして、処理全体を高速化しました。
Cloud TPUを導入することで、特に全件処理時のパフォーマンスが著しく向上しました。GPUサーバを用いて処理を行なっていた時には、1週間程度の時間がかかっていましたが、Cloud TPUを用いることにより1日以内に処理を終えることができるようになりました。
以前は、日々のオペレーションの課題は、新規のニュース記事に対する日次のバッチ処理をいかに安定的に行うかということにありました。しかし、最近ではチームの人数の増加やプロダクトの成長により、新しいモデルを開発したり、既存のモデルを改善をする頻度が増えており、この際に発生するデータの全件処理をいかに早く行い、開発サイクルを早く回すかが重要な課題になっています。処理時間が大幅に短縮されたことで、お客様に新しい機能をより迅速に提供できるようになるといった効果が期待されます。
BERTの推論処理を行う上で、TPUがGPUに比べるとどの程度、処理時間やコストを削減できるかについても検証しました。ここでは、ベンチマークとして約50万のニュース記事を用いて、BERTの推論処理にかかる時間とコストをTPUとGPUで比較しました。タスクとしては、文章中から企業名を抽出するタスク(詳しくは過去の記事をご覧ください)を用いました。結果は下のようにまとめられ、TPUを用いるとGPUと比較して処理時間が73%、コストが51%削減することができることがわかり、大規模なデータの処理でのTPUの優位性がわかりました。
** 学習時のTPUとGPUの性能比較については過去のブログをご覧ください。
今回はTransformersのTensorFlowのBERTを用いました。まず、GPUを用いる場合には、モデルのロードは例えば以下のような感じです。
from transformers import TFBertForTokenClassification
model = TFBertForTokenClassification.from_pretrained(MODEL_PATH)
TPUを用いる場合には、上のコードを以下のように変えるだけです。(参考:https://www.tensorflow.org/guide/tpu)
import tensorflow as tf
from transformers import TFBertForTokenClassification
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=NODE_NAME)
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
with strategy.scope():
model = TFBertForTokenClassification.from_pretrained(MODEL_PATH)
NODE_NAMEはGCP上で作成したTPUノードの名前です。
TPUで最適なパフォーマンスを得るために下のような点に気をつけました。
TPUは8個のコア、64GBのメモリを持っており、GPUに比べると多くのデータを一度に処理できます。そのため、バッチサイズをできるだけ大きくすることで、処理が効率的になります。
GPUと比べてTPUは初期化やモデルのロードに時間がかかることにも注意が必要です。以前に、GPUで処理をしていた時は、その時に処理するデータをいくつかのファイル(例えば1万記事ごと)に分割して、それぞれのファイルにスクリプトを実行していました。しかし、TPUの場合には、データを分割して処理すると、それぞれのファイルを処理するたびに、上記のオーバーヘッドが発生してしまい、非効率です。そこで処理するデータはできるだけまとめて処理するようにしました。この処理により処理時間が50%削減できました。
BERTにデータを入力する際には、入力する系列の長さを揃える必要があり、短い系列にはダミーのトークンを足して、長さを一定にするPaddingという処理を行います。多くの場合、データの中で最長のものか、ある決められた長さに揃えます。しかしながら十分短い系列は、Paddingをすると本来の処理とは本質的には関係のないダミーのトークンが大部分を占めてしまうため、処理が非効率になってしまいます。そこでデータを長さが短い順に並び替えてから、バッチ化することにしました。この処理により、最初の方の短い系列のバッチは短い長さで揃えることができるので、Paddingによって挿入されるダミーのトークンの数を減らすことができます。この処理により、さらに処理時間が46%削減できました。
Stockmarkでは日々、膨大な数のニュース記事に対してBERTの推論処理を行なっており、そのための強力な処理基盤が必要でした。そこで、Google Cloud PlattformでCloud TPUを用いたBERT推論のための基盤インフラを開発し、より効率的にデータ処理をすることが可能になりました。
機械学習システムを開発する際には、機械学習のアルゴリズムの特性に応じたインフラ設計が、最適なパフォーマンス得るために重要です。そのためには、インフラ技術と機械学習技術のそれぞれを深く理解し、それをうまく組み合わせることが必要となってきます。このようなインフラと機械学習のコラボレーションは今後ますます重要になってくると考えられ、Stockmarkでは今後も引き続き取り組んでいきたいと思います。
現在、Stockmarkではインフラや機械学習のエンジニアを募集しています!興味のある方は是非応募してください!