大規模言語モデルに最適化されたDatabricksモデルサービング

[Model serving optimized for large language models \| Databricks on AWS](https://docs.databricks.com/en/machine-learning/model-serving/llm-optimized-model-serving.html) [2023/9/28時点]の翻訳です。

:::note warn
本書は抄訳であり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
:::

:::note info
**プレビュー**
本機能は[パブリックプレビュー](https://docs.databricks.com/release-notes/release-types.html)です。
:::

本書では、GPUを用いた大規模言語モデル(LLM)に最適化されたモデルサービングの有効化方法を説明します。

LLM最適化モデルサービングは、従来のサービングアプローチと比較して3-5倍改善されたスループットとレーテンシーを提供します。現時点では、以下のLLMファミリーおよびプロプライエタリなデータでファインチューニングされた派生モデルをサポートしています:

- MPTファミリーのモデル
- LLaMA-V2モデル

# 要件

- LLM最適化モデルサービングは、[GPUデプロイメント](https://docs.databricks.com/en/machine-learning/model-serving/create-manage-serving-endpoints.html#gpu)のパブリックプレビューの一部としてサポートされています。
- お使いのモデルは、MLflow 2.4以降、Databricks機械学習ランタイム 13.2以降で記録される必要があります。
- 大規模なモデルのアップロードとダウンロードを高速にするために、Unity Catalogでモデルを管理することをお勧めします。
- モデルをデプロイする際には、あなたのモデルのパラメーターのサイズに適した計算資源のサイズを選択することが重要です。推奨値に関しては以下の表をご覧ください。500億以上のパラメータのモデルに関しては、必要なGPUにアクセスするためにDatabricks担当者に連絡いただくか、[プレビュー登録フォーム](https://docs.google.com/forms/d/1-GWIlfjlIaclqDz6BPODI2j1Xg4f4WbFvBXyebBpN-Y/edit?ts=65124ee0)を提出してください。
![Screenshot 2023-10-04 at 9.05.05.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/df54afcf-20ad-8746-d3ce-72ad321798b8.png)

# 大規模言語モデルの記録

最初に、あなたのモデルをMLflowの`transformes`フレーバーで記録し、MLflowのメタデータで`metadata = {"task": "llm/v1/completions"}`をタスクフィールドに指定します。ここでは、モデルサービングエンドポイントで使用されるAPIのシグネチャを指定します。

DatabricksのLLM最適化モデルサービングは、Databricks AIゲートウェイでサポートされているルートタイプと互換性があります。現時点では、`llm/v1/completions`です。サポートされていないモデルファミリーあるいはタスクタイプがある場合には、Databricks担当者にお問い合わせください。

```py:Python
model = AutoModelForCausalLM.from_pretrained("mosaicml/mpt-7b-instruct",device_map="auto",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(
        "model",
        transformers_model=MPT7BInstruct(),
        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='mpt'
  )
```

モデルが記録されたら、以下のように`CATALOG.SCHEMA.MODEL_NAME`をモデルの3レベルの名前で置き換えて、Unity Catalogにモデルを登録します。

```py:Python
mlflow.set_registry_uri("databricks-uc")

registered_model_name=CATALOG.SCHEMA.MODEL_NAME
```

# モデルサービングエンドポイントの作成

次に、モデルサービングエンドポイントを作成します。あなたのモデルがLLM最適化フレームワークでサポートされているのであれば、サーブしようとした際にDatabricksが自動で最適化モデルサービングエンドポイントを作成します。

```py:Python
import requests
import json

# Set the name of the MLflow endpoint
endpoint_name = "llama2-3b-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

# Specify the type of compute (CPU, GPU_SMALL, GPU_MEDIUM, etc.)
workload_type = "GPU_MEDIUM_4"

# Specify the scale-out size of compute (Small, Medium, Large, etc.)
workload_size = "Small"

# Specify Scale to Zero (only supported for CPU endpoints)
scale_to_zero = False

# 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()

# send the POST request to create the serving endpoint

data = {
    "name": endpoint_name,
    "config": {
        "served_models": [
            {
                "model_name": model_name,
                "model_version": model_version,
                "workload_size": workload_size,
                "scale_to_zero_enabled": scale_to_zero,
                "workload_type": workload_type,
            }
        ]
    },
}

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))
```

# 入出力スキーマのフォーマット

最適化されたLLMサービングエンドポイントには、Databricksがコントロールする入出力スキーマがあります。4つのフォーマットがサポートされています。

- `dataframe_split`は、`split`オリエンテーションのJSONでシリアライズされたPandasデータフレームです。

```json:JSON
{
  "dataframe_split":{
    "columns":["prompt"],
    "index":[0],
    "data":[["Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instructions:\nWhat is Apache Spark?\n\n### Response:\n"]]
  },
  "params": {
    "temperature": 0.5,
    "max_tokens": 100,
    "stop": ["word1","word2"],
    "candidate_count": 1
  }
}
```

- `dataframe_records`は、`records`オリエンテーションのJSONでシリアライズされたPandasデータフレームです。

```json:JSON
{
  "dataframe_records": [{"prompt": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instructions:\nWhat is Apache Spark?\n\n### Response:\n"}],
  "params": {
    "temperature": 0.5,
    "max_tokens": 100,
    "stop": ["word1","word2"],
    "candidate_count": 1
  }
}
```

- `instances`

```json:JSON
{
  "instances": [
   {
     "prompt": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instructions:\nWhat is Apache Spark?\n\n### Response:\n"
   }
  ],
  "params": {
  "temperature": 0.5,
  "max_tokens": 100,
  "stop": ["word1","word2"],
  "candidate_count": 1
  }
}
```

- `inputs`

```json:JSON
{
  "inputs": {
    "prompt": ["Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instructions:\nWhat is Apache Spark?\n\n### Response:\n"],
  },
  "params": {
    "temperature": 0.5,
    "max_tokens": 100,
    "stop": ["word1","word2"],
    "candidate_count": 1
  }
}
```

# エンドポイントへのクエリー

エンドポイントの準備ができると、APIリクエストを用いてクエリーを行うことができます。モデルのサイズと複雑性によりますが、エンドポイントの準備ができるまで30分以上を要します。

```py:Python

data = {
  "inputs" : ["Hello, I'm a language model,"],
  "params" : {"max_new_tokens": 100, "temperature": 1}
}

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

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

print(json.dumps(response.json()))
```

# 制限

- GPUでサービングされるモデルのインストール要件が多いため、GPUサービングのコンテナイメージ作成は、CPUサービングのイメージ作成よりも時間を要します。
    - また、モデルのサイズがイメージ作成にインパクトを与えます。例えば、300億以上のパラメーターを持つモデルは構築に少なくとも1時間を要します。
    - 同じバージョンのモデルが次回デプロイされた際には、同じコンテナを再利用するので、以降のデプロイメントの時間は短縮されます。
- GPUコンピュートでサービングされるモデルはセットアップに時間を要するので、GPUサービングのオートスケーリングはCPUサービングよりも時間を要します。リクエストのタイムアウトを避けるには多めにプロビジョニングすることをお勧めします。

# サンプルノートブック

以下のノートブックでは、最適化サービングエンドポイントの作成方法を説明しています。

## Llama2モデルに対する最適化LLMサービングノートブック

https://docs.databricks.com/_extras/notebooks/source/machine-learning/large-language-models/optimized-llama-serving.html

## MPTモデルに対する最適化LLMサービングノートブック

https://docs.databricks.com/_extras/notebooks/source/machine-learning/large-language-models/optimized-mpt-serving.html

### Databricksクイックスタートガイド

[Databricksクイックスタートガイド](https://www.amazon.co.jp/dp/B09V1YXFVQ/)


### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)