DatabricksでLlama3.1 70BのマルチGPU推論を行う - CCCMKホールディングス TECH LABの Tech Blog

はじめに

こんにちは。テックラボの高橋です。本記事ではDatabricks上でLlama3.1のマルチGPU推論ができるかどうか試していきます。 8Bのモデルだと16G 1枚で実行できてしまうので、あえて70Bのモデルに挑戦してみます。

※モデル毎の速度比較については下記リンク先をご参照ください。

techblog.cccmkhd.co.jp

環境

Databricks runtime 15.3 ML
GPU T4 16G x 4

llama.cppのインストール

今回はモデルの量子化を活用した推論高速化ツールであるllama.cppを用います。

Databricksにllama.cppをインストールする場合は、init scriptを使うと良いでしょう。

mkdir -p ${YOUR_WORK_DIR}
cd ${YOUR_WORK_DIR}

# https://github.com/ggerganov/llama.cpp/blob/c8a0090922bad576623de4aae227717085249262/docs/build.md?plain=1#L172
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp
PATH="/usr/local/cuda/bin:$PATH" make GGML_CUDA=1

なお、デフォルトでは/databricks/init_scripts/ にログが出力されるので、ビルドが上手くいかない場合は参照してみてください。

CUDAビルドの確認

llama.cppはオプション無しでビルドするとCPU推論で実行してしまいます。実際に推論を行って、llama.cppが実行時にCUDAを利用しているか見てみます。

確認用にMeta-Llama-3.1-8Bのgguf変換済みモデルを実行します。

https://huggingface.co/bullerwins/Meta-Llama-3.1-8B-Instruct-GGUF

!curl -L https://huggingface.co/bullerwins/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q2_K.gguf?download=true -o ${YOUR_WORK_DIR}/llama.cpp/models/Meta-Llama-3.1-8B-Instruct-Q2_K.gguf

!cd ${YOUR_WORK_DIR} && ./llama-cli -m ${YOUR_WORK_DIR}/models/Meta-Llama-3.1-8B-Instruct-Q2_K.gguf -p "[INST]美味しい白菜の育て方を教えて。日本語で。[/INST]"

CPUではなく、GPUが利用されているでしょうか？

llama.cppの実行時に、別途nvidia-smiを行うとGPUメモリが使われているかチェックできるかと思います。

（出力が少し変かもしれませんが、後ほど説明があります）

マルチGPUでの実行

先程のモデルをT4 4枚の環境で実行してみます

!cd ${YOUR_WORK_DIR}/llama.cpp && ./llama-cli -m ${YOUR_WORK_DIR}/llama.cpp/models/Meta-Llama-3.1-8B-Instruct-Q2_K.gguf -c 2048 -n 128 -p "[INST]美味しい白菜の育て方を教えて。日本語で。[/INST]"

もしかすると、推論中にnvidia-smiの結果を見ても、GPUが1枚しか使われていないかもしれません。

ここでllama.cppのドキュメントを見ると、

github.com

-ngl N, --n-gpu-layers N: When compiled with GPU support, this option allows offloading some layers to the GPU for computation. Generally results in increased performance.

という記載があります。

実は、llama.cppはこの設定が無いとGPUを推論時に利用しません。今まで実行したnvidia-smiの結果においても、 GPUメモリは利用されていても GPU-Utilの値は上がっていなかったのではないでしょうか。

このオプションを設定してみましょう。

!cd ${YOUR_WORK_DIR}/llama.cpp_gpu && ./llama-cli -m ${YOUR_WORK_DIR}/llama.cpp/models/Meta-Llama-3.1-8B-Instruct-Q2_K.gguf -ngl 81 -c 2048 -n 128 -p "[INST]美味しい白菜の育て方を教えて。日本語で。[/INST]"

nvidia-smiの結果が変わっているでしょうか。

それでは、さらに大きいモデルを推論してみましょう。

CyberAgent社が公開しているLlama-3.1-70B-Japanese-Instruct-2407の、gguf変換モデルを利用します。

https://huggingface.co/mmnga/Llama-3.1-70B-Japanese-Instruct-2407-gguf/tree/main

!cd  ${YOUR_WORK_DIR}/llama.cpp && ./llama-cli -m ${YOUR_WORK_DIR}/llama.cpp/models/Llama-3.1-70B-Japanese-Instruct-2407-gguf/Llama-3.1-70B-Japanese-Instruct-2407-IQ4_XS.gguf --temp 0.1 -ngl 81 -c 2048 -n 128 -p "[INST]美味しい白菜の育て方を教えて。箇条書きで[/INST]"

美味しい白菜の育て方を箇条書きでご紹介します。 - 白菜は冷涼な気候を好むので、涼しい季節に種をまきます。一般的には、春まき（3月下旬～4月上旬）と秋まき（8月下旬～9月上旬）が適しています。 - 白菜は日当たりが良く、風通しの良い場所を好みます。日照時間が短いと結球しないので、日当たりが良い場所……

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.161.07             Driver Version: 535.161.07   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla T4                       Off | 00000001:00:00.0 Off |                  Off |
| N/A   37C    P0              44W /  70W |   9829MiB / 16384MiB |      7%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   1  Tesla T4                       Off | 00000002:00:00.0 Off |                  Off |
| N/A   37C    P0              27W /  70W |   9387MiB / 16384MiB |     34%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   2  Tesla T4                       Off | 00000003:00:00.0 Off |                  Off |
| N/A   38C    P0              27W /  70W |   9387MiB / 16384MiB |     39%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   3  Tesla T4                       Off | 00000004:00:00.0 Off |                  Off |
| N/A   37C    P0              57W /  70W |   9765MiB / 16384MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

無事にマルチGPUで推論されているようですね。