MacBook ProでAirLLMを試してみる
「70B の大規模言語モデルを 1 枚の 4GB GPU カードで推論を実行できる」と言われているAirLLMを、MacBook Proで試してみます。
（注）「実行できる」とはあるけれど、まともなスピードで動く…とは書いていない点に注意しましょう。
試したマシンは、MacBook Pro M3 Proチップ、メモリ18GBです。
1. 準備
venv構築
python3 -m venv air_llm
cd $_
source bin/activateパッケージインストール
pip install mlx airllmlistはこちら。
% 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.42. コードの作成
使用するモデル
なお、使用したモデルは、Llama2ベースで70Bのgarage-bAInd/Platypus2-70B-instructです。
流し込むコード
こんな感じのコードになりました。
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_history3. 試してみる
では、聞いてみましょう。
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.
必要なレイヤーを実行するときにディスクからそのレイヤーをロードし、すべての計算を行った後、メモリを完全に解放できます。
transformerレイヤー毎にメモリのロード・アンロードを繰り返して推論しているのね…。
当然メモリの使用量は抑えられるし、1トークン生成するのに毎回Disk I/Oが伴うから、3トークンで147秒かかると。
4. リソース使用状況
まずはメモリ。pythonは8.91GBとなっています。
続いてCPU。64.0%あたりをうろうろ。
5. まとめ
確かに70Bが動きはしました。
ただ、1トークンあたり約49秒とかなりの時間を要します…。