こんにちは。レトリバの飯田(@HIROKIIIDA7)です。TSUNADE事業部 研究チームのリーダーをしており、分類エンジンの開発・マネジメント、検索分野の研究、チームマネジメントを行っています。
今回は、前回の記事から自己紹介に追加されている「分類エンジンの開発・マネジメント」について書いていきます。これは、チームで取り組みました。
経緯
レトリバでは、その前身のPreferred Infrastructure時代から自然言語処理に取り組んでいました。文書・文分類は基本タスクの一つであり、応用でも非常に多く取り組む機会があります。そのため、実験管理も含めたエンジンを10年以上前から作っていました。
このエンジンはコア部分がC++で作られていました。そのため、モデルの追加はC++で行うことになります。また、パラメータ管理をRailsで行っていた結果、新しいモデルのパラメータ設定はRailsで設定する必要がありました。しかし、機械学習や自然言語処理はpythonベースでエコシステムが発展しています。そのため、エンジンの刷新に高い開発リテラシーが必要となっており、長らく手がつけられない状態に陥っていました。
そこで、python及び付随する機械学習に関するライブラリを使用し、エンジンを刷新していく活動を行うことにしました。
いきなりつまづく
エンジンの利用方法としては、ユーザが自分達のデータで学習を行うことが必須条件でした。そのため、データ登録→モデルチューニング→実験結果管理→推論までを一貫して行えるようにする必要がありました。しかし、これら全てを全面的に作り替えるにはかなりの工数を要するだろうという見込みでした。また、本チームのメンバは開発に不慣れな面があることから、一部だけを作り替えられないかということで、コア部分を差し替えることで検討しました。
始めてみるとパラメータ管理がRailsによって行われるため、一部だけの作り替えが困難なことがわかり、全面的に作り替えることにしました。
あらためて計画
パラメータ管理部分も作る必要が出てきたため、MLOpsツールの使用を検討しました。先行調査していた中で、MLflowが今回のユースケースには最も合致していると考えました。
MLOpsツールの区分としてハイパーパラメータチューニング、実験管理、ワークフロー区分が「小さく始めて大きく育てるMLOps2020」で紹介されています。
この区分で考えた際、ユーザが自分達のデータで学習を行うため、最も重要な点は実験管理でした。なぜならば、既に機能として存在しており、かつユーザが色々やってこれがベストということで納得感が上がるためです。さらに、現時点ではパイプラインはそこまで重要ではありませんでした。なぜならば、データサイエンティストが作ったモデルを定期的に更新するというユースケースはなかったためです。また、データの管理がユーザ側にあるため、データの更新を踏まえた設計が極めて困難でした。さらに、オンプレミスでも使用することから、クラウドベンダーの提供するMLOpsツールに完全に依存することができないという制約もありました。以上から、実験管理に優れ、オンプレミスでも使用出来るMLflowを用いました。
次に MLflowをベースとして学習や推論をどのように一連の流れとして構築することが可能かを調査していきました。まず、実験管理部分である MLflow Trackingを使用することに決めました。 MLflow Trackingを用いることで、モデルの追加に伴った変数の追加の工数が不要になるため、Railsが管理しているという問題を解決することができました。
モデル管理、デプロイ周りについては、 MLflow Modelを一部使用することにしました。具体的には、MLflowにおけるpython-function形式のインタフェースに則りモデルを保存しておくことです。これによって、load時に実行idを指定して読み込むことが可能なので管理が楽になりうることがわかりました。
MLflow ProjectsやMLflow Model Registryはモデル更新に関わる部分と考えており、こちらはあまり使われる想定がないため、採用を見送りました。
アーキテクチャ
MLflowとして使う部分を決めた（実際はやりながら決まっていきました）ため、アーキテクチャを決めていきました。基本的には以下のように学習を行うサーバー、推論を行うサーバー、 MLflow Tracking、データを保存するDBの４つのサービスからなるアーキテクチャとしました。また、それぞれはdockerで隠蔽されています。
推論サーバーを一つのコンテナにしておくことで、スケールアウトにも対応しうる構造にしました。機械学習システムデザインパターンによるとデプロイ時のモデルの配布パターンとしてはモデルインイメージパターンとモデルロードパターンがありますが、後者を選択しています（というかやっていたらそうなっていました）。ユーザが自分達で学習することを前提にしているため、前者にすると毎回コンテナをbuildする必要があり、時間もかかる上システムも複雑になってしまうことが予見されたためです。
学習・推論のサーバーはFastAPIで実装されています。FastAPIは簡単に使用できて高速な処理が可能なため採用しました。APIドキュメントについては、Swaggerを使用できるため、ドキュメント作成の手間が大幅に減っています。なお、推論については MLflow models serve や MLflow deployments create -t sagemaker -m my_model などの利用もありうると考えており、場合によってはこれらに変える可能性もありそうです。
マネジメントスタイル
開発なので、基本的にはスクラムに則り行っていきました。長期展望やマイルストーンをラフに作りながら、具体的な計画はスプリント毎に落とし込んでいきました。メンバーは諸々の研究やアルゴリズムの検証を行いながら開発を実施する必要があったため、スプリントは長めに２週間とってやりました。始めはストーリーポイントの見積りと実時間のばらつきも大きかったですが、今ではかなり安定してきています。開発の実践としてのスクラムとしては、そこまで高度なことはできていませんが、可能な範囲で導入することで、ラフなマイルストーンから大幅な遅れもなくできています。
今後
弊社のYOSHINAではAWSを基盤として利用しているため、S3をモデルのストレージにすることや、GPUのスポット実行を行うことでBERTなどを学習する際も一時的にGPUを使うことでコスト低減になるように取り組んでいます。大規模言語モデルが大流行の昨今ですので、今後は少パラメータでの学習を取り入れて、少ないデータでも高精度なモデルを学習・推論できるようにしていきたいと考えています。
