🗿
"LLM-jp-13B" をDatabricks上でサービングする
"LLM-jp-13B" をDatabricks上でサービングする
LLM勉強会（LLM-jp）が開発しているLLM「LLM-jp-13B」を多くの方に活用いただきたい！ということでサービングしてみましょう。
Databricks上のマネージドMLFlowを使うと、LLMを比較的簡単にMLOpsパイプラインに乗せることができ、サービングエンドポイントとしてサーバーレスGPU環境上にデプロイ可能です。
LLM勉強会およびLLM-jp-13Bについてはこちらをご参照ください。
MLFlowについてはこちらの４回ものが分かりやすいです。
環境
- ノートブック用
- Databricks Runtime: 14.1 ML GPU
- ノードタイプ: g5.12xlarge (A10g x 4)
- サービング用
- サーバーレス: GPU Medium x4 (4xA10G), Small 4 concurrency (112 DBU)
大まかな流れ
- MLFlowのTransformersフレーバーを使用してモデルをMLFlow Trackingに記録
- Trackingに記録したモデルをモデルレジストリー（今回はUnity Catalog）へ登録
- Unity Cactalogに登録したモデルをサービングエンドポイントとしてデプロイ
前提
- Unity Catalogが有効化されている。
- まだの方はこちら を参照して有効化ください
1. MLFlow Transformersフレーバーを使用してMLFlow Trackingにモデルを記録
まずはHuggingFaceからモデルをダウンロードします。
今回はDollyのデータセットも使われているllm-jp/llm-jp-13b-instruct-full-dolly-oasst-v1.0を使用します。
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_name = "llm-jp/llm-jp-13b-instruct-full-dolly-oasst-v1.0"
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16, cache_dir="/local_disk0/.cache/huggingface/")
tokenizer = AutoTokenizer.from_pretrained(model_name)
続いて、MLFlowが提供するTransformersフレーバーを使用して、モデルをMLFlow Trackingへ記録します。
import mlflow
from mlflow.models import infer_signature
from mlflow.models.signature import ModelSignature
from mlflow.types import DataType, Schema, ColSpec
import pandas as pd
# params を含むモデルシグネチャ（スキーマ）を定義
text = "自然言語処理とは何か"
text = text + "### 回答："
input_example = {"prompt": text}
inference_config = {
"temperature": 1.0,
"max_new_tokens": 100,
"do_sample": True,
"top_p": 0.95,
}
signature = infer_signature(
model_input=input_example,
model_output="Machien Learning is...",
params=inference_config
)
# アーティファクト、パイプ要件、入力例など、モデルの詳細をログに記録する。
# 完了するまでに数分かかる可能性があります。
with mlflow.start_run() as run:
mlflow.transformers.log_model(
transformers_model={
"model": model,
"tokenizer": tokenizer,
},
task = "text-generation",
artifact_path="model",
pip_requirements=["torch", "transformers", "accelerate", "sentencepiece"],
input_example=input_example,
signature=signature,
# メタデータ・タスクを追加して、後で作成されるモデル・サービング・エンドポイントが最適化されるようにする。
metadata={"task": "llm/v1/completions"}
)
2. Trackingに記録したモデルをUnity Catalogへ登録
続いて、モデルをモデルレジストリーへ移していきます。
現在Databricks上ではMLFlow Model RegistryとUnity Catalogのどちらかをモデルのレジストリーとして選択いただけます。Unity Catalogのモデルを使用すると、ワークスペース間での一元的なアクセス制御、監査、リネージ、モデルディスカバリーなど、Unity Catalogの利点をMLモデルに拡張できるため、基本的にはUnity Catalogの活用をお勧めしています。したがって、今回はUnity Catalogへモデルを登録します。
import mlflow
# モデルレジストリーとしてUnity Catalogをセット
mlflow.set_registry_uri("databricks-uc")
# UCモデル名は<カタログ名>.<スキーマ名>.<モデル名>のパターンに従っており、カタログ名、スキーマ名、登録モデル名に対応していることに注意してください
registered_name = <カタログ名>.<スキーマ名>.<モデル名>
result = mlflow.register_model(
"runs:/"+run.info.run_id+"/model",
registered_name,
)
3. Unity Cactalogに登録したモデルをサービングエンドポイントにデプロイ
モデルをエンドポイントとしてデプロイします。
この際にDatabricksのサービスへアクセスするためのトークンが必要です。
import requests
import json
endpoint_name = $YOUR_ENDPOINT_NAME
databricks_token = $YOUR_DATABRICKS_TOKEN
databricks_url = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiUrl().getOrElse(None)
deploy_headers = {'Authorization': f'Bearer {databricks_token}', 'Content-Type': 'application/json'}
deploy_url = f'{databricks_url}/api/2.0/serving-endpoints'
model_version = result
endpoint_config = {
"name": endpoint_name,
"config": {
"served_models": [{
"name": f'{model_version.name.replace(".", "_")}_{model_version.version}',
"model_name": model_version.name,
"model_version": model_version.version,
"workload_type": "MULTIGPU_MEDIUM",
"workload_size": "Small",
"scale_to_zero_enabled": "False"
}]
}
}
endpoint_json = json.dumps(endpoint_config, indent=' ')
# APIにPOSTリクエストを送る
deploy_response = requests.request(method='POST', headers=deploy_headers, url=deploy_url, data=endpoint_json)
if deploy_response.status_code != 200:
raise Exception(f'Request failed with status {deploy_response.status_code}, {deploy_response.text}')
# POSTリクエストのレスポンスを表示する
# 最初にサービングエンドポイントを作成する際、'ready'の状態が 'NOT_READY'であることが表示される
# Databricksモデルのサービングエンドポイントページでステータスを確認可能
print(deploy_response.json())
（オプション） RESTクライアントからEndpointをCallする
import os
import requests
import numpy as np
import pandas as pd
import json
def score_model(workspace_name: str, endpoint_name: str, databricks_token: str):
url = f'https://{workspace_name}.cloud.databricks.com/serving-endpoints/{endpoint_name}/invocations'
headers = {'Authorization': f'Bearer {databricks_token}', 'Content-Type': 'application/json'}
ds_dict = {
"dataframe_split": {
"columns": [
"prompt"
],
"data": [
[
"自然言語処理とは何か### 回答："
]
]
}
}
data_json = json.dumps(ds_dict, allow_nan=True)
response = requests.request(method='POST', headers=headers, url=url, data=data_json)
if response.status_code != 200:
raise Exception(f'Request failed with status {response.status_code}, {response.text}')
return response.json()
print(score_model(YOUR_WORKSPACE_NAME, YOUR_ENDPOINT_NAME, YOUR_DATABRICKS_TOKEN))
以下が出力。
{
"predictions": [
"自然言語処理とは何か### 回答：自然言語処理（NLP）とは、人間が知覚できる言語を扱うコンピューター技術である。テキストや音声入力の認識、要約、検索、要約、テキストやデータの要約などを行う。NLPの例には、チャットアシスタント、音声認識、会話エージェント、ビデオチャットの録画などがある。"
]
}
まとめ
Databricks上のマネージドMLFlowを使用して、LLM-jp-13bのサービングに成功しました。ここまでくるとRESTでモデルへプロンプトを投げられるようになるので、様々なアプリとの統合も現実的になってまいります。
現状、推論インフラはサーバーレスではありますが、使用されているGPUはA10g×4です。今後モデルの最適化および量子化などにより、GPUの枚数を減らすことができる、または、CPUでも推論できるものか、継続して検証していこうと思います。
BFN!
参考
- NIIによるニュースリリース
- HUggingFace 「llm-jp」
Databricks無料トライアル
Discussion
大変参考になりました、ありがとうございます！
loraチューニングされたモデルの登録は今の所サポートされていないようなのですが、
このissueの解決するのを待つ以外にloraのモデルを登録する方法はありますでしょうか？
フィードバックありがとうございます！
ご質問に関してですが、こちらのサンプルが参考になるかもしれません。 （このサンプルはこのままだと見づらいですが、Databricksへインポートいただくと、ノートブックへ変換されます。。。）
LLaMa v2をQLoraでFine Tuningした後に、MLFlowに登録しています。
上記Issueが解決された結果できるようになったのか否かは定かではないですが、いずれにしてもMLFLow登録時にモデルを「mlflow.pyfunc.PythonModel」でラップしてあげるとモデル本体とLoraのアダプター含めてMLFlowへ登録できるように見えます。
お役に立てば幸いです。