🙇‍♀️

MLC-LLM で Vulkan など Graphics API + C/C++ で LLM が流行ってほしいメモ

2023/04/30に公開

Vulkan

いきなり彗星のごとく(?)現れた MLC-LLM

Vicuna-7b が iPhone でオンデバイスでうごく... いいね!

https://github.com/mlc-ai/mlc-llm

とりま

で cli サンプル試せます.

3090(160W powerlimit) での結果

cli サンプルでは word 単位での streamling ではないため, tokens/sec 推測しずらいですが,
RWKV ほど爆速というわけではないものの, そこそこ実用的な速度にはできそう?

あと日本語はダメダメでした. Tokenizer 書き直さないとダメかも?

GPU メモリ消費は 3.5 GB くらいでした.

しくみ

Apache TVM をベースにしています.

TVM, なんかあんまり流行って無い気もしていましたが, 着実に発展していたのですね.
(MLIR はどうなるじゃろか... https://qiita.com/syoyo/items/6de8a9f6f5374fd23d94 )

ML モデルを Vulkan などで動くようにするのは

MLC https://mlc.ai/ で native code にコンパルし, C/C++ などから使えるようにする感じですかね.
AITemplate や CTranslate2 みたいな感じですが, backend として Vulkan や Metal などグラフィックス API を選択できるようになっています.
(CUDA もあるよ)

サーバー GPU で Vulan

ちな Tesla P100 などサーバ系 GPU でも Vulkan 動きます.
(最新の H100 は不明であるが...)

ビルド

Linux でビルドしてみます.
現状まともなビルドドキュメント無いです...

を参考に対応でしょうか.

mlc-llm は MLC での LLM 変換だけのようなので, まず TVM をビルドする必要があります.

しかも, 普通の TVM(pip の apache-tvm) ではなく, TVM Unity(ややこしい名前...) https://tvm.apache.org/2021/12/15/tvm-unity

の fork(?) である relax (これも名前ややこしすぎ...) をビルドしてインストールする必要があります.

# まず dependency
$ python -m pip install build pytest

$ git clone https://github.com/mlc-ai/relax.git
$ cd relax
$ git submodule update --init --recursive
$ mkdir build
$ cp cmake/config.cmake build

↑のサイト参考に config.cmake 編集して,

USE_LLVM, USE_VULKAN を有効にしておきます(それぞれ LLVM, Vulkan SDK は別途セットアップしておきます)

少なくとも USE_LLVM は必須です.
そうしないと mlc での変換時,

 Check failed: (bf != nullptr) is false: target.build.llvm is not enabled

とエラーがでます.

あとは make でビルドと LLM モデル変換していけばいけるはずです...

著者環境では, CUDA 版を変換すると

...
free(): invalid pointer                                                             Aborted (core dumped)

と出ました... 多分後処理で abort 発生していると思われますので, 無視でよいでしょう...

あと, mlc-llm の tests/chat.py 実行すると, vicuna-v1-7b でも

FileNotFoundError: [Errno 2] No such file or directory: 'dist/vicuna-v1-7b/float32/params/ndarray-cache.json'

と言われました...

tests/chat.py は,

にある cli コマンドとはまた別っぽい模様?

mlc-chat-nightly のビルドスクリプトはどうなっているのだろうか...
(conda-forge を調べないとだめっぽそう)

vicuna-13b 変換(W.I.P.)

デフォでは vicuna-13b は対応していません.
そのうち対応されるでしょうが, とりま無理やり以下の対応してみました
(しかし変換はできるが, 実行時うまくいかなかった)

diff --git a/mlc_llm/relax_model/llama.py b/mlc_llm/relax_model/llama.py
index ac6be6b..9b00ba6 100644
--- a/mlc_llm/relax_model/llama.py
+++ b/mlc_llm/relax_model/llama.py
@@ -50,6 +50,7 @@ class LlamaConfig:

 MODEL_CONFIG = {
     "vicuna-v1-7b": {},
+    "vicuna-v1-13b": {'hidden_size': 5120, 'num_hidden_layers': 40, 'num_attention_heads': 40, 'intermediate_size': 13824 },
     "llama-7b": {},
 }

diff --git a/mlc_llm/utils.py b/mlc_llm/utils.py
index 898080a..cc7c8ca 100644
--- a/mlc_llm/utils.py
+++ b/mlc_llm/utils.py
@@ -15,6 +15,7 @@ def argparse_add_common(args: argparse.ArgumentParser) -> None:
         default="vicuna-v1-7b",
         choices=[
             "vicuna-v1-7b",
+            "vicuna-v1-13b",
             "dolly-v2-3b",
             "dolly-v2-7b",
             "dolly-v2-12b"

変換は 1 分もあれば完了, メモリ消費は 40 GB くらいのようでした.

vicuna-13b Vulkan 実行

メモリ消費は 6.2 GB ほど.

|    0   N/A  N/A     44809    C+G   mlc_chat_cli                     6214MiB |

ただ, 実行すると返答がかえってきません. 何かしらパラメータ設定が 13b 用にできていないところがあるようです..

T.B.C.

GPU のグラフィックス Unit 活用で LLM on GPU 飛躍ワンチャンあるか!?

Vulkan, Metal などグラフィックス系 API 使うことで, HW 圧縮テクスチャユニットが活用できます.
LLM の int4 や int8 などの量子化はこの HW 圧縮テクスチャユニット(特に HDR の圧縮フォーマット)を活用してより効率化できるのが期待できそう(TVM などで Texture unit 活用は対応まだされてなさそうであるが...).

あとはレイトレ BVH コアで embedding 計算(kNN 探索などを BVH で効率化)の効率化とかですかね.

Graphics API で LLM 流行ってきたらそのうち誰かがやってくれるでしょう.

...誰かやってくれるよね...

学習は無理かも?

学習(training)はサポートしてなさそげですね.

pytorch から backward のネットワークを export して頑張ればワンチャンあるかもはしれませんが...

TODO

Vicuna-13b の変換と実行を調査する

しくみ