DatabricksによるDeepSeek R1(distilled Llama 8B)のサービング

Last updated at 2025-02-01Posted at 2025-02-01

ブログ記事訳していたら @isanakamishiro2 さんに先を越されました。速い！

こちらのノートブックを翻訳しながら私も動かします。

プロビジョニングスループットを使用してDeepSeek R1(distilled Llama 8B)をさーびんぐす

このノートブックでは、DeepSeek R1 distilled Llama 8BモデルをUnity Catalogにダウンロードして登録し、Foundation Model APIのプロビジョニングスループットエンドポイントを使用してデプロイする方法を示します。

HuggingFaceの`transformers`ライブラリをインストールする

huggingface transformersをインストールしましょう。

!pip install transformers==4.44.2 mlflow
%restart_python

DeepSeek R1 distilled Llama 8B をダウンロードする

次のコードは、DeepSeek R1 distilled Llama 8B モデルをローカルマシンにダウンロードします。

model_id = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"

huggingfaceのキャッシュフォルダをローカルSSDドライブに設定します。

import os

LOCAL_DISK_HF = "/local_disk0/hf_cache"
os.makedirs(LOCAL_DISK_HF, exist_ok=True)
os.environ["HF_HOME"] = LOCAL_DISK_HF
os.environ["HF_DATASETS_CACHE"] = LOCAL_DISK_HF
os.environ["TRANSFORMERS_CACHE"] = LOCAL_DISK_HF

最初にデプロイするチェックポイントをダウンロードします。

from huggingface_hub import snapshot_download
snapshot_download(model_id)

ダウンロードしたモデルをUnity Catalogに登録する

以下のコードは、ダウンロードしたモデルをUnity Catalogに登録するためのランを開始してログを記録する方法を示しています。

import mlflow
import transformers

mlflow.set_registry_uri("databricks-uc")

model_id = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
uc_model_name = "deepseek_r1_distilled_llama8b_v1"

task = "llm/v1/chat"
model = transformers.AutoModelForCausalLM.from_pretrained(model_id)
tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)

transformers_model = {"model": model, "tokenizer": tokenizer}

with mlflow.start_run():
    model_info = mlflow.transformers.log_model(
        transformers_model=transformers_model,
        artifact_path="model",
        task=task,
        registered_model_name=f"users.takaaki_yayoi.{uc_model_name}",
        metadata={
            "task": task,
            "pretrained_model_name": "meta-llama/Llama-3.3-8B-Instruct",
            "databricks_model_family": "LlamaForCausalLM",
            "databricks_model_size_parameters": "8b",
        },
    )

MLfowエクスペリメントに記録されます。

そして、Unity Catalogにモデルが登録されます。

モデル提供のためのプロビジョニングスループットエンドポイントを作成する

次のコードは、Unity Catalogにダウンロードして登録したLlama 70Bを提供するためのプロビジョニングスループットモデル提供エンドポイントを作成する方法を示しています。

from mlflow.deployments import get_deploy_client


client = get_deploy_client("databricks")


endpoint = client.create_endpoint(
    name=uc_model_name,
    config={
        "served_entities": [{
            "entity_name": f"users.takaaki_yayoi.{uc_model_name}",
            "entity_version": model_info.registered_model_version,
             "min_provisioned_throughput": 0,
             "max_provisioned_throughput": 9500,
            "scale_to_zero_enabled": True
        }],
        "traffic_config": {
            "routes": [{
                "served_model_name": f"{uc_model_name}-{model_info.registered_model_version}",
                "traffic_percentage": 100
            }]
        }
    }
)