DatabricksでLlama.cppを手軽に動かす

軽めの内容です。

導入

Llama.cppはLLM用のランタイムであり、ローカルLLM界隈では非常によく利用されている推論パッケージです。
CPU推論/GPU推論両方をカバーし、多くの量子化にも対応しています。

最近使ってみて、いつのまにかDatabricksのノートブック上でもGPU推論が手軽に利用できるようになっていたので試してみました。

検証環境はDatabricks on AWS、DBRは15.4ML、インスタンスタイプはg5.xlargeを利用しました。

Step1. パッケージインストール

ノートブックを作成し、Llama-cpp-pythonパッケージをインストール。
この際、CUDA対応のプレビルド版をインストールするように指定します。
(DBR 15.4MLの場合、CUDA 12.1がインストール済みですのでcu121のURLを指定します)

%pip install llama-cpp-python --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu121

dbutils.library.restartPython()

Step2. モデルのロード

事前にダウンロードしておいたGGUF形式のモデルをロードします。
今回は以下のリポジトリのモデルを利用させていただきました。

from llama_cpp import Llama

llm = Llama(
      model_path="/Volumes/training/llm/model_snapshots/models--mmnga--HODACHI-Borea-Phi-3.5-mini-Instruct-Common-gguf/HODACHI-Borea-Phi-3.5-mini-Instruct-Common-Q4_K_M.gguf",
      n_gpu_layers=-1, 
      n_ctx=2048, 
)

Step3. チャット形式での推論

Chat形式でのCompletionを生成します。

from pprint import pprint

output = llm.create_chat_completion(
    messages=[
        {"role": "system", "content": "あなたは親切な日本語AIアシスタントです。"},
        {"role": "user", "content": "Databricksとは何ですか？"},
    ],
    max_tokens=512,
)
pprint(output)

出力

Llama.generate: 45 prefix-match hit, remaining 1 prompt tokens to eval

llama_print_timings:        load time =     303.78 ms
llama_print_timings:      sample time =      21.78 ms /   512 runs   (    0.04 ms per token, 23508.88 tokens per second)
llama_print_timings: prompt eval time =       0.00 ms /     0 tokens (    -nan ms per token,     -nan tokens per second)
llama_print_timings:        eval time =    4217.96 ms /   512 runs   (    8.24 ms per token,   121.39 tokens per second)
llama_print_timings:       total time =    4575.69 ms /   512 tokens
{'choices': [{'finish_reason': 'length',
              'index': 0,
              'logprobs': None,
              'message': {'content': 'Databricksは、データ分析、機械学習、およびビッグデータプラットフォームの開発を支援するクラウドサービスです。データサイエンティストやデータアナリスト向けに、統合されたデータサイエンス環境を提供し、Python、R、Scalaなどのさまざまなプログラミング言語でアクセスできるクラウドベースのデータ科学（Data '
                                     'Science as a Service）プラットフォームを提供しています。\n'
                                     '\n'
                                     'Databricksは以下の機能を提供しています：\n'
                                     '\n'
                                     '1. '
                                     '**データレイク**：大規模なデータセットを保存し、管理するための高度なデータレイクサービス。\n'
                                     '2. '
                                     '**データクレンディング**：データを整形し、欠損値を処理し、カテゴリ変数をエンコードするためのツールとAPI。\n'
                                     '3. '
                                     '**データ分析**：SQLやデータウェアハウス（DW）ツールを使用してデータを分析し、洞察を得るための統合された環境。\n'
                                     '4. **機械学習**：Databricks MLflowやDatabricks '
                                     'Deep '
                                     'Learningなどの機械学習フレームワークを使用してモデルを開発、テスト、保存、および実行するための環境。\n'
                                     '5. **データビジュアリング**：データを視覚的に表現し、洞察を',
                          'role': 'assistant'}}],
 'created': xxxxxxx,
 'id': 'chatcmpl-63544d90-d89f-463b-a954-2d1a934a802a',
 'model': '/Volumes/training/llm/model_snapshots/models--mmnga--HODACHI-Borea-Phi-3.5-mini-Instruct-Common-gguf/HODACHI-Borea-Phi-3.5-mini-Instruct-Common-Q4_K_M.gguf',
 'object': 'chat.completion',
 'usage': {'completion_tokens': 512, 'prompt_tokens': 46, 'total_tokens': 558}}

出力形式としてはOpenAI APIと同形式で結果が問題なく得られました。

まとめ

DatabricksでGPU推論をする場合に結構面倒くさかった印象があるのですが、事前ビルドのパッケージが提供されていたりと手軽に利用できるようになっていました。
個人的にExLlamaV2のような他の推論エンジンを使うことが多かったのですが、使い分け含めていろいろ試してみようと思います。