MacでLLMをファインチューニング｜mlx-lmでLoRAを試す

Apple Silicon搭載のMacでは、mlx-lm（MLX）を活用することで、手軽にローカルLLMのファインチューニングが可能です。本記事では、LoRA（Low-Rank Adaptation） を用いた軽量ファインチューニングの手法を、Mac環境で実践する方法を解説します。

MLXとは

MLXはAppleの提供するApple Silicon向けの機械学習用のフレームワークです。ここでは、MLXを用いて大規模言語モデル（LLM）を利用するためのユーティリティ MLX-LM を使って、ファインチューニングを行いたいと思います。

MLX-LMには、mlx_lm.generate, mlx_lm.loraといったコマンドラインツールが提供されていてこれを使うことでプログラミングなしでファインチューニング可能です。

インストール

今回は、huggingfaceのデータセットとmlx_lmを利用するので以下の２つをインストールします。

pip install mlx-lm datasets

venvやcondaなどで仮想環境を作ってから試すことをお勧めします

condaを使う場合

コマンド早見表(Conda/docker/docer-compose)

uvを使う場合

Macで開発環境関連の初期設定する手順（VSCode, brew, iterm2, python, go）

データセットの準備

学習にあたってデータセットを準備する必要があります。今回はhuggingfaceのデータセットを使いますが、mlx_lm.loraで使えるデータセットのフォーマットについて解説しておきます。

データセットのフォーマット

準備するデータは、train.jsonl, valid.jsonl test.jsonlの３つのファイルになります。これらを同じフォルダ（例えばdatasetsなど）に格納しておく必要があります。

test.jsonlはなくても一応動きます

公式ページを見るとフォーマットは以下のようになっています。チャットの場合は、system, user, assistantの3つに対応するテキストが必要なようです。試した感じではinstructが付いたモデルはこちらのフォーマットで学習した方が良さげでした。

{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello."}, {"role": "assistant", "content": "How can I assistant you today."}]}

下記は、文章を与える場合のフォーマットです。

{"text": "This is an example for the model."}

データセットでは、上記のフォーマットで１行を作り、これが１つの学習データとなります。各ファイルには、これを複数行並べます。例えば、textが1000個ある場合は、{"text":...}を1000行並べます。

HuggingFaceのデータセットの利用方法

mlx_lmを使った学習では、huggingfaceのデータセットを使った学習も可能です。huggingfaceの学習を行う場合は、後で説明するconfig.yamlに以下のフォーマットで設定を書きます

hf_dataset:
  name: "Open-Orca/OpenOrca"
  train_split: "train[:90%]"
  valid_split: "train[-10%:]"
  prompt_feature: "question"
  completion_feature: "response"

nameはデータセット名です。ローカルに存在しない場合は自動的にダウンロードします。

train_split, valid_split, test_splitを使って、データセット中の項目を学習、検証、テスト用データに分割することができます。

hf_dataset:
  name: "billsum"
  prompt_feature: "text"
  completion_feature: "summary"

テキストの例はシンプルな例です。billsumデータセットは、train, valid, testのデータセットがあらかじめ分割されているのでtrain_splitなどは設定していません。

データセットの中身はhuggingfaceで確認してください

今回利用するデータセット

今回はLoRAの効果がわかりやすそうなdatabricks-dolly-15k-ja-gozaruを利用してみます。このデータセットはQ＆A形式のデータで、回答の末尾に「ござる」がついているのが特徴です。

データセットを見るとcategory, instruction, input, output, indexの列がありますが、今回はinstructionとoutputを利用して学習させることにします。この場合のconfig.yamlの設定は以下になります。データセットにはtrainしかないので、これを80%:10%:10%で分割してtrain, valid, testにしています。

hf_dataset:
 name: "bbz662bbz/databricks-dolly-15k-ja-gozaru"
 train_split: "train[:80%]"
 valid_split: "train[80%:90%]"
 test_split: "train[-10%:]"
 prompt_feature: "instruction"
 completion_feature: "output"

mlx_lm.loraを使ったファインチューニング

mlx_lm.loraでは、コマンドラインオプションを使って各種設定することも可能ですが、ここでは、config.yamlを使って一気に設定する方法を取りました。

学習に利用した環境

学習に利用した環境は以下になります。

• MacBook Pro M4Max 128GB

16コアCPU、40コアGPUです（学習中はGPUがフル稼働していました）

学習は大量のメモリを必要とするのでなるべく大きなものを準備した方が良いです。また、小さなモデル（0.5Bや1Bモデル）で試すことをお勧めします。

config.yamlによる設定

以下config.yamlの設定です。MLX形式のモデルであればなんでもよかったのですが、今回は、mlx-community/Mistral-7B-Instruct-v0.2-4bit-mlxを利用しました。

MLXでないモデルもコンバートして利用することが可能です（mlx_lm.convert）

また、fine_tine_typeはloraにしています。

# The path to the local model directory or Hugging Face repo.
model: "mlx-community/Mistral-7B-Instruct-v0.2-4bit-mlx"

# Whether or not to train (boolean)
train: true

# The fine-tuning method: "lora", "dora", or "full".
fine_tune_type: lora

# Directory with {train, valid, test}.jsonl files
data: "bbz662bbz/databricks-dolly-15k-ja-gozaru"

# The PRNG seed
seed: 42

# Number of layers to fine-tune
num_layers: 4

# Minibatch size.
batch_size: 1

# Iterations to train for.
iters: 1000

# Number of validation batches, -1 uses the entire validation set.
val_batches: 2

# Adam learning rate.
learning_rate: 1e-5

# Number of training steps between loss reporting.
steps_per_report: 10

# Number of training steps between validations.
steps_per_eval: 100

# Load path to resume training with the given adapter weights.
resume_adapter_file: null

# Save/load path for the trained adapter weights.
adapter_path: "adapter"

# Save the model every N iterations.
save_every: 100

# Evaluate on the test set after training
test: ture

# Number of test set batches, -1 uses the entire test set.
test_batches: 2

# Maximum sequence length.
max_seq_length: 2048

# Use gradient checkpointing to reduce memory use.
grad_checkpoint: false

# LoRA parameters can only be specified in a config file
lora_parameters:
  # The layer keys to apply LoRA to.
  # These will be applied for the last lora_layers
  # keys: ["self_attn.q_proj", "self_attn.v_proj"]
  keys: ["self_attn.v_proj"]
  rank: 8
  scale: 16.0
  dropout: 0.0

# Schedule can only be specified in a config file, uncomment to use.
#lr_schedule:
#  name: cosine_decay
#  warmup: 100 # 0 for no warmup
#  warmup_init: 1e-7 # 0 if not specified
#  arguments: [1e-5, 1000, 1e-7] # passed to scheduler

hf_dataset:
 name: "bbz662bbz/databricks-dolly-15k-ja-gozaru"
 train_split: "train[:80%]"
 valid_split: "train[80%:90%]"
 test_split: "train[-10%:]"
 prompt_feature: "instruction"
 completion_feature: "output"  

以下設定のポイントです

気になった設定部分をChatGPTなどに問い合わせながらまとめています

LoRAのランク設定

lora_parametersのrank設定についてです。以下のような設定が目安だそうです。今回は7Bモデルなのでrank=8にしました。

なお、scaleは目安はrank*2だそうです。

LoRAのレイヤ数設定

num_layersは最終層から何層を再学習させるかの設定です。目安は以下のようになるそうです。今回は語尾のござるだけ覚えさせたいのでlayers=2~6ということで設定しました。

対象とするレイヤー

LoRAを行う場合、対象とする層としては、以下の３つがあります。

• self_attn.q_proj
• self_attn.v_proj
• self_attn.k_proj

用途によってどの部分を学習対象にすれば良いか変わるようです。

上記は、あくまでも目安なので、実際には試行錯誤が必要です。今回もある程度試行錯誤しました。

学習の実行

学習には以下のコマンドを実行します（yamlファイルはカレントフォルダにある前提です）

mlx_lm.lora --config config.yaml

学習が進むと、以下のような出力が行われます。ちなみに、学習ですが、省電力モードにしていたので結構かかりました（15〜20分くらい）。

学習データはtrain, valid, test合計で15,000行くらいありますので、割と多い方だと思います。これを見ると1回の学習が1秒弱（1.0 It/se）と、かなり高速に処理できていることがわかります。

なお、メモリはPeak mem 24.729 GBとなっていました。rankを上げたり、num_layersを増やすとメモリ消費量は多くなります。また、モデルを4bitから8bitにしても大きくなります。学習には結構なメモリが必要なようです。

色々なパラメータで試した結果、最大で100GB程度まで上昇しました。学習する場合はMacのメモリは64GB程度は最低でも必要かもしれません。

Loading configuration file config.yaml
Loading pretrained model
Fetching 5 files: 100%|████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 22215.59it/s]
Loading datasets
Loading Hugging Face dataset bbz662bbz/databricks-dolly-15k-ja-gozaru.
Training
Trainable parameters: 0.018% (1.311M/7241.732M)
Starting training..., iters: 1000
Iter 1: Val loss 1.915, Val took 4.712s
Iter 10: Train loss 2.664, Learning Rate 1.000e-05, It/sec 0.502, Tokens/sec 186.015, Trained Tokens 3705, Peak mem 13.574 GB
Iter 20: Train loss 2.352, Learning Rate 1.000e-05, It/sec 0.875, Tokens/sec 181.598, Trained Tokens 5781, Peak mem 13.574 GB
Iter 30: Train loss 2.075, Learning Rate 1.000e-05, It/sec 1.380, Tokens/sec 172.615, Trained Tokens 7032, Peak mem 13.574 GB
Iter 40: Train loss 1.735, Learning Rate 1.000e-05, It/sec 0.909, Tokens/sec 181.033, Trained Tokens 9023, Peak mem 13.574 GB
Iter 50: Train loss 1.762, Learning Rate 1.000e-05, It/sec 0.786, Tokens/sec 185.421, Trained Tokens 11383, Peak mem 13.574 GB
Iter 60: Train loss 1.683, Learning Rate 1.000e-05, It/sec 1.164, Tokens/sec 179.083, Trained Tokens 12921, Peak mem 13.574 GB
Iter 70: Train loss 1.604, Learning Rate 1.000e-05, It/sec 0.652, Tokens/sec 180.608, Trained Tokens 15692, Peak mem 13.574 GB
Iter 80: Train loss 1.684, Learning Rate 1.000e-05, It/sec 1.279, Tokens/sec 172.320, Trained Tokens 17039, Peak mem 13.574 GB
Iter 90: Train loss 1.661, Learning Rate 1.000e-05, It/sec 0.537, Tokens/sec 172.984, Trained Tokens 20260, Peak mem 20.648 GB
Iter 100: Val loss 1.396, Val took 0.654s
Iter 100: Train loss 1.734, Learning Rate 1.000e-05, It/sec 12.428, Tokens/sec 3717.218, Trained Tokens 23251, Peak mem 20.648 GB
Iter 100: Saved adapter weights to adapter/adapters.safetensors and adapter/0000100_adapters.safetensors.
 :
 :
（略）
 :
 :
Iter 970: Train loss 1.533, Learning Rate 1.000e-05, It/sec 1.031, Tokens/sec 180.512, Trained Tokens 221270, Peak mem 24.729 GB
Iter 980: Train loss 1.390, Learning Rate 1.000e-05, It/sec 0.797, Tokens/sec 184.951, Trained Tokens 223591, Peak mem 24.729 GB
Iter 990: Train loss 1.363, Learning Rate 1.000e-05, It/sec 1.050, Tokens/sec 176.123, Trained Tokens 225269, Peak mem 24.729 GB
Iter 1000: Val loss 1.911, Val took 1.194s
Iter 1000: Train loss 1.783, Learning Rate 1.000e-05, It/sec 12.341, Tokens/sec 3667.604, Trained Tokens 228241, Peak mem 24.729 GB
Iter 1000: Saved adapter weights to adapter/adapters.safetensors and adapter/0001000_adapters.safetensors.
Saved final weights to adapter/adapters.safetensors.
Testing
Test loss 1.615, Test ppl 5.026.

Testの結果が末尾にありますが、ppl（perplexity）が5程度ということで、一応「次の単語の候補を平均 5 個程度に絞れている」ということになります。まぁまぁ学習できたように見えます。

mlx_lm.generateを使って確認

mlx_lm.generateコマンドを使ってファインチューニングの効果を確認してみます。

ファインチューニング前

まずは、ファインチューニング前の回答です。

mlx-test % mlx_lm.generate --model mlx-community/Mistral-7B-Instruct-v0.2-4bit-mlx --temp 0.9  --prompt "日本で一番大きな山は？" 

ファインチューニング前は上記のような回答でした。少し最後の方が壊れています。また、小さいモデルではありがちですが「東ストッカートルク山やピレーヌ山」というちょっとわからない回答になっています。

おそらく、Mistral-7Bは日本語での学習が不十分なんだとおもいます

ファインチューニング後

ファインチューニング後の回答です。loraの学習結果を使う場合は、引数--adapter-pathに学習したloraモデルのフォルダを指定します（ここでは、adapter）。

mlx_lm.generate --model mlx-community/Mistral-7B-Instruct-v0.2-4bit-mlx --temp 0.9 --adapter-path adapter --seed 42 --prompt  "日本で一番大きな山は？"

結果を見ると、末尾に「ござる」がちゃんとついて回答されました。

いろいろなプロンプトで試しましたが、うまく行っている場合とダメな場合が混在していました。やはりLoRAによる学習は結構難しい印象です

まとめ

Macでファインチューニングを行ってみました。思ったより手軽でしたが、メモリがかなり必要になります（4bit量子化の7Bモデルの学習でも48GBほど必要でした）。まだまだ、ファインチューニングは手軽とまではいかない印象です。

ファインチューニングできれば、例えば、特定の分野に特化したLLMなどを作ることができるのでローカルLLMの活用の範囲が広がります。

ある/Aru

IT＆機械学習エンジニア/ファイナンシャルプランナー（CFP®)

XContact

専門分野は並列処理・画像処理・機械学習・ディープラーニング。プログラミング言語はC, C++, Go, Pythonを中心として色々利用。現在は、Kaggle, 競プロなどをしながら悠々自適に活動中