見出し画像

MLflow AI Gatewayを用いたプロンプトエンジニアリングの紹介

Atsushi Hara

先日、社内で大規模言語モデル（以下、LLMとします。）のパフォーマンスを比較するための取り組みについて社内勉強会で発表し、MLflowとMLflow AI Gatewayの紹介をしました。この記事では、そこで発表した内容のうち、どのような場面でMLflow AI Gatewayが役立つのかについて紹介したいと思います。


LLMのパフォーマンスの観点

LLMを用いた生成AI（以下、GenAIとします。）のトレンドは、今もなお大きな反響があり、より発展的な分野になると期待されています。 一般的に、GenAIが生成するアウトプットは、入力のプロンプトやtemperatureなどのパラメータに影響されると知られています。このアウトプットの質は、どのようなタスクを解くのかに依存しますが、様々な測られ方が考えられます。アウトプットに対して、LLM以前から自動生成された文章に対して使われてきたROUGEやBLEU、BERT SCOREなど計算可能な指標もあります。一方で、プロダクトに関わる人が評価するなど考えられます。 他にはどのようなパフォーマンスが考えられるでしょうか？よくあるのは、レスポンスのレイテンシー、スループットなどシステムも含めたパフォーマンスなどが挙げられると思います。 このブログでは、GenAIのアウトプットをツールを使って効率的に評価するかに焦点を当てていきたいと思います。

MLflow AI Gatewayとは

MLflowは、Databricks社が主に開発、メンテナンスしているOSSで、機械学習モデルの継続的な開発、パフォーマンスの向上、デプロイメントをサポートをサポートしてくれます。そのため、MLの開発現場でよく使われるツールの一つです。LLMが登場してきたことにより、これまでの機能ではカバーできない場面が増えてきました。そこで新たな機能として開発されたのがMLflow AI Gatewayです。

なぜLLMの開発でmlflow AI Gatewayが役立つのか

MLflow AI Gatewayで提供される主機能は2つだと私は解釈しています。

1つ目は、SaaSとして提供されているLLMへのアクセスのサポートです。
これまでのMLとLLMで大きく異なるのは、LLMの特徴の一つである汎用性です。汎用性が高いため、従来のMLよりもSaaS APIとして利用できるサービスが増えている印象があります。一方で、同じプロンプトを与えても各サービスのアウトプットは異なり、解きたい課題に対してどのサービスがよいかは異なります。また、運用しているサービスが特定のクラウドベンダー上で動作している場合、組み込みやすさなどを考慮したい場合なども考えられます。
上記のような背景を踏まえ、様々なAPIを試したい場合、都度コードを変更する必要があるかと思います。各APIのインターフェースは、当然ですが、似て非なる部分があります。どこまで汎用性をもたせつつ実装するのかは、悩みどころですよね。
この部分をMLflow AI Gatewayはサポートしてくれます。有名なOpenAI社のChatGPTやAWS BedrockなどのLLM APIへのアクセスはサポートされているので、必要な設定情報とコンフィグファイルを記述することで、MLflow AI GatewayがAPIへのアクセスをハンドリングしてくれます。また、Llama2のようにOSSとして提供されているLLMもHugging Faceのインターフェースに沿って提供することで、MLflow AI Gatewayからアクセスすることができます。

2つ目は、一つのツールで比較できるUIです。
MLflowは長らくMLの開発現場をサポートしてきました。LLMを開発する場合においても、使い慣れたツールで一貫して比較できることのほうが望ましいと考えられます。また、LLM開発スタイルはこれまでのMLの開発スタイルから大きく変化すると私は考えており、ここに1つのツールを使うメリットがあると思います。

LLMの開発は従来のMLの開発と何が異なるのか

これまでのMLの開発は、大部分がデータエンジニアやMLエンジニアの専門領域でした。なぜなら、大まかなMLモデルの開発のプロセスは、課題の設定から始まり、その設定された「特定の課題」に対して大規模なデータから規則性を見つけるMLモデルを開発し、想定される利用範囲での未学習のデータに対する汎用性を高めるためにMLモデルをチューニングする。というようなものです。これらの一連のプロセスでは、学習後に外部から入力を与える自由度は決して高くはありません。
LLMの場合、その汎用性と柔軟性の高さからMLモデルに内包されていた「特定の課題」は、ある程度外部からコントロールできるものに変わりました。そのため、LLMが課題を解くパフォーマンスは、モデルに内包されているパラメータだけでなく、学習後のモデルに対しても可変なプロンプトやTempertureやTop Kなどのパラメータにも大きく左右される部分になりました。もちろん、LLM自体のファインチューニングや学習プロセスなどは依然としてデータエンジニアやMLエンジニアの専門領域です。しかし、これらの学習後に与える各入力がLLMにどのような影響を与えるのかを理解していれば、開発チームの誰しもがLLMのパフォーマンスチューニングに関われるようになったといえます。
これは素晴らしいことである一方で、一つのペインポイントにもなり得ます。それは、各々が各々の環境で実験してしまうことで、比較と再現が難しくなる。ということです。このようなケースが発生することを未然に防ぐことは重要だと考えられます。一つのツールに実験したプロンプトやプロンプトテンプレートがあり、プロジェクトに関わるメンバー誰しもが比較できることは継続的な開発の観点からとても大切であると言えます。

MLflow AI Gatewayのシステムの構築について

今回紹介した内容を実際に検証するために、AWSのMLflow と Amazon SageMaker による機械学習のライフサイクル管理の記事を参考にしながら、社内環境でMLflow AI Gatewayを扱える環境を構築しました。社内環境に建てたMLflowのUIを載せておきます。

社内環境に建てたMLflow endpoint

残念ながら実際のコードをここで紹介することはできませんが、構築したシステムのアーキテクチャ図と概略について紹介したいと思います。

システムアーキテクチャ図

今回は、以下の3点の工夫を行なってシステムを構築しました。

  1. 他のプロジェクトでもすぐに利用できるようにするために、インフラをコード化しておきました。今回は、自分の学習も兼ねてAWS CDKを使いました。

  2. コアとなるMLflowとMLflow AI Gatewayは、それぞれがエンドポイントを建て、相互にアクセスします。いくつかのアプローチが考えられますが、今回は一つのコンテナにまとめ、それをAWS Fargate上で実行されるようにしました。

  3. MLflow AI GatewayにはドキュメントAPIが用意されています。そのため、MLflowとMLflow AI Gatewayの2つにアクセスできるようにロードバランサーを作成しました。MLflowのエンドポイントには5000番ポート、MLflow AI Gatewayには5100番ポートを使うように設定しました。

今後は、AWS環境だけでなく、そのほかのクラウド環境でもすぐに利用できるように引き続き改善していきたいと考えています。

まとめ

この記事では、継続的にLLMを開発するために必要になると思われるツールとしてMLflow AI Gatewayを紹介し、どのように役立つのかなどについて紹介をしました。 今回は、MLflow AI Gatewayに焦点をあてましたが、今後、他にも同様のツールが登場すると考えられます。それぞれのプロジェクトで構築しているシステム環境、利用しているツールを鑑みて、必要なツールとして考慮されてはいかがでしょうか？

この記事が気に入ったらサポートをしてみませんか？