Apple silicon専用機械学習フレームワークでLLMのファインチューニングをやってみた

# 概要

Apple silicon専用の機械学習用フレームワークである [MLX](https://github.com/ml-explore/mlx) と、MLXを用いた大規模言語モデル（LLM）を利用するためのユーティリティ群 [MLX-LM](https://github.com/ml-explore/mlx-examples/tree/main/llms/mlx_lm) を使って、
[llm-jp/llm-jp-3-1.8b](https://huggingface.co/llm-jp/llm-jp-3-1.8b) のファインチューニングをやってみました。  
学習に使ったのは、円城塔さん小説データセットです。

# データセットについて

芥川賞作家の円城塔さんは、[カクヨムで公開している小説](https://kakuyomu.jp/users/enjoetoh) を、機械学習に利用することを想定してGitHub でも公開してくれています。  
* [「通信記録保管所」に相当するデータセット](https://github.com/EnJoeToh/stories_2000): 1ファイル2000字程度で、1作品をなす。
* [「花とスキャナー」に相当するデータセット](https://github.com/EnJoeToh/stories_10000): 4ファイル1万字程度で、1作品をなす。

いずれも [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/deed.ja) での公開であり、さらに
> 個人的な読書、非営利での朗読
> 機械学習のデータなどに利用することができます

とのことです。

# 必要なライブラリのインストール

`uv` を使ってインストールしました。
※ 依存ライブラリの一つ `sentencepiece=0.2.0` が、Python 3.13 ではインストールに失敗するようです。
```terminal
$ python --version
Python 3.12.7
$ uv add mlx-lm datasets
```

# データセットのダウンロードと整形

まず、文章をGitHubからローカルにクローンします。
今回は、「通信記録保管所」の100作品を使います。
```terminal
$ git clone https://github.com/EnJoeToh/stories_2000.git
```
ローカルに `stories_2000` というフォルダができます。
Hugging Face の [Datasets](https://huggingface.co/docs/datasets/index) ライブラリを使うことで、学習しやすい形式に整えることができます。  
今回は空行ごとに分割して、学習用データセットに整えます。

```python:split_paragraph.py
from datasets import load_dataset

# sample_by パラメータに paragraph を指定することで、
# 空行ごとに一つのデータになる
dataset_enjoe = load_dataset("text", data_files="stories_2000/0*.txt", sample_by="paragraph")

# 9:1 で、バリデーション用のデータを取り分けておく。
dataset_enjoe = dataset_enjoe['train'].train_test_split(test_size=0.1, shuffle=False)

# jsonl 形式で保存
dataset_enjoe["train"].to_json("stories_2000_sample_by_paragraph/train.jsonl")
dataset_enjoe["test"].to_json("stories_2000_sample_by_paragraph/test.jsonl")

print(dataset_enjoe.shape)  # {'train': (378, 1), 'test': (43, 1)}
```
```terminal
$ python split_paragraph.py
```
train.jsonl が 378データ、test.jsonl が 43データになりました。

# mlx_lm.lora による学習

MLX_LM をインストールすると使えるようになる `mlx_lm.lora` というCLIツールが便利です。

```terminal
$ mlx_lm.lora --help                            
usage: mlx_lm.lora [-h] [--model MODEL] [--train] [--data DATA] [--fine-tune-type {lora,dora,full}]
                   [--num-layers NUM_LAYERS] [--batch-size BATCH_SIZE] [--iters ITERS]
                   [--val-batches VAL_BATCHES] [--learning-rate LEARNING_RATE]
                   [--steps-per-report STEPS_PER_REPORT] [--steps-per-eval STEPS_PER_EVAL]
                   [--resume-adapter-file RESUME_ADAPTER_FILE] [--adapter-path ADAPTER_PATH]
                   [--save-every SAVE_EVERY] [--test] [--test-batches TEST_BATCHES]
                   [--max-seq-length MAX_SEQ_LENGTH] [-c CONFIG] [--grad-checkpoint] [--seed SEED]

LoRA or QLoRA finetuning.

options:
  -h, --help            show this help message and exit
  --model MODEL         The path to the local model directory or Hugging Face repo.
  --train               Do training
  --data DATA           Directory with {train, valid, test}.jsonl files or the name of a Hugging Face dataset
                        (e.g., 'mlx-community/wikisql')
  --fine-tune-type {lora,dora,full}
                        Type of fine-tuning to perform: lora, dora, or full.
  --num-layers NUM_LAYERS
                        Number of layers to fine-tune. Default is 16, use -1 for all.
  --batch-size BATCH_SIZE
                        Minibatch size.
  --iters ITERS         Iterations to train for.
  --val-batches VAL_BATCHES
                        Number of validation batches, -1 uses the entire validation set.
  --learning-rate LEARNING_RATE
                        Adam learning rate.
  --steps-per-report STEPS_PER_REPORT
                        Number of training steps between loss reporting.
  --steps-per-eval STEPS_PER_EVAL
                        Number of training steps between validations.
  --resume-adapter-file RESUME_ADAPTER_FILE
                        Load path to resume training from the given fine-tuned weights.
  --adapter-path ADAPTER_PATH
                        Save/load path for the fine-tuned weights.
  --save-every SAVE_EVERY
                        Save the model every N iterations.
  --test                Evaluate on the test set after training
  --test-batches TEST_BATCHES
                        Number of test set batches, -1 uses the entire test set.
  --max-seq-length MAX_SEQ_LENGTH
                        Maximum sequence length.
  -c CONFIG, --config CONFIG
                        A YAML configuration file with the training options
  --grad-checkpoint     Use gradient checkpointing to reduce memory use.
  --seed SEED           The PRNG seed
```

学習用データセットとして、 `{train, valid,test}.jsonl` が必要です。
先ほど `train.jsonl` と `test.jsonl` は作ったので、`valid.jsonl` として `test.jsonl` をコピーします。

```terminal
$ cp stories_2000_sample_by_paragraph/test.jsonl stories_2000_sample_by_paragraph/valid.jsonl
```

学習モードである `--train` の指定と、ファインチューニングしたいモデル名を示す `--model` パラメータ、
先ほど用意したデータセットのパスを示す `--data` パラメータ、
イテレーションの回数示す `--iters` を指定すればほぼ十分です。  
今回はこれに加えて、デフォルトのミニバッチサイズが `4` であり、大きいモデルを使うとメモリが厳しいので、`--batch-size 1` と減らし、  
再現性のために `--seed 42` を指定しました。  
結果の保存先を指定する `--adapter-path` も指定しています。
```terminal
$ mlx_lm.lora --model llm-jp/llm-jp-3-1.8b \
--train --data stories_2000_sample_by_paragraph \
--batch-size 1 \
--iters 100 \
--adapter-path it100_seed42 \
--seed 42
```

モデルを少し大きな `llm-jp/llm-jp-3-3.7b` に変更したり、
イテレーションの回数を増やすと、性能が上がるかもしれません。

# mlx_lm.generate による文書生成

MLX_LM は文書生成についても、 `mlx_lm.generate` という手軽なCLIツールを意してくれています。

```terminal
$ mlx_lm.generate --help
usage: mlx_lm.generate [-h] [--model MODEL] [--adapter-path ADAPTER_PATH] [--trust-remote-code]
                       [--eos-token EOS_TOKEN] [--prompt PROMPT] [--max-tokens MAX_TOKENS] [--temp TEMP]
                       [--top-p TOP_P] [--seed SEED] [--ignore-chat-template] [--use-default-chat-template]
                       [--verbose VERBOSE] [--colorize] [--max-kv-size MAX_KV_SIZE]
                       [--prompt-cache-file PROMPT_CACHE_FILE] [--kv-bits KV_BITS]
                       [--kv-group-size KV_GROUP_SIZE] [--quantized-kv-start QUANTIZED_KV_START]

LLM inference script

options:
  -h, --help            show this help message and exit
  --model MODEL         The path to the local model directory or Hugging Face repo. If no model is specified,
                        then mlx-community/Llama-3.2-3B-Instruct-4bit is used.
  --adapter-path ADAPTER_PATH
                        Optional path for the trained adapter weights and config.
  --trust-remote-code   Enable trusting remote code for tokenizer
  --eos-token EOS_TOKEN
                        End of sequence token for tokenizer
  --prompt PROMPT       Message to be processed by the model ('-' reads from stdin)
  --max-tokens MAX_TOKENS, -m MAX_TOKENS
                        Maximum number of tokens to generate
  --temp TEMP           Sampling temperature
  --top-p TOP_P         Sampling top-p
  --seed SEED           PRNG seed
  --ignore-chat-template
                        Use the raw prompt without the tokenizer's chat template.
  --use-default-chat-template
                        Use the default chat template
  --verbose VERBOSE     Log verbose output when 'True' or 'T' or only print the response when 'False' or 'F'
  --colorize            Colorize output based on T[0] probability
  --max-kv-size MAX_KV_SIZE
                        Set the maximum key-value cache size
  --prompt-cache-file PROMPT_CACHE_FILE
                        A file containing saved KV caches to avoid recomputing them
  --kv-bits KV_BITS     Number of bits for KV cache quantization. Defaults to no quantization.
  --kv-group-size KV_GROUP_SIZE
                        Group size for KV cache quantization.
  --quantized-kv-start QUANTIZED_KV_START
                        When --kv-bits is set, start quantizing the KV cache from this step onwards.
```

`--model` パラメータに、学習時と同じ `llm-jp/llm-jp-3-1.8b` を、
`--adapter-path` パラメータに、先ほど保存先とした `it100_seed42` を、それぞれ指定します。
`--temp` パラメータが小さいと、同じ文字列の繰り返しが現れやすいので、 `0.9` を指定しました。
`--prompt` パラメータに示した文字列に続く文字列を生成してくれます。今回は、学習に使っていない[花とスキャナー](https://kakuyomu.jp/works/16817139558560459035)の冒頭を使ってみました。

```terminal
$ mlx_lm.generate --model llm-jp/llm-jp-3-1.8b --temp 0.9 --adapter-path it100_seed42 --seed 42 --prompt "　花が見えた。             
　コンクリートの壁をジグザグに駆け登る非常階段の中ほどに、"
Fetching 6 files: 100%|████████████████████████████████████████████████████████| 6/6 [00:00<00:00, 73156.47it/s]
==========
Prompt: 　花が見えた。
　コンクリートの壁をジグザグに駆け登る非常階段の中ほどに、
ぽつんと一輪。
　その色は何にも似ていない。
　この種の「花」の色を言うなら、数は少ないはずだ。
　だが、ここには季節ごとに、雑草のように同じ花が生えてくる。
　紫の花。
==========
Prompt: 24 tokens, 299.467 tokens-per-sec
Generation: 55 tokens, 27.743 tokens-per-sec
Peak memory: 3.539 GB
```

`--adapter-path` パラメータを消すと、ファインチューニングをしない状態の出力を見ることができます。

```terminal:（参考）ファインチューニングをしない場合の出力
$ mlx_lm.generate --model llm-jp/llm-jp-3-1.8b --temp 0.9 --seed 42 --prompt "　花が見え た。            
　コンクリートの壁をジグザグに駆け登る非常階段の中ほどに、"
Fetching 6 files: 100%|████████████████████████████████████████████████████████| 6/6 [00:00<00:00, 72944.42it/s]
==========
Prompt: 　花が見えた。
　コンクリートの壁をジグザグに駆け登る非常階段の中ほどに、

鉄製の手摺に花が揺れている。
それはまるで、アザミの花が咲き誇る草原から漂ってくる匂いのようだ、
奈津美はそう思いながら手摺を登り、花のほうに顔を向けた。
ふいに、下のほうにヘリがやってきた。
それと共に、徐々に下のほうに視線を移していく。
そして、下のほうにどんどん見えてくる、下のほうに視認できる自動車の数。
上のほうでも止まっていた。
==========
Prompt: 24 tokens, 310.004 tokens-per-sec
Generation: 100 tokens, 28.880 tokens-per-sec
Peak memory: 3.534 GB
```

この例では、ファインチューンをした場合の方が、より自然な小説らしい文章になっているように思えます。

# 補足（2024-11-24）

`mlx_lm.convert` を使うと、Hugging Faceで公開されているモデルを簡単に quantize することができます。
[llm-jp/llm-jp-3-13b](https://huggingface.co/llm-jp/llm-jp-3-13b) で試してみたら、だいたい1/4強、28%くらいになりました。
メモリ32GiBのM4 Macmini だと、llm-jp-3-13bはメモリ不足で動きませんでしたが、quantize したモデルは動きました。
LoRAのベースモデルとして、 quantize したモデルを使うこともできます。

```terminal
$ mlx_lm.convert -q --hf-path llm-jp/llm-jp-3-13b --mlx-path llm-jp-3-13b-q
$ du -hs ~/.cache/huggingface/hub/models--llm-jp--llm-jp-3-13b llm-jp-3-13b-q
 26G	~/.cache/huggingface/hub/models--llm-jp--llm-jp-3-13b
7.2G	llm-jp-3-13b-q
```


# 参考リンク

* [MLXの公式ドキュメント](https://ml-explore.github.io/mlx/build/html/index.html)
* [MLX-LMでRoLAを行う際のReadme](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/LORA.md) 
* [LLM勉強会](https://llm-jp.nii.ac.jp/) オープンソースかつ日本語に強いLLMの構築を行っている団体。
* [Datasets ライブラリのテキストデータのロードに関する公式ドキュメント](https://huggingface.co/docs/datasets/nlp_load)  
  `sample_by` パラメータについて言及している。
* [円城塔を近似する](https://colab.research.google.com/drive/1oXxBIYJvvUYsVZP6WYAUCb3QK09zTJtO?usp=sharing) 
  僕が2023年11月ごろに書いた、Transformerライブラリを使ったファインチューニングの Colab ノートブック