MacbookでLoRAチューニングが簡単に実行できる時代になりました。自分のPCが熱を上げて、学習を回している体験は感動ものですね…。

この記事は、以下のリポジトリの LoRA に関するサンプルコードを参考にしています。結構簡単ですので、Macbookで LoRA チューニングしてみたい方は、ぜひ最後まで読んでみてください。

この記事で紹介すること

• MLXによる LoRA チューニングの実行の仕方（Mistral-7bで試します）

• カスタムデータによる LoRA チューニング

前提：動作環境

Macbook Pro、Apple M1 Maxチップ、メモリ64GBを使って検証してます。

Apple シリコンチップであれば、動作できると思いますが、メモリ不足で実行できないという場合があるかもしれません。ご了承ください。

やり方

mlx-examplesに、MLXを用いて様々なモデルを動かすためのコードが格納されています。今回はここにある LoRA に関するコードを利用していきます。

https://github.com/ml-explore/mlx-examples/tree/main/lora

環境構築

適当なディレクトリで以下を実行します。

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

続いて、LoRA のコードがあるディレクトリに移動して、環境構築をしていきます。

cd mlx-examples/lora

# .lora という名前の仮想環境を作成
python3 -m venv .lora

# 仮想環境をアクティベート
source .lora/bin/activate

# パッケージインストール
pip install -r requirements.txt

次に、モデルをダウンロードして変換します。Mistral-7b のパラメータは次のコマンドでダウンロードできます。

curl -O https://files.mistral-7b-v0-1.mistral.ai/mistral-7B-v0.1.tar
tar -xf mistral-7B-v0.1.tar

ダウンロードには少し時間がかかるので、気長に待ちましょう☕️

ダウンロードが完了したら、MLX用にパラメータを変換していきます。先ほどダウンロードして格納してあるパラメータのパスと、MLX用のパラメータを出力するパス（ここでは `mlx-mstral-7b-v0.1` ）を指定します。

python convert.py \
    --torch-model ./mistral-7B-v0.1/ \
    --mlx-model ./mlx-mistral-7B-v0.1/

LoRAチューニング

続いては、おまちかねの LoRA チューニングです！以下のコマンドを実行すると学習が開始します。（--iters 600 は結構時間かかりました。1時間くらい？）

python lora.py --model mlx-mistral-7B-v0.1 --train --iters 600

> Loading pretrained model
Total parameters 7243.436M
Trainable parameters 1.704M
Loading datasets
Training
Iter 1: Val loss 2.361, Val took 54.728s
Iter 10: Train loss 2.342, It/sec 0.336, Tokens/sec 134.829
Iter 20: Train loss 2.192, It/sec 0.299, Tokens/sec 121.071
Iter 30: Train loss 2.085, It/sec 0.268, Tokens/sec 110.425
Iter 40: Train loss 1.858, It/sec 0.313, Tokens/sec 121.175
Iter 50: Train loss 1.605, It/sec 0.295, Tokens/sec 117.277
Iter 60: Train loss 1.426, It/sec 0.298, Tokens/sec 118.012
Iter 70: Train loss 1.410, It/sec 0.214, Tokens/sec 86.716
Iter 80: Train loss 1.426, It/sec 0.301, Tokens/sec 120.817
Iter 90: Train loss 1.367, It/sec 0.303, Tokens/sec 114.800
Iter 100: Train loss 1.382, It/sec 0.281, Tokens/sec 107.002
Iter 110: Train loss 1.266, It/sec 0.314, Tokens/sec 121.551
Iter 120: Train loss 1.263, It/sec 0.302, Tokens/sec 118.015
Iter 130: Train loss 1.258, It/sec 0.289, Tokens/sec 114.392
Iter 140: Train loss 1.279, It/sec 0.137, Tokens/sec 53.242
Iter 150: Train loss 1.266, It/sec 0.266, Tokens/sec 108.363
Iter 160: Train loss 1.202, It/sec 0.294, Tokens/sec 115.531
Iter 170: Train loss 1.217, It/sec 0.305, Tokens/sec 125.783
Iter 180: Train loss 1.211, It/sec 0.311, Tokens/sec 112.980
Iter 190: Train loss 1.215, It/sec 0.304, Tokens/sec 121.622
Iter 200: Train loss 1.298, It/sec 0.312, Tokens/sec 118.595
Iter 200: Val loss 1.228, Val took 53.067s
...
Iter 550: Train loss 1.010, It/sec 0.309, Tokens/sec 121.676
Iter 560: Train loss 1.065, It/sec 0.305, Tokens/sec 121.496
Iter 570: Train loss 0.964, It/sec 0.314, Tokens/sec 121.277
Iter 580: Train loss 0.944, It/sec 0.308, Tokens/sec 119.934
Iter 590: Train loss 0.920, It/sec 0.297, Tokens/sec 117.197
Iter 600: Train loss 0.988, It/sec 0.305, Tokens/sec 125.505

ちなみにデフォルトでは data/ にある test.jsonl, train.jsonl, valid.jsonl を利用しているようで、サンプルコードには自然言語からSQLを生成するデータセットがありました。

また学習中の使用状況については、以下のような感じになっていました。GPUが結構使われていますね！PCも高揚してます🔥

評価する

テストセットのパープレキシティを計算するには、次を使用します。

python lora.py --model mlx-mistral-7B-v0.1 --adapter-file adapters.npz --test

> Loading pretrained model
Total parameters 7243.436M
Trainable parameters 1.704M
Loading datasets
Testing
Test loss 1.351, Test ppl 3.862.

出力結果について、あまりピンと来なかったのでGPT-4に解説してもらいました。

ここら辺は、イテレーション回数を対照実験的に変化させて見ていくのが良さそうですね。

実際に生成してみる

LoRAチューニングしたモデルを早速使ってみましょう。以下のようにモデルとアダプタファイルを指定し、先ほど学習させたデータセットの形式でプロンプトも記述します。

python lora.py --model mlx-mistral-7B-v0.1 \
               --adapter-file adapters.npz \
               --num-tokens 50 \
               --prompt "table: 1-10015132-16
columns: Player, No., Nationality, Position, Years in Toronto, School/Club Team
Q: What is terrence ross' nationality
A: "

> Loading pretrained model
Total parameters 7243.436M
Trainable parameters 1.704M
Loading datasets
Generating
table: 1-10015132-16
columns: Player, No., Nationality, Position, Years in Toronto, School/Club Team
Q: What is terrence ross' nationality
A: SELECT Position FROM 1-10015132-16 WHERE Nationality = 'Terrence Ross'9 SELECT No. FROM 1-10015132-16 WHERE National

良い感じですね！最後の方でループしてしまっているのは、--num-tokens で指定されている50トークン分が生成されてしまうからですかね。

モデルが回答する際に、末尾に必ず終端トークン（End of Sentence）的なものを入れられたら、システムに組み込むなどできそうな気がします。

カスタムデータセットでLoRA

というわけで課題も見えたところで、今度は終端トークンも生成してもらえるようにしたいと思います。文字列はなんでもいいのですが、生成結果の末尾に <EOS> を出力してもらうようにしてみましょう。

data/ を複製して、custom-data/ を作成し、一括置換を使ってjsonlの各データセットの文字列末尾に <EOS> を追加します。

そして以下のコマンドを実行します。

python lora.py --model mlx-mistral-7B-v0.1 --train --iters 600 \
               --data custom-data --adapter-file custom-adapters.npz

完了したら、実際に生成結果がどう変化するかみてみます。

python lora.py --model mlx-mistral-7B-v0.1 \
               --adapter-file custom-adapters.npz \
               --num-tokens 50 \
               --prompt "table: 1-10015132-16
columns: Player, No., Nationality, Position, Years in Toronto, School/Club Team
Q: What is terrence ross' nationality
A: "

> Loading pretrained model
Total parameters 7243.436M
Trainable parameters 1.704M
Loading datasets
Generating
table: 1-10015132-16
columns: Player, No., Nationality, Position, Years in Toronto, School/Club Team
Q: What is terrence ross' nationality
A: SELECT No. FROM 1-10015132-16 WHERE Nationality = 'Terrence Ross'<EOS>

やりました！ちゃんと <EOS> が末尾についています。これでシステム的にもパースがしやすくなり、扱えるレベルになった気がしますね。

おまけ：メモリが足りない場合

READMEには次のような記述もありました。もしメモリが不足している際はお試しください。

おわりに

Mistral-7b を LoRAチューニングしてみました。想像以上に結構サクサクできてびっくりです。カスタムデータセットの学習も、意図した通りの挙動を得られて、とても良きでした👏

今回は、mlx-examplesにあるサンプルコードと手順に沿って行う形でしたが、Mistral-7b 以外のモデルも LoRA チューニングできるのかについては、今後検証していきたいと思います！