尝试在Databricks的LLM优化模型服务中部署Llama2模型

1 年 ago
新, 韵
3 minutes

我会为您进行浏览这个示例笔记本。

 

Llama2服务的优化示例

通过最优化LLM保存，您可以将Databricks模型服务部署到具有自动优化功能的最新OSS LLM和GPU，以提供改进的延迟和吞吐量。目前，我们支持Llama2和Mosaic MPT类模型的优化。

在这个示例中，我们将演示以下内容:

    1. 从Hugging Face transformers下载模型

 

    1. 将模型记录在Databricks的Unity Catalog或Workspace Registry中，以优化的serving支持格式

 

    启用模型的优化serving。

前置条件

    • ノートブックを十分なメモリーを持つクラスターにアタッチ

 

    • MLflowバージョン2.7.0以降がインストールされていること

 

    7B以上のサイズのモデルを取り扱う際には特にUCでのモデル管理を行うこと

步骤1：为了优化LLM存储，记录模型。

# 必要な依存関係を更新/インストール
!pip install -U mlflow
!pip install -U transformers
!pip install -U accelerate
!pip install -U flash_attn # DBR 13.3 MLでエラーになったので追加
dbutils.library.restartPython()

输入 Hugging Face 的令牌。更多详情请点击此处。

import huggingface_hub
# 既にhugging faceにログインしている際にはスキップ
huggingface_hub.login()

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf", 
    torch_dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf"
)

import mlflow
from mlflow.models.signature import ModelSignature
from mlflow.types.schema import ColSpec, Schema
import numpy as np

# モデルの入出力スキーマの定義
input_schema = Schema([
    ColSpec("string", "prompt"),
    ColSpec("double", "temperature", optional=True),
    ColSpec("integer", "max_tokens", optional=True),
    ColSpec("string", "stop", optional=True),
    ColSpec("integer", "candidate_count", optional=True)
])

output_schema = Schema([
    ColSpec('string', 'predictions')
])

signature = ModelSignature(inputs=input_schema, outputs=output_schema)

# 入力サンプルの定義
input_example = {
    "prompt": np.array([
        "Below is an instruction that describes a task. "
        "Write a response that appropriately completes the request.\n\n"
        "### Instruction:\n"
        "What is Apache Spark?\n\n"
        "### Response:\n"
    ])
}

要启用最优化服务，您需要在记录模型时调用mlflow.transformers.log_model，并包含以下方式的字典作为附加的元数据:

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

在这里，我们指定了用于模型服务端点的API签名。

import mlflow

# UCでのモデル管理を行わない場合には以下の行をコメントアウトします
# 3レベルの名前空間ではなく、シンプルにモデル名を指定します
mlflow.set_registry_uri('databricks-uc')
CATALOG = "quickstart_catalog_taka"
SCHEMA = "llm-catalog"
registered_model_name = f"{CATALOG}.{SCHEMA}.llama2-7b"

# 新規MLflowランをスタート
with mlflow.start_run():
    components = {
        "model": model,
        "tokenizer": tokenizer,
    }
    mlflow.transformers.log_model(
        transformers_model=components,
        task = "text-generation",
        artifact_path="model",
        registered_model_name=registered_model_name,
        signature=signature,
        input_example=input_example,
        metadata={"task": "llm/v1/completions"}
    )
Screenshot 2023-10-08 at 17.12.08.png
Screenshot 2023-10-08 at 17.11.42.png

步骤2：设置和创建模型保存的GPU端点。

要更改终端节点名，请修改以下单元格。调用终端节点API之后，记录的Llama2模型将自动部署到优化的LLM服务中。workload_type需要指定环境中可用规格的GPU。

# MLflowエンドポイント名の設定
endpoint_name = "taka-llama2-7b"

# 登録されたMLflowモデル名
model_name = registered_model_name

# MLflowモデルの最新バージョンの取得
model_version = 1

# 計算資源タイプの指定 (CPU, GPU_SMALL, GPU_MEDIUM, etc.)
workload_type = "GPU_MEDIUM"

# 計算資源のスケールアウトサイズの指定 (Small, Medium, Large, etc.)
workload_size = "Small" 

# ゼロへのスケールの指定 (CPUエンドポイントでのみサポート)
scale_to_zero = False 

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

import requests
import json

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

参考端点

Screenshot 2023-10-09 at 15.00.59.png

步骤3：查询终节点

当准备好终端点后，通过进行API请求可以进行查询。这取决于模型的大小和复杂性，但准备时间需要30分钟以上。

GPU Medium的情况下，大约30分钟就会准备好。
data = {
    "inputs": {
        "prompt": [
            "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"
        ]
    },
    "params": {
        "max_tokens": 100, 
        "temperature": 0.0
    }
}

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(), indent=4))

查询成功！

{
    "predictions": [
        {
            "candidates": [
                {
                    "text": "Apache Spark is an open-source data processing engine that can handle large-scale data processing tasks. It was developed at the University of California, Berkeley and is now maintained by Apache Software Foundation. Spark provides high-level APIs in Java, Python, Scala, and R, and supports a wide range of data sources, including Hadoop Distributed File System (HDFS), Cassandra, HBase, and Hive. Spark's core features include in-memory computing, fault",
                    "metadata": {
                        "finish_reason": "length"
                    }
                }
            ],
            "metadata": {
                "input_tokens": 41,
                "output_tokens": 100,
                "total_tokens": 141
            }
        }
    ]
}

Databricks快速入门指南

Databricks快速入门指南

数据脉冲免费试用

Databricks免费试用