🦙

【cuBLAS】llama-cpp-pythonでのGPU推論入門

2024/01/31に公開

Python

LLM

llama.cpp

tech

はじめに

llama.cppは実はpythonでも使える。

これの良いところはpythonアプリに組み込むときに使える点。GPUオフロードにも対応しているのでcuBLASを使ってGPU推論できる。一方で環境変数の問題やpoetryとの相性の悪さがある。
「llama-cpp-python+cuBLASでGPU推論させる」を目標に、簡易的な備忘録として残しておく。
今回はUbuntuなので、Windowsは適宜READMEのWindws Notesを見ておくこと。

事前準備

なお、nvidia-smiとnvccを各自使える状況にしておく。
cuBLASはCUDAのみのため、いわゆるNVIDIAのGPUしか動かない。
nvidia-smiとnvccは以下記事を参考にインストールすること。

ライブラリのインストール

llama.cppでGPUオフロードする場合はビルドが必要だったが、最近はあらかじめビルドされたものをインストールすることができるようになった。

pip install llama-cpp-python --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu121 --upgrade --force-reinstall --no-cache-dir
cu121の部分は各自のcudaのバージョンにする必要があるので、nvidia-smiで確認する。
もしこれでインストールできた場合、自分でビルドする場合をスキップしても良い。

ライブラリのインストール（自分でビルドする場合）

下記は自分でライブラリをビルドする方法。
CMAKE_ARGSという環境変数の設定を行った後、llama-cpp-pythonをクリーンインストールする。
CMAKE_ARGS="-DLLAMA_CUBLAS=on"
pip install llama-cpp-python --upgrade --force-reinstall --no-cache-dir

ここで大事なのは「pip install」であること。どうやらinstall時にmakeが回っており、poetryでのinstallではcuBLAS対応版としてインストールすることができなかった。

動作確認

あとはいつもと同じ。huggingfaceでgguf落としてきてロードして推論。

from llama_cpp import Llama
llm = Llama(
      model_path="./models/calm2-7b-chat.Q5_K_M.gguf",
      n_gpu_layers=-1, # Uncomment to use GPU acceleration
)
output = llm(
      "USER:こんにちは\nASSISTANT:", # Prompt
      max_tokens=32
)
print(output)