📖

MLX SwiftでMeta Llama 3を動かす

2024/04/28に公開

Swift

macOSやiOSデバイスでLLMの推論を動かすにはllama.cpp^[1]やMLX^[2]が利用できます

MLXはAppleによるプロジェクトでSwift APIも公開されています^[3]

MLX Swiftを使ってデバイス上でLLMを実行するにはMLX Swift Examplesリポジトリで公開されているソースコードが参考になります

この中のLLMEvalというアプリは、Hugging Faceから任意のモデルをダウンロードしてきてテキスト生成を実行します。MacとiOSでも動作します

LLMEvalは標準で以下のモデルに切り換えて実行できます

Llama 3を動かす

リストされているllamaはCodeLlamaなので、ここにLlama 3を追加して動かしてみます

LLMEvalが依存しているLLMライブラリ(MXLL)のソースコードを更新します

Models.swiftに以下のように記述します

+     public static let llama38bInstruct4bit = ModelConfiguration(
+        id: "mlx-community/Meta-Llama-3-8B-Instruct-4bit",
+        overrideTokenizer: "LlamaTokenizer"
+    )

指定するのはLLMEval/ContentView.swiftです

- let modelConfiguration = ModelConfiguration.phi34bit
+ let modelConfiguration = ModelConfiguration.llama38bInstruct4bit

初回起動時にダウンロードがはじまります

Macでは会話ができました

※iPhone 15 Proではメモリが不足しアプリが落ちました

ModelConfigurationについて

gemma-2b-itを例にModelConfigurationの仕組みについて説明します

ModelConfigurationのイニシャライザが以下です

まずidというのは量子化済みモデルのhuggingface.co上のURLパスです

モデルは https://huggingface.co/mlx-community で公開されており、この中に該当するモデルがない場合は自分で変換してアップデートする必要があります

tokenizerIdはトークナイザーのURLです。meta-llama/Llama-2-7b-hf のように指定してダウンロードできます

overrideTokenizerはトークナイザーをダウンロードせずswift-transformersの実装クラスを使う方法です。以下に実装があります

prompt引数は各モデルに応じたテンプレートテキストを入れる場所です

Llama-3の場合は以下に公開情報がありました

これを参考にすると、こんな感じの定義になりそうです

    public static let openelm3BInstruct4bit = ModelConfiguration(
        id: "mlx-community/OpenELM-3B-instruct-4bit",
        overrideTokenizer: "LlamaTokenizer"
    )
    { prompt in
        "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\(prompt)<|eot_id|><|start_header_id|>assistant"
    }

リストにないモデルの対応は?

リストにないモデルをLLMEvalでロードするには、モデルのconfig.jsonを参考に内部パラメータ調整して読み込むためのコードを定義する必要があるようです

筆者はOpenELMを動かしたかったのですが、この対応関係の知識がなくできませんでした

ステップとしては以下のような流れになると思います

Libraries/LLM/OpenELM.swiftを定義
ModelTypeを追加
OpenELMConfigurationにconfig.jsonから受け取った値をセットアップする
LoRAModelの適応などが必要ならそこも調整

脚注

llama.cpp https://github.com/ggerganov/llama.cpp ↩︎
MLX https://ml-explore.github.io/mlx/build/html/index.html ↩︎
MLX Swift https://github.com/ml-explore/mlx-swift ↩︎

Llama 3を動かす

ModelConfigurationについて

リストにないモデルの対応は?

Discussion