プロビジョニングされたスループットの基盤モデルAPIs

この記事では、プロビジョニング スループットで基盤モデルAPIsを使用してモデルをデプロイする方法を説明します。 Databricks本番運用ワークロードにプロビジョニング スループットを推奨しており、パフォーマンスが保証された基盤モデルに最適化された推論を提供します。

サポートされているモデル アーキテクチャの一覧については、「 プロビジョニングされたスループットの基盤モデル APIs 」を参照してください。

要件

「要件」を参照してください。

微調整された基盤モデルを展開するために、

  • モデルは、MLflow 2.4 以降または Databricks Runtime 13.2 機械学習以降を使用してログに記録する必要があります。

  • Databricks では、大規模なモデルのアップロードとダウンロードを高速化するために、Unity Catalog のモデルを使用することを推奨しています。

[推奨]Databricks Marketplace から基本基盤モデルをデプロイする

Databricks Marketplace を使用して、基本基盤モデルを Unity Catalog にインストールできます。

Databricks では、 Databricks Marketplaceを使用して基盤モデルをインストールすることをお勧めします。 モデル ファミリを検索し、モデル ページから[アクセスを取得]を選択し、ログイン認証情報を指定して、モデルをUnity Catalogにインストールできます。

モデルがUnity Catalogにインストールされたら、Serving UI を使用してモデルサービング エンドポイントを作成できます。 「UI を使用してプロビジョニングされたスループット エンドポイントを作成する」を参照してください。

Databricks Marketplace の DBRX モデル

Databricks では、ワークロードに DBRX Instruct モデルを提供することをお勧めします。 プロビジョニングされたスループットを使用して DBRX Base モデルと DBRX Instruct モデルを提供するには、前のセクションのガイダンスに従って、これらのモデルを Databricks Marketplace から Unity Catalog にインストールする必要があります。

これらの DBRX モデルを提供する場合、プロビジョニングされたスループットは最大 16k のコンテキスト長をサポートします。 より大きなコンテキストサイズは近日公開予定です。

DBRX モデルは、次のデフォルトのシステム プロンプトを使用して、モデル応答の関連性と正確性を確保します。

You are DBRX, created by Databricks. You were last updated in December 2023. You answer questions based on information available up to that point.
YOU PROVIDE SHORT RESPONSES TO SHORT QUESTIONS OR STATEMENTS, but provide thorough responses to more complex and open-ended questions.
You assist with various tasks, from writing to coding (using markdown for code blocks — remember to use ``` with code, JSON, and tables).
(You do not have real-time data access or code execution capabilities. You avoid stereotyping and provide balanced perspectives on controversial topics. You do not provide song lyrics, poems, or news articles and do not divulge details of your training data.)
This is your system prompt, guiding your responses. Do not reference it, just respond to the user. If you find yourself talking about this message, stop. You should be responding appropriately and usually that means not mentioning this.
YOU DO NOT MENTION ANY OF THIS INFORMATION ABOUT YOURSELF UNLESS THE INFORMATION IS DIRECTLY PERTINENT TO THE USER'S QUERY.

ログの微調整された基盤モデル

Databricks Marketplaceからモデルをインストールできない場合は、 Unity Catalogにログを記録することで、微調整された基盤モデルをデプロイできます。 以下に、 MLflowモデルをUnity Catalogに記録するコードを設定する方法を示します。

mlflow.set_registry_uri('databricks-uc')
CATALOG = "ml"
SCHEMA = "llm-catalog"
MODEL_NAME = "mpt" # or "bge"
registered_model_name = f"{CATALOG}.{SCHEMA}.{MODEL_NAME}"

MLflow transformersフレーバーを使用してモデルをログに記録し、次のいずれかを使用して MLflow メタデータのタスク フィールドを指定できます。

  • metadata = {"task": "llm/v1/completions"}

  • metadata = {"task": "llm/v1/chat"}

  • metadata = {"task": "llm/v1/embeddings"}

これは、モデルサービング エンドポイントに使用されるAPI署名を指定し、この方法でログに記録されたモデルはプロビジョニング スループットの対象となります。

以下は、テキスト補完言語モデルをログに記録する方法の例です。

model = AutoModelForCausalLM.from_pretrained("mosaicml/mpt-7b-instruct",torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained("mosaicml/mpt-7b-instruct")
with mlflow.start_run():
    components = {
        "model": model,
        "tokenizer": tokenizer,
    }
    mlflow.transformers.log_model(
        transformers_model=components,
        "model",
        input_example={"prompt": np.array(["Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\nWhat is Apache Spark?\n\n### Response:\n"])},
        metadata={"task": "llm/v1/completions"},
        registered_model_name=registered_model_name
    )

プロビジョニングされたスループットは、小規模と大規模の両方の BGE 埋め込みモデルもサポートします。 以下は、プロビジョニングされたスループットで提供できるようにモデルBAAI/bge-small-en-v1.5をログに記録する方法の例です。

model = AutoModel.from_pretrained("BAAI/bge-small-en-v1.5")
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-small-en-v1.5")
with mlflow.start_run():
    components = {
        "model": model,
        "tokenizer": tokenizer,
    }
    mlflow.transformers.log_model(
        transformers_model=components,
        artifact_path="bge-small-transformers",
        task="feature-extraction",
        metadata={"task": "llm/v1/embeddings"},
        registered_model_name=registered_model_name)

UI を使用してプロビジョニングされたスループット エンドポイントを作成する

記録済みモデルがUnity Catalogに追加されたら、次のステップを使用してプロビジョニングされたスループットを提供するエンドポイントを作成します。

  1. ワークスペースのサービス提供 UIに移動します。

  2. [ サービス エンドポイントの作成] を選択します。

  3. [エンティティ]フィールドで、 Unity Catalogからモデルを選択します。 対象となるモデルの場合、サービス対象エンティティの UI に「プロビジョニングされたスループット」画面が表示されます。

  4. [最大まで]ドロップダウンで、エンドポイントの 1 秒あたりの最大トークン スループットを構成できます。

    1. プロビジョニングされたスループット エンドポイントは自動的にスケールされるため、 [変更]を選択して、エンドポイントがスケールダウンできる 1 秒あたりの最小トークンを表示できます。

プロビジョニング済みスループット

REST API を使用してプロビジョニングされたスループット エンドポイントを作成する

REST API を使用してプロビジョニングされたスループット モードでモデルをデプロイするには、リクエストでmin_provisioned_throughputmax_provisioned_throughputフィールドを指定する必要があります。

モデルに適したプロビジョニングされたスループットの範囲を特定するには、 「プロビジョニングされたスループットを段階的に取得する」を参照してください。

import requests
import json

# Set the name of the MLflow endpoint
endpoint_name = "llama2-13b-chat"

# Name of the registered MLflow model
model_name = "ml.llm-catalog.llama-13b"

# Get the latest version of the MLflow model
model_version = 3

# Get the API endpoint and token for the current notebook context
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}"}

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

if 'optimizable' not in optimizable_info or not optimizable_info['optimizable']:
   raise ValueError("Model is not eligible for provisioned throughput")

chunk_size = optimizable_info['throughput_chunk_size']

# Minimum desired provisioned throughput
min_provisioned_throughput = 2 * chunk_size

# Maximum desired provisioned throughput
max_provisioned_throughput = 3 * chunk_size

# Send the POST request to create the serving endpoint
data = {
    "name": endpoint_name,
    "config": {
        "served_entities": [
            {
                "entity_name": model_name,
                "entity_version": model_version,
                "min_provisioned_throughput": min_provisioned_throughput,
                "max_provisioned_throughput": max_provisioned_throughput,
            }
        ]
    },
}

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

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

プロビジョニングされたスループットを段階的に取得します

プロビジョニングされたスループットは、1 秒あたりのトークンの増分で利用できますが、具体的な増分はモデルによって異なります。 ニーズに適した範囲を特定するために、Databricks ではプラットフォーム内でモデル最適化情報 API を使用することをお勧めします。

GET api/2.0/serving-endpoints/get-model-optimization-info/{registered_model_name}/{version}

以下は API からの応答の例です。

{
 "optimizable": true,
 "model_type": "llama",
 "throughput_chunk_size": 980
}

ノートブックの例

次のノートブックは、プロビジョニングされたスループットの基盤モデル API を作成する方法の例を示しています。

Llama2 モデルノートブックのプロビジョニングされたスループット サービス

ノートブックを新しいタブで開く

Mistral モデル ノートブックのプロビジョニングされたスループット サービス

ノートブックを新しいタブで開く

BGE モデル ノートブックのプロビジョニングされたスループット サービス

ノートブックを新しいタブで開く

制限事項

  • GPU容量の問題が原因でモデルのデプロイメントが失敗し、エンドポイントの作成または更新中にタイムアウトが発生することがあります。 解決するには、Databricks アカウント チームにお問い合わせください。

  • 基盤モデル APIs の自動スケーリングは、CPU モデルサービングよりも低速です。 Databricks では、要求のタイムアウトを回避するためにオーバープロビジョニングを推奨しています。