All Articles

LLM (Cerebras-GPT) を MLflow Models でラップしてノーコードデプロイした話

はじめに

ChatGPT 以来の LLM ブームは 2023 年 4 月 5 日の現時点で留まるところを知らず、火付け役の OpenAI は GPT-4 をリリースして更に狂騒に拍車をかけ、 OpenAI に計算リソースと資金を提供して協業している弊社 Microsoft は各種製品に「Copilot」と呼ばれる自然言語補助 UI を搭載する方向で Azure OpenAI Service という OpenAI モデルのプロバイダーと OpenAI モデルのユーザーの 2 つの立場で爆走しており、私もキャッチアップするだけでひぃひぃ言っている状況です。

そんな LLM ブームの中にある 1 つのサブブームとして、大規模言語モデルを作成し OSS として公開する動きがあります。Meta が LLaMA を公開して以来、それをベースにした言語モデルが多数リリースされているようです。

また LLaMA とは別枠で独自に OSS なモデルを訓練する動きもあります。さながら Stable Diffusion が公開されて以来の画像生成モデルブームとも似たような事態が起きており、ここでも深層学習界隈あるあるだった画像が先行して言語が後追いするいつもの奴が起きてる気配を感じます。

というわけで、かつて rinna が日本語 GPT-2 を公開したときと比べて、更に数多の事前学習済みモデルを利用できる状況が整ってきています。シンプルにやりたいのであれば Azure OpenAI Service を使う方が簡単ですし性能も良さそうですが、fine-tuning の自由度や必要とされるスケールの事情、セキュリティ上の理由などにより自前でホストしたいケースが今後出てくる可能性もありそうですし、先んじて AzureML で OSS 言語モデルをホストする方法を検討しておこうと思います。

方針

使用する LLM

今回は大規模 Cerebras-GPT もその 1 つで、このモデルは Apache 2.0 ライセンスで公開されていることと、パラメーター数が異なるモデルが複数用意されていることでお試しにはちょうど良いということで、今回はこれを使ってみようと思います。

MLflow Models

今回は大規模言語モデルを Hugging Face から引っ張ってきて、MLflow Models でラップして手軽に利用できる API としてデプロイしてみようと思います。どうして MLflow Models でラップするかと言うと、そうすることで AzureML の Managed Online Endpoint にノーコードデプロイが可能になるからです。加えて Hugging Face が transformers ライブラリによって言語モデルの種類が変わってもほぼ同じコードを使いまわせるような抽象化レイヤーを提供してくれていますし、MLflow Models の抽象化と合わせれば言語モデルの API デプロイが非常に単純なタスクになりそうです。

MLflow Models については現在技術書典 14 で出す予定の書籍でも書いていて、その記述を以下に貼っておきます。

MLflow Models は機械学習モデルに対して統一的なフォーマットとインターフェースを提供する抽象化レイヤーです。

機械学習プロジェクトでは、モデルの学習や評価を行った後、そのモデルを実際の本番環境やテスト環境にデプロイすることが一般的です。しかし世の中には様々な機械学習フレームワークや環境が存在しているため、最終的にやりたいこと (デプロイ) は同じなのにモデルごとに異なる処理を記述する必要が生じてきます。

一例として、PyTorch と XGBoost で作成したモデルについて、それぞれモデルを保存し推論環境上で再ロードする処理を考えてみます。PyTorch モデルの場合はまずtorch.save(model.state_dict(), PATH)としてモデルのパラメーターを保存し、その後モデルを定義したクラスから作ったインスタンスであるmodelに対し、model.load_state_dict(torch.load(PATH))としてモデルのパラメーターをロードするという手順になります。必然的にモデルのパラメーターを保存したファイルだけでなく、モデルを定義したクラスやその周辺関数・クラス群もセットで推論環境に配置する必要があります。XGBoost の場合は学習済みモデルmodelに対し、model.save_model('<filename>.json')のようにしてモデルを保存し、model.load_model('<filename>.json')でモデルをロードします。このとき、XGBoost のライブラリさえあれば問題はなく、PyTorch と違ってクラス定義などを持ち込む必要はありません。

MLflow Models は、このようなフレームワーク間の取り回し方の違いや API の違いを吸収し、MLflow が提供する統一的なインターフェースで様々なフレームワークで作成したモデルを取り扱えるようにすることを可能にします。

手順

環境構築

以下 environment.yaml に基づいて、環境を構築します。

name: py310-chapter6-env
channels:
  - pytorch
  - nvidia
  - conda-forge
  - defaults
dependencies:
  - python=3.10
  - pytorch=2.0.0
  - torchvision=0.15.0
  - torchaudio=2.0.0
  - pytorch-cuda=11.7
  - pip
  - pip:
      - azure-ai-ml==1.5.0
      - mlflow==2.2.2
      - azureml-mlflow==1.49.0
      - ipykernel
      - numpy==1.23.5
      - pandas==1.5.3
      - xlrd==2.0.1
      - transformers==4.27.4
      - accelerate==0.18.0
conda env create -f environment.yaml
conda activate py310-chapter6-env
ipython kernel install --user --name=py310-chapter5-env

モデルテスト

まずはモデルを動かしてみます。後ほどモデル実体を AzureML にアップロードするために、ローカルにモデルを落としてそのモデルを読み込むように記述しています。

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
import mlflow
import pandas as pd

REPO_ID = "cerebras/Cerebras-GPT-111M"
download_path = snapshot_download(repo_id=REPO_ID)

tokenizer = AutoTokenizer.from_pretrained(download_path)
model = AutoModelForCausalLM.from_pretrained(download_path)

pipe = pipeline(
    task="text-generation",
    model=model,
    tokenizer=tokenizer,
    device=torch.device(type='cuda', index=0)
)

prompts = {"prompts": ["Generative AI is", "So, today we are"]}
input_df = pd.DataFrame(prompts)

generated_text = pipe(
    input_df["prompts"].values.tolist(),
    max_length=256,
    do_sample=False,
    no_repeat_ngram_size=2
)

outputs = [text[0]['generated_text'] for text in generated_text]
output_df = pd.DataFrame({"outputs": outputs})
print(output_df)

pyfunc フレーバーによるラップ

mlflow.pyfunc.PythonModelを継承した Class を実装します。この Class はまず predict 関数を実装する必要があり、オプションで load_context 関数を実装することでモデルの読み込みを自由に記述することができます。

signature = mlflow.models.signature.infer_signature(
    input_df,
    output_df
)

artifacts = {"cached_model_path": download_path}

class LLMWrapper(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
        import torch
        self.tokenizer = AutoTokenizer.from_pretrained(
            context.artifacts["cached_model_path"]
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            context.artifacts["cached_model_path"]
        )
        self.pipe = pipeline(
            task="text-generation",
            model=model,
            tokenizer=tokenizer,
            device=torch.device(type='cuda', index=0)
        )
    def predict(self, context, model_input):
        generated_text = self.pipe(
            model_input["prompts"].values.tolist(),
            max_length=256,
            do_sample=False,
            no_repeat_ngram_size=2
        )
        outputs = [text[0]['generated_text'] for text in generated_text]
        output_df = pd.DataFrame({"outputs": outputs})
        return outputs

load_contextはモデルをロードする際に呼び出されるコールバックです。LLMWrapperクラスのインスタンスがモデル実体として振る舞うわけですが、このインスタンスは内部的には cloudpickle によってバイナリ化されて保存されます。これをインスタンスに戻した後、load_contextが勝手に呼ばれ、このときにモデルが復元する処理が走ることが期待されます。load_contextの引数であるcontextに渡されるのは上で定義しているartifacts辞書と同じ key-value 構成を持った辞書です。contextとartifactsはこの時点では独立していますが、後ほど MLflow でモデルを記録する際に紐づける記述が登場します。

モデル登録

AzureML に接続します。

from azure.ai.ml import MLClient
from azure.identity import DefaultAzureCredential

subscription_id = "SUBSCRIPTION_ID"
resource_group = "RESOURCE_GROUP"
workspace = "AML_WORKSPACE_NAME"

ml_client = MLClient(
    DefaultAzureCredential(),
    subscription_id,
    resource_group,
    workspace,
)

azureml_mlflow_uri = ml_client.workspaces.get(
    ml_client.workspace_name
).mlflow_tracking_uri

mlflow.set_tracking_uri(azureml_mlflow_uri)

exp = mlflow.set_experiment("chapter6-llm-notebook")

続いてモデルの登録を行います。

with mlflow.start_run() as run:
    mlflow_model_dir = 'llm_model'
    mlflow.pyfunc.log_model(
        artifact_path=mlflow_model_dir,
        python_model=LLMWrapper(),
        conda_env='environment.yaml',
        artifacts=artifacts,
        signature=signature,
    )

ここでpython_modelに渡しているのが先ほど実装したラッパ－クラスのインスタンスで、artifactsに渡しているのがファイルパスを記した辞書オブジェクトです。この指定により、load_contextの引数であるcontextにartifacts同様の key-value 構造を持った辞書が渡されるようになります。

なお、artifactsの key は何でも良いのですが、value の方は何らかのファイルかディレクトリを指定する必要があります。value で指定したファイル/ディレクトリは MLflow Models の実体ファイル群の 1 つとしてアップロードされます。value の値はアップロード後のパスに置換され、key は同じままその value はアップロードされたファイル/ディレクトリとなった辞書オブジェクトが実際にload_contextに渡されるcontextとなります。

登録後、モデルのテストをします。

loaded_model = mlflow.pyfunc.load_model(
    model_uri=f"runs:/{run.info.run_id}/{mlflow_model_dir}/",
)

sample_prompts = {"prompts": ["Generative AI is", "So, today we are"]}
sample_input_df = pd.DataFrame(prompts)

print(loaded_model.predict(sample_input_df))

何か結果が返ってこれば成功です。仕上げにモデルを AzureML Model Registry に登録します。

mlflow.register_model(
    model_uri=f"runs:/{run.info.run_id}/{mlflow_model_dir}/",
    name='chapter6-cerebras-gpt'
)

ノーコードデプロイ

登録したモデルをデプロイします。

GUI からモデルを選択し、デプロイから「リアルタイム エンドポイント」を選択します。MLflow Models でラップした恩恵で、何らスクリプトを書くことなく API デプロイを実行することができます。NVIDIA GPU を積んだマシンを選ぶ必要がある点にのみ注意する必要があります。

image

10 ～ 20 分程度待つとデプロイが完了します。 (たまに conda 周りで通信エラー起こしてこけます、そのときはやり直してください)

モデルのテストをしてみます。111M パラメーターの限界か日本語では酷い出力ですが、ひとまず何か出ているのでヨシとします。

image

おわりに

OSS LLM の 1 つである Cerebras-GPT モデルを Hugging Face の transformers ライブラリを使用してローカルに落として読み込んで、MLflow Models でラップすることでノーコードデプロイまで持っていくことができました。

大抵の言語モデルは今回のスクリプトで使用している AutoTokenizer と AutoModelForCausalLM とそれらをさらにラップした pipeline によってロードから出力までを行うことができますし、モデル固有の要素はほぼ排除して実装したので、REPO_ID さえ変更すれば他の言語モデルにも適用できるものと思います。

これで自分だけの LLM API を AzureML を使用して簡単にホストできます。LLM の fine-tuning などもそのうちやってみようと思いますが、課金がしんどいので今度にします。

Published 2023/04/05

ShuntaIto による技術ブログ