🧠

MLX で Llama2 を動かしてみる

2023/12/26に公開

Appleシリコン上での実行に対応した機械学習ライブラリMLXが公開されました。
今回は公式が公開している"mlx-examples"リポジトリの"llama"を使って、llama2-7b-chatの実行を試してみます。

commit: 3cf436b529ea58d6c0c0a29c0dd799908cd4497d 2023/12/22

検証環境

MacBook Pro
- Apple M3 Pro
- メモリ 36GB
公式要件としては、以下の環境が示されています。以下要件を満たしていてもメモリが少ない場合、実行するモデルによっては推論ができない可能性があります。
- Using an M series chip (Apple silicon)
- Using a native Python >= 3.8
- macOS >= 13.3

環境構築

まず"mlx-examples"のリポジトリをローカルにクローンします。

git clone https://github.com/ml-explore/mlx-examples.git

その後llamaを実行するコードがあるディレクトリまで移動し、requirements.txtから必要なパッケージをインストールします。
またrequirements.txtに記載がないですが、numpyが必要なのでこちらも合わせてインストールしておきます。

# ディレクトリ移動
cd mlx-examples/llms/llama
# パッケージの一括インストール
pip install -r requirements.txt
# numpyのインストール
pip install numpy

llama2-7b-chatの取得

README内では、Llamaのモデル取得方法として2つの手段が説明されていますが、ここではLlama開発元のMetaからダウンロードする方法を紹介します。

1.利用申請

モデルのダウンロードには、以下の使用申請ページから申請が必要になります。

人によって差があるようですが、自分の場合は10分ほどで申請が通り、"Get started with Llama 2"というタイトルの案内メールが届きました。

2.モデルのダウンロード

届いた案内メールの"How to download the models"に従ってダウンロードを進めます。

まず、llamaのgithubリポジトリをクローンします。

git clone https://github.com/facebookresearch/llama.git

その後クローンしたディレクトリに移動し、ダウンロードスクリプトを実行します。

# ディレクトリ移動
cd llama
# ダウンロードスクリプトの実行
bash download.sh

download.shでは"wget"を使用してファイルのダウンロードを行います。
wgetをインストールしていない場合以下のコマンドでインストールしてください。

brew install wget

実行すると以下のようなメッセージが表示されcustom URLを要求されるので、メールに記載されているアドレスを入力してください。

Enter the URL from email: <メールに記載されている、https://download.llamameta.net/〜から始まるアドレスを入力>

入力が通ると以下のメッセージが表示され、どのLlamaモデルをダウンロードするのか選択できます。
今回はllama2-7b-chatの実行を試すので、"7B-chat"と入力しEnterを押します。

Enter the list of models to download without spaces (7B,13B,70B,7B-chat,13B-chat,70B-chat), or press Enter for all: <ダウンロードしたいモデル名を入力>

以上でllama2-7b-chatのモデルダウンロードまで完了しました。

モデルの変換

変換を行うコード"convert.py"は、上記でダウンロードした"llama/llama-2-7b-chat"ディレクトリへのパスを指定して実行しますが、
"llama/llama-2-7b-chat"ディレクトリ下に"tokenizer.model"が存在することを前提としています。
従ってここでは"llama/llama-2-7b-chat"ディレクトリ下に"llama/tokenizer.model"をコピーして対応します。

# tokenizer.modelをコピーする
cp llama/tokenizer.model llama/llama-2-7b-chat

ダウンロードしたllama2-7b-chatのモデルをMLXで利用できる形式に変換します。
変換を行うコード"convert.py"は、"mlx-examles/llms/llama"にあるのでこちらを利用します。

# ディレクトリ移動
cd mlx-examples/llms/llama
# モデルの変換
python convert.py --torch-path <先ほどダウンロードした"llama/llama-2-7b-chat"ディレクトリへのパスを指定します>

少し待つと変換が完了し、"mlx-examples/llms/llama"ディレクトリ内に"mlx_model"ディレクトリが作成されます。
本ディレクトリ内には変換したモデルファイル類が保存されます。

実行

準備が整ったので実行してみます。
--prompt の後にプロンプトを入力します。

# llama-2-7b-chatの実行
python llama.py --prompt "<プロンプトを入力>"
# 実行するとモデルのロードが始まり、`Press enter to start generation`表示後Enterを押すことで出力が開始される

私の環境では以下のような結果になりました。

python llama.py --prompt "What is XR?"                                         
[INFO] Loading model from disk.
Press enter to start generation
------
What is XR?

XR stands for extended reality, which is a term used to describe a group of technologically advanced and immersive mediums that blend the digital and physical Worlds. XR encompasses a wide spectrum of technological advancements, from virtual reality (VR) and augmented reality (AR) to mixed reality (MR) and beyond.
XR technologists are constantly pushing the limits of what is possible in the digital World, creating new and
------
[INFO] Prompt processing: 0.665 s
[INFO] Full generation: 11.515 s

Prompt processingに0.665秒、Full generationに11.515秒となかなかな処理速度です。
デフォルトでは最大トークン数が100に設定されているので生成文が途中で終わっていますが、実行時に"--max-tokens <好きな数>"を付け加えることでトークン数の変更も可能です。

まとめ

mlx-exampleリポジトリは現在頻繁に更新が行われており整備が進んでいる最中です。
MLXの登場により、今後他のApple Siliconを搭載したデバイスにおいても、オンデバイスで機械学習モデルを実行するならMLXを活用する事がベストな選択肢になるかもしれません。

弊社ではMacBookやiPhone,iPadを始め、2024年発売予定のVision Pro上での活用も視野に入れながらMLXの動向をウォッチしていきたいと思います。

書いた人

ttt-tkmr

髙村哲太

5年間ロボティクスと機械学習を活用した研究開発やサービス・プロダクト開発を中心に従事。
AI技術とXR領域の融合による未来を具現化するためMESONへ入社。

X

MESON Works

MESONの制作実績一覧もあります。ご興味ある方はぜひ見てみてください。

MESON Works