More than 1 year has passed since last update.

DatabricksのAI Playground上でSwallow-MS-7b-instruct-v0.1を試す

Posted at 2024-04-30

導入

以下で紹介されているDatabricksのAI Playground機能が、自分の環境下でも利用できるようになっていました。

現状、日本リージョンだと自分で用意した「プロビジョニング済みスループットのエンドポイント」か「外部モデルのサービング」のみが利用可能なようなので、自前でプロビジョニングされたスループットの基盤モデルを用意して試してみます。

せっかくなので、先日ひっそり（？）とリリースされていた、TokyoTech-LLMのSwallow-MS-7b-instruct-v0.1を使ってます。

Step1. 準備

Databricks上でノートブックを作成し、以下の処理を実行することで「プロビジョニングされたスループットの基盤モデル」を準備します。
(DBRは14.3ML、インスタンスタイプはg5.xlarge(AWS)で実行しました)

パッケージインストール

transformersとmlflowを最新化します。
MLflowは2.12.1以降をインストールしてください。

%pip install -U "mlflow>=2.12.1" transformers accelerate

dbutils.library.restartPython()

MLflow登録

こちらのモデルを事前にダウンロードしておき、transformersで読み込みます。

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_path = "/Volumes/training/llm/model_snapshots/models--tokyotech-llm--Swallow-MS-7b-v0.1/"

model = AutoModelForCausalLM.from_pretrained(
    model_path, 
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(
    model_path,
)

MLflow Model Registryにロギングします。
注意点として、MLflowのバージョンが2.12.1以降とその前で指定の仕方が異なります。
今回は2.12.1以降のやり方です。

詳しくは公式Docを確認ください。

import mlflow
import numpy as np

# Models in UCを使っていない場合には下の行をコメントアウトします 
# そして、3レベル名前空間の名称ではなくシンプルにモデル名を指定します
mlflow.set_registry_uri('databricks-uc')
CATALOG = "training"
SCHEMA = "llm"
registered_model_name = f"{CATALOG}.{SCHEMA}.swallow-ms-7b-instruct"

# サンプル入力の定義
input_example = {
  "messages": [
    {
      "role": "system",
      "content": "You are a helpful assistant."
    },
    {
      "role": "user",
      "content": "Hello!"
    }
  ],
  "max_tokens": 10,
  "stream": false
}

# 新規MLflowランをスタート
with mlflow.start_run():
    components = {
        "model": model,
        "tokenizer": tokenizer,
    }
    mlflow.transformers.log_model(
        transformers_model=components,
        task="llm/v1/chat",
        artifact_path="model",
        registered_model_name=registered_model_name,
        input_example=input_example,
        example_no_conversion=True,
    )

モデルサービングへの登録

まず、登録したモデルが「プロビジョニングされたスループットの基盤モデル」対応かどうかを確認。

import requests
import json
from mlflow import MlflowClient


# 登録されたMLflowモデルの名前
CATALOG = "training"
SCHEMA = "llm"
registered_model_name = f"{CATALOG}.{SCHEMA}.swallow-ms-7b-instruct"

model_name = registered_model_name

# MLflowモデルの最新バージョンを取得
client=MlflowClient()
versions = [mv.version for mv in client.search_model_versions(f"name='{registered_model_name}'")]
model_version = versions[0]

# 現在のノートブックのコンテキストからAPIエンドポイントとトークンを取得
API_ROOT = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiUrl().get() 
API_TOKEN = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

headers = {"Context-Type": "text/json", "Authorization": f"Bearer {API_TOKEN}"}

response = requests.get(url=f"{API_ROOT}/api/2.0/serving-endpoints/get-model-optimization-info/{model_name}/{model_version}", headers=headers)

print(json.dumps(response.json(), indent=4))

出力

{
    "optimizable": true,
    "model_type": "mistral",
    "throughput_chunk_size": 970,
    "dbus": 24
}

optimizableがTrueのため、対象です。

では、こちらのモデルを使ってモデルサービングのエンドポイントを登録します。

# MLflowエンドポイント名を設定
endpoint_name = "swallow-ms-7b-instruct-endpoint"

# プロビジョンされるスループットの最小値を指定 
min_provisioned_throughput = response.json()['throughput_chunk_size']

# プロビジョンされるスループットの最大値を指定
max_provisioned_throughput = response.json()['throughput_chunk_size']


data = {
    "name": endpoint_name,
    "config": {
        "served_entities": [
            {
                "entity_name": model_name,
                "entity_version": model_version,
                "min_provisioned_throughput": min_provisioned_throughput,
                "max_provisioned_throughput": min_provisioned_throughput,
            }
        ]
    },
}

headers = {"Context-Type": "text/json", "Authorization": f"Bearer {API_TOKEN}"}

response = requests.post(url=f"{API_ROOT}/api/2.0/serving-endpoints", json=data, headers=headers)

print(json.dumps(response.json(), indent=4))