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

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

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

要件

要件を参照してください。微調整された基盤モデルをデプロイするには、 「微調整された基盤モデルをデプロイする」を参照してください。

[推奨] Unity Catalogから基盤モデルをデプロイする

プレビュー

この機能は パブリックプレビュー版です。

Databricks では、Unity Catalog にプリインストールされている基盤モデルを使用することをお勧めします。 これらのモデルは、スキーマai (system.ai) のカタログsystemの下にあります。

基盤モデルをデプロイするには:

  1. 「カタログエクスプローラ」(Catalog Explorer) で「 system.ai 」() に移動します。

  2. デプロイするモデルの名前をクリックします。

  3. モデル ページで、[ このモデルを提供 ] ボタンをクリックします。

  4. [ サービス エンドポイントの作成 ] ページが表示されます。 「UI を使用してプロビジョニングされたスループット エンドポイントを作成する」を参照してください。

Databricks Marketplaceから基盤モデルをデプロイする

あるいは、 Databricks Marketplaceから Unity Catalog に基盤モデルをインストールすることもできます。

モデル ファミリを検索し、モデル ページから[アクセスを取得]を選択してログイン資格情報を入力し、モデルを Unity Catalog にインストールできます。

モデルをUnity Catalogにインストールしたら、Serving UI を使用してモデルサービングエンドポイントを作成できます。

DBRX モデルのデプロイ

Databricks では、ワークロードに DBRX Instruct モデルを提供することをお勧めします。 プロビジョニングされたスループットを使用して DBRX Instruct モデルを提供するには、 「[推奨] 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.

きめ細かく調整された基盤モデルをデプロイする

system.aiスキーマ内のモデルを使用できない場合、または Databricks Marketplace からモデルをインストールできない場合は、Unity Catalog にログ記録することで、微調整された基盤モデルをデプロイできます。 このセクションと次のセクションでは、コードを設定して MLflow モデルを Unity Catalog に記録し、UI または REST API を使用してプロビジョニングされたスループット エンドポイントを作成する方法を示します。

要件

  • 微調整された基盤モデルのデプロイは、MLflow 2.11 以降でのみサポートされます。 Databricks Runtime 15.0 ML 以降では、互換性のある MLflow バージョンがプリインストールされます。

  • 埋め込みエンドポイントの場合、モデルは小規模または大規模の BGE 埋め込みモデル アーキテクチャである必要があります。

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

カタログ、スキーマ、モデル名の定義

微調整された基盤モデルをデプロイするには、ターゲットの Unity Catalog カタログ、スキーマ、および選択したモデル名を定義します。

mlflow.set_registry_uri('databricks-uc')
CATALOG = "catalog"
SCHEMA = "schema"
MODEL_NAME = "model_name"
registered_model_name = f"{CATALOG}.{SCHEMA}.{MODEL_NAME}"

モデルのログ記録

モデル エンドポイントのプロビジョニングされたスループットを有効にするには、MLflow transformersフレーバーを使用してモデルをログに記録し、次のオプションから適切なモデル タイプ インターフェースを使用してtask引数を指定する必要があります。

  • "llm/v1/completions"

  • "llm/v1/chat"

  • "llm/v1/embeddings"

これらの引数は、モデルサービング エンドポイントに使用されるAPI署名を指定します。 これらのタスクと対応する入力/出力スキーマの詳細については、MLflow のドキュメントを参照してください。

以下は、MLflow を使用してログに記録されたテキスト補完言語モデルをログに記録する方法の例です。

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,
        artifact_path="model",
        # Specify the llm/v1/xxx task that is compatible with the model being logged
        task="llm/v1/completions",
        # Specify an input example that conforms to the input schema for the task.
        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"])},
        # By passing the model name, MLflow automatically registers the Transformers model to Unity Catalog with the given catalog/schema/model_name.
        registered_model_name=registered_model_name
        # Optionally, you can set save_pretrained to False to avoid unnecessary copy of model weight and gain more efficiency
        save_pretrained=False
    )

注:

2.12 より前のバージョンのMLflowを使用している場合は、代わりに同じ mlflow.transformer.log_model() 関数の metadata 構造体内でタスクを指定する必要があります。

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

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

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

プロビジョニングされたスループットは、小規模および大規模の 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="model",
        task="llm/v1/embeddings",
        registered_model_name=registered_model_name,
        # model_type is required for logging a fine-tuned BGE models.
        metadata={
            "model_type": "bge-large"  # Or "bge-small"
        }
    )

モデルがUnity Catalogにログインしたら、 UI を使用して「 タイムアウト エンドポイントの作成」に進み、タイムアウトを持つモデルサービングエンドポイントを作成します。

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 では、要求のタイムアウトを回避するためにオーバープロビジョニングを推奨しています。