MacBook ProでAirLLMを試してみる

「70B の大規模言語モデルを 1 枚の 4GB GPU カードで推論を実行できる」と言われているAirLLMを、MacBook Proで試してみます。

（注）「実行できる」とはあるけれど、まともなスピードで動く…とは書いていない点に注意しましょう。

Unbelievable! Run 70B LLM Inference on a Single 4GB GPU with This NEW Technique

試したマシンは、MacBook Pro M3 Proチップ、メモリ18GBです。

1. 準備

venv構築

python3 -m venv air_llm
cd $_
source bin/activate

パッケージインストール

pip install mlx airllm

listはこちら。

% pip list
Package            Version
------------------ ------------
accelerate         0.25.0
aiohttp            3.9.1
aiosignal          1.3.1
airllm             2.8.3
async-timeout      4.0.3
attrs              23.1.0
certifi            2023.11.17
charset-normalizer 3.3.2
coloredlogs        15.0.1
datasets           2.16.0
dill               0.3.7
filelock           3.13.1
frozenlist         1.4.1
fsspec             2023.10.0
huggingface-hub    0.20.1
humanfriendly      10.0
idna               3.6
Jinja2             3.1.2
MarkupSafe         2.1.3
mlx                0.0.6
mpmath             1.3.0
multidict          6.0.4
multiprocess       0.70.15
networkx           3.2.1
numpy              1.26.2
optimum            1.16.1
packaging          23.2
pandas             2.1.4
pip                23.3.2
protobuf           4.25.1
psutil             5.9.7
pyarrow            14.0.2
pyarrow-hotfix     0.6
python-dateutil    2.8.2
pytz               2023.3.post1
PyYAML             6.0.1
regex              2023.12.25
requests           2.31.0
safetensors        0.4.1
scipy              1.11.4
sentencepiece      0.1.99
setuptools         58.0.4
six                1.16.0
sympy              1.12
tokenizers         0.15.0
torch              2.1.2
tqdm               4.66.1
transformers       4.36.2
typing_extensions  4.9.0
tzdata             2023.3
urllib3            2.1.0
xxhash             3.4.1
yarl               1.9.4

2. コードの作成

使用するモデル

なお、使用したモデルは、Llama2ベースで70Bのgarage-bAInd/Platypus2-70B-instructです。

garage-bAInd/Platypus2-70B-instruct · Hugging Face

流し込むコード

こんな感じのコードになりました。

from airllm import AutoModel
import mlx.core as mx
import time
 #llm  = "elyza/ELYZA-japanese-Llama-2-13b-instruct"
llm = "garage-bAInd/Platypus2-70B-instruct"

# モデルの準備
model = AutoModel.from_pretrained(llm)

# Llama 2 based
B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"

def build_prompt(user_query, chat_history=None):
    prompt = "{chat_history}{b_inst} {system}{prompt} {e_inst} ".format(
        chat_history=chat_history,
        b_inst=B_INST,
        system=f"{B_SYS}{DEFAULT_SYSTEM_PROMPT}{E_SYS}",
        prompt=user_query,
        e_inst=E_INST,
    )
    return prompt

def q(user_query, chat_history=None):
    start = time.process_time()
    # 推論の実行
    prompt = build_prompt(user_query, chat_history)
    input_ids = model.tokenizer(
        model.tokenizer.bos_token + prompt,
        return_tensors="np",
        return_attention_mask=False,
        truncation=True,
        max_length=128,
        padding=False
    )
    output_ids = model.generate(
        mx.array(input_ids['input_ids']),
        max_new_tokens=3,
        use_cache=True,
        return_dict_in_generate=True
    )
    print(output_ids)
    chat_history = prompt + output_ids
    end = time.process_time()
    print(end - start)
    return chat_history

3. 試してみる

では、聞いてみましょう。

chat_history = q("What is Doraemon?")

応答はこちら。

running layers: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 80/80 [01:16<00:00,  1.05it/s]
running layers: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 80/80 [01:12<00:00,  1.10it/s]
running layers: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 80/80 [01:12<00:00,  1.10it/s]


D
147.511253

うん？
running layers 80/80が3つ？
147秒？

あー、そういうことか。

We can load whichever layer is needed from disk when executing that layer, do all the calculations, and then completely free the memory after.

必要なレイヤーを実行するときにディスクからそのレイヤーをロードし、すべての計算を行った後、メモリを完全に解放できます。

Unbelievable! Run 70B LLM Inference on a Single 4GB GPU with This NEW Technique | by Gavin Li | Nov, 2023 | AI Advances (gopubby.com)

transformerレイヤー毎にメモリのロード・アンロードを繰り返して推論しているのね…。
当然メモリの使用量は抑えられるし、1トークン生成するのに毎回Disk I/Oが伴うから、3トークンで147秒（＝49 tps）かかると。

4. リソース使用状況

まずはメモリ。pythonは8.91GBとなっています。

アクティビティモニタのメモリより

続いてCPU。64.0%あたりをうろうろ。

アクティビティモニタのCPUより

5. まとめ

確かに70Bが動きはしました。
ただ、1トークンあたり約49秒とかなりの時間を要します…。