预配吞吐量基础模型 API

本文演示如何使用 基础模型 API 预配的吞吐量部署模型。 Databricks 建议为生产工作负荷预配吞吐量,并为基础模型提供优化推理,并提供性能保证。

什么是预配吞吐量?

预配吞吐量是指可以同时向某个终结点提交的请求所对应的令牌数量。 预配置的吞吐量服务端点是根据每秒可发送到端点的令牌范围进行配置的专用端点。

有关详细信息,请参阅以下资源:

有关支持的模型体系结构列表,请参阅 马赛克 AI 模型服务中的支持基础模型,这些模型适用于指定吞吐量终端。

要求

请参阅 要求。 有关部署优化的基础模型,请参阅 部署微调的基础模型

重要说明

此功能目前为公共预览版

Databricks 建议使用预安装在 Unity 目录中的基础模型。 可以在架构 systemai)中的目录 system.ai 下找到这些模型。

部署基础模型:

  1. 在目录资源管理器中导航到 system.ai
  2. 单击要部署的模型的名称。
  3. 在模型页上,单击“服务此模型”按钮
  4. 此时会显示“创建服务终结点”页面。 请参阅 使用 UI创建预配的吞吐量终结点。

注释

若要从 system.ai Unity 目录中部署 Meta Llama 模型,必须选择适用的 “指示 ”版本。 在 Unity Catalog 中,来自 system.ai 的部署不支持 Meta Llama 模型的基本版本。 请参阅 Databricks 上托管的以及支持的 Meta Llama 模型变体的基础模型。

从 Databricks 市场部署基础模型

或者,可以从 Databricks Marketplace将基础模型安装到 Unity 目录。

你可以搜索模型系列,然后从模型页面选择“获取访问权限”并提供登录凭据,以将模型安装到 Unity Catalog

将模型安装到 Unity 目录后,可以使用服务 UI 创建提供终结点的模型。

部署经过微调的基础模型

如果无法在 system.ai 架构中使用模型,或者从 Databricks Marketplace 安装模型,则可以通过将模型记录到 Unity 目录来部署经过微调的基础模型。 本部分和以下部分演示如何设置代码以将 MLflow 模型记录到 Unity 目录,以及使用 UI 或 REST API 创建预配的吞吐量终结点。

请参阅 Databricks 上托管的基础模型 ,了解受支持的 Meta Llama 3.1、3.2 和 3.3 微调模型及其区域可用性。

要求

  • MLflow 2.11 或更高版本仅支持部署微调的基础模型。 Databricks Runtime 15.0 ML 及更高版本预安装兼容的 MLflow 版本。
  • Databricks 建议在 Unity 目录中使用模型,以便更快地上传和下载大型模型。

定义目录、架构和模型名称

若要部署经过微调的基础模型,请定义所选的目标 Unity 目录、架构和模型名称。

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("meta-llama/Llama-3.3-70B-Instruct", torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.3-70B-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,则必须改为在同一 metadata 函数的 mlflow.transformer.log_model() 参数中指定任务。

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

预配吞吐量也支持基本和大型 GTE 嵌入模型。 下面的示例演示了如何记录模型 Alibaba-NLP/gte-large-en-v1.5,以便可以使用预配吞吐量为其提供服务:

model = AutoModel.from_pretrained("Alibaba-NLP/gte-large-en-v1.5")
tokenizer = AutoTokenizer.from_pretrained("Alibaba-NLP/gte-large-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": "gte-large"
        }
    )

在模型记录到 Unity Catalog 中之后,继续参阅使用 UI 创建预配吞吐量终结点,以便创建具有预配吞吐量的模型服务终结点。

使用 UI 创建预配的吞吐量终结点

记录的模型在 Unity 目录中后,使用以下步骤创建预配的吞吐量服务终结点:

  1. 导航到工作区中的服务 UI
  2. 选择 创建服务终结点
  3. 实体 字段中,从 Unity 目录中选择模型。 对于符合条件的模型,服务实体的 UI 会显示“预配吞吐量”屏幕。
  4. 在“最多”下拉列表中,您可以配置终结点的每秒最大令牌吞吐量。
    1. 预配吞吐量终结点会自动缩放,因此可以选择“修改”以查看终结点可以纵向缩减到的最小每秒令牌数

预配的吞吐量

使用 REST API 创建预配的吞吐量终结点

若要使用 REST API 在预配吞吐量模式下部署模型,必须在请求中指定 min_provisioned_throughputmax_provisioned_throughput 字段。 如果要使用 Python,也可以使用 MLflow 部署 SDK 创建终结点

若要确定模型的合适预配吞吐量范围,请参阅以增量方式获取预配吞吐量的方法

import requests
import json

# Set the name of the MLflow endpoint
endpoint_name = "prov-throughput-endpoint"

# Name of the registered MLflow model
model_name = "ml.llm-catalog.foundation-model"

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

# Get the API endpoint and token for the current notebook context
API_ROOT = "<YOUR-API-URL>"
API_TOKEN = "<YOUR-API-TOKEN>"

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

聊天补全任务的对数概率

对于聊天补全任务,可以使用 logprobs 参数来提供某个令牌在大型语言模型生成过程中被采样的对数概率。 可以将 logprobs 用于各种方案,包括分类、评估模型不确定性和运行评估指标。 有关参数详细信息,请参阅 聊天任务

以递增方式获取调配吞吐量

预配的吞吐量以每秒令牌增量提供,具体增量因模型而异。 为了根据需求确定合适的范围,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": 1580
}

笔记本示例

以下笔记本演示了如何创建预配的吞吐量基础模型 API 的示例:

GTE 模型笔记本的预配吞吐量服务

获取笔记本

BGE 模型笔记本的预配吞吐量服务

获取笔记本

以下笔记本演示如何在 Unity 目录中下载和注册 DeepSeek R1 精炼的 Llama 模型,以便使用基础模型 API 预配置的吞吐量端点部署它。

用于 DeepSeek R1 蒸馏 Llama 模型笔记本的预配吞吐量

获取笔记本

限制

  • 模型部署可能会由于 GPU 容量问题而失败,这会导致终结点创建或更新期间超时。 联系 Databricks 帐户团队以帮助解决。
  • 基础模型 API 的自动缩放速度比 CPU 模型服务慢。 Databricks 建议过度预配以避免请求超时。
  • GTE v1.5(英语)不生成规范化嵌入。