メモリ32GBのMacでLLMファインチューニングを試す(MLX-LM LoRA)

LLMの独自学習のニーズが高まってきていることもあり、感触をつかむためにMacで試してみることにしました

まずLLMの独自学習についてはこちらの記事がよくまとまっているため、簡単に(安価に)できるだろうと思っているかたは是非お読みください
LLM学習の現実:GPU選びから学習コストまで徹底解説

ということで、手軽にできるのはモデルの学習ではなく、既存モデルのファインチューニングということになります
ただ、これもかなりのGPUメモリを必要とします

普段はNVIDIAのcuda環境でいろいろやっており、GPUメモリが128GBというDGX Sparkを買う直前まで行きました
ただ、DGXシリーズはUbuntuベースとは言えDGX OSというパッケージでドライバ類も含めた形で提供されており、NVIDIA社がアップデートをやめたら使いどころのない箱となります
もちろんUbuntuをインストールしてドライバ類を独自に探してということも可能ですが、コスト的にちょっとどうなの…と個人的には思うところがありまして…

そこで思いついたのは、速度さえ許せるならユニファイドメモリとしてGPUもメモリを共用できるMacを使うのは有りでは?という発想です
Mac Studioならメモリ128GBのモデルもあり、DGX Sparkぐらいの価格で買うことができます
しかもMacなら他の用途への転用も可能です

とはいえそれなりの価格ではあるので、本当に感触をつかむためだけに手元にあるメモリ32GBのMacでさわりの部分だけでも試すことができないものかと調べました
すると、7Bのモデルを4bit量子化したものであれば扱えそうということがわかります

うまく動かなくても、とりあえずちょっとやってみようということで進めてみることにしました

MacでPythonを使った検証を行うための環境構築

LLMとは直接関係ない内容も含みますが、以下、環境構築としてやったことを記載しておきます

リモートから操作する際の日本語文字化け対策

外部ホストからssh接続する場合の日本語の文字化けを防ぐため、環境変数LANGを設定します
ssh接続を有効にする方法はここで詳しく触れませんが、共有設定でsshを許可する設定を入れておく必要があります

# .zprofileを編集します(viではなくGUIから操作しても問題ありません)
vi ~/.zprofile

# 以下の1行を入れて保存、終了します
export LANG=ja_JP.UTF-8

今のセッションには反映されないので、ログインをし直しておきます

source ~/.zprofile

でも良いかもしれません

uvを使ったPython仮想環境の構築

uvをインストールします
brewでも入るようですが、公式のスクリプトを使用しています

curl -LsSf https://astral.sh/uv/install.sh | sh

プロジェクトを作成します
自分の環境ではhomeディレクトリ配下にprojectsディレクトリを作っていますので、その下にプロジェクトを作っていきます

cd ~/projects
uv init mlx-test # ここでmlx-testディレクトリが作られ、pyproject.tomlなど必要なファイルも生成される

cd mlx-test
uv venv --python 3.12 # ここでpython3.12を使った.venvが作られる

source .venv/bin/activate # venvのアクティベート(念のため、抜けるときは deactivate です)
uv python pin 3.12 # 一応3.12に固定しておく
uv add mlx-lm datasets # 必要なライブラリをインストールする

ファインチューニングしてみる

まずは感触をつかみたいだけだったので、ネット上で見つけたこちらの記事を参考に同じことをしてみました
MacでLLMをファインチューニング|mlx-lmでLoRAを試す

記事を書かれているかたはメモリ128GBの環境のようですが、この内容なら32GBでも動くだろうという感覚があったのでそのまま試してみました
いくつか書かれている内容と変えた部分があり、それらの部分を記載しておきます

まずはこのままのconfigでと思ったのですが、

File "/Users/jun/projects/mlx-test/.venv/lib/python3.12/site-packages/mlx_lm/tuner/datasets.py", line 256, in load_custom_hf_dataset
    ds_path = ds["path"]
              ~~^^^^^^^^
KeyError: 'path'

とのエラーになります

configに設定するデータセットのnameというのがpathに変わっているようで、ネット上にたくさんある過去の記事はnameを使ったものが多いようです
エラーメッセージがわかりにくいですが、データセットそのものの問題ではないためconfigを修正します

hf_dataset:
 name: "bbz662bbz/databricks-dolly-15k-ja-gozaru"

hf_dataset:
 path: "bbz662bbz/databricks-dolly-15k-ja-gozaru"

に変更します

また、モデルは元の記事で使用されている mlx-community/Mistral-7B-Instruct-v0.2-4bit-mlx というのがなくなっており、 mlx-community/Mistral-7B-Instruct-v0.3-4bit を使用しました

使用したMacのスペックですが、メモリは32GBと書いているとおりで、プロセッサはM2 ProのMac miniです
実行時のスクリーンショットはこんな感じで、1000 iterで約20分弱ぐらいでした
メモリ32GBのMacでMLX-LMを使ったファインチューニング時のログ

学習時のメモリ使用状況はこんな感じです
しっかり使いきりぐらいの感じですが、スワップは発生していないようです
docker等環境的に入っていますが、今回は使用していません
メモリ32GBのMacでMLX-LMを使ったファインチューニング時のメモリ使用状況

1000 iter回したときの結果がこちら
メモリ32GBのMacでMLX-LMを使ったファインチューニング 1000iterの結果

ちょっと微妙な感じがしたので、2000 iterまで回してみました
メモリ32GBのMacでMLX-LMを使ったファインチューニング 2000iterの結果

メモリ32GBのMacでLLMファインチューニングを試す(MLX-LM LoRA)にコメントする

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です