概要
Apple silicon専用の機械学習用フレームワークである MLX と、MLXを用いた大規模言語モデル（LLM）を利用するためのユーティリティ群 MLX-LM を使って、
llm-jp/llm-jp-3-1.8b のファインチューニングをやってみました。
学習に使ったのは、円城塔さん小説データセットです。
データセットについて
芥川賞作家の円城塔さんは、カクヨムで公開している小説 を、機械学習に利用することを想定してGitHub でも公開してくれています。
- 「通信記録保管所」に相当するデータセット: 1ファイル2000字程度で、1作品をなす。
- 「花とスキャナー」に相当するデータセット: 4ファイル1万字程度で、1作品をなす。
いずれも CC BY-NC 4.0 での公開であり、さらに
個人的な読書、非営利での朗読
機械学習のデータなどに利用することができます
とのことです。
必要なライブラリのインストール
uv を使ってインストールしました。
※ 依存ライブラリの一つ sentencepiece=0.2.0 が、Python 3.13 ではインストールに失敗するようです。
$ python --version
Python 3.12.7
$ uv add mlx-lm datasets
データセットのダウンロードと整形
まず、文章をGitHubからローカルにクローンします。
今回は、「通信記録保管所」の100作品を使います。
$ git clone https://github.com/EnJoeToh/stories_2000.git
ローカルに stories_2000 というフォルダができます。
Hugging Face の Datasets ライブラリを使うことで、学習しやすい形式に整えることができます。
今回は空行ごとに分割して、学習用データセットに整えます。
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)}
$ python split_paragraph.py
train.jsonl が 378データ、test.jsonl が 43データになりました。
mlx_lm.lora による学習
MLX_LM をインストールすると使えるようになる mlx_lm.lora というCLIツールが便利です。
$ 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 をコピーします。
$ 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 も指定しています。
$ 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ツールを意してくれています。
$ 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 パラメータに示した文字列に続く文字列を生成してくれます。今回は、学習に使っていない花とスキャナーの冒頭を使ってみました。
$ 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 パラメータを消すと、ファインチューニングをしない状態の出力を見ることができます。
$ 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 で試してみたら、だいたい1/4強、28%くらいになりました。
メモリ32GiBのM4 Macmini だと、llm-jp-3-13bはメモリ不足で動きませんでしたが、quantize したモデルは動きました。
LoRAのベースモデルとして、 quantize したモデルを使うこともできます。
$ 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の公式ドキュメント
- MLX-LMでRoLAを行う際のReadme
- LLM勉強会 オープンソースかつ日本語に強いLLMの構築を行っている団体。
-
Datasets ライブラリのテキストデータのロードに関する公式ドキュメント
sample_byパラメータについて言及している。 -
円城塔を近似する
僕が2023年11月ごろに書いた、Transformerライブラリを使ったファインチューニングの Colab ノートブック