DatabricksでDeepSeek R1(Distilled Llama 8B)を動かす

# 導入

こちらのDatabricks公式blogにて、DeepSeek R1を使った蒸留モデルのサービングについて記載されています。

https://www.databricks.com/blog/deepseek-r1-databricks

既に日本語訳記事も上げられていますね。

https://qiita.com/taka_yayoi/items/30ff228f1e5d73f03465

最初タイトルを見たときは「Databricks公式のDeepSeek-R1モデルエンドポイントが公開されたのか！？」と思ったのですが、どうやらまだそういうわけではなさそう。
記事の内容としては蒸留によって構築されたLlamaベースのモデルをDatabricks上でサービングするものです。

せっかくなので、サンプルノートブックをベースにウォークスルーしてみたいと思います。

https://docs.databricks.com/en/_extras/notebooks/source/machine-learning/large-language-models/prov-throughput-deepseek-r1-distill-llama.html

環境はDatabricks on AWS、DBRは15.4(機械学習)、クラスタのインスタンスタイプはr6id.2xlargeです。GPUクラスタは使わずにCPUクラスタで実施してみました。

# Step1. モデルのロギング&サービング

まずは蒸留で構築されたモデルをMLflowにロギングし、Databricks Mosaic AI Model Servingにてサービングします。

まず、必要なパッケージをインストール。
mlflowのバージョンを指定していませんが、今回は```2.20.1```がインストールされています。

```python
!pip install transformers==4.44.2 mlflow
%restart_python
```

次にモデルをhuggingfaceから取得。
今回はこちらのLlama3.1 8Bモデルをベースに構築されたモデルを利用します。

https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-8B

```python
import os
from huggingface_hub import snapshot_download

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

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

snapshot_download(model_id)
```

モデルをMLflowを使ってロギングします。
この後Mosai AI Model Servingから利用するため、Unity Catalogへも登録します。

```python
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"training.llm.{uc_model_name}",
        metadata={
            "task": task,
            "pretrained_model_name": "meta-llama/Llama-3.3-8B-Instruct",
            "databricks_model_family": "LlamaForCausalLM",
            "databricks_model_size_parameters": "70b",
        },
    )
```

最後にMosaic AI Model Servingへエンドポイントを作成します。

```python
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"training.llm.{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
            }]
        }
    }
)
```

これで利用する準備ができました。

# Step2. 使ってみる

では使ってみましょう。
DatabricksのPlaygroundからチャットしてみます。

せっかくなので、日本語でよくある数字比較の質問をしてみます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/3004accf-27e6-6049-2623-4247861da854.png)

<think>タグの間で推論過程の出力がされ、最終結果として正しい内容を回答されてました。
推論過程も日本語になったのは少し意外でした。英語 or 中国語になると思っていたので。
とはいえ、いくつか質問を試す限りだとそこまで日本語性能は高く無さそうです。

※　実行にあたって、温度(temperature)は0.6、最大トークン数を4000にしています。

# まとめ

Databricks上でDeepSeek R1を使って蒸留されたLlamaベースのモデルをサービングして使ってみました。
LlamaベースだとDatabricks上で高速推論が可能なので相性がいいですね。
8Bサイズの軽量モデルでも推論性能は高く、ユースケースによっては面白いことができそうです。

DeepSeek V3/R1登場後、ローカルLLM界隈はいろいろ盛り上がっていて面白いですね。
特に蒸留された日本語モデルも多く登場しており、いろいろ試してみたいと思います。

また、DeepSeek-R1がAWSやNVidiaのプラットフォーム上でもホスティングされて提供され始めていますね。Databricksの基盤モデルとしてもPay-per-Tokenのサービスとして(お安く)公開されるとユーザとしてありがたいなあ。