1. きっかけ
コマンドラインで学習することが可能になりました。
しかし、引数で渡せず、yamlファイルを用意しなければいけなかったりする点が面倒なので、Notebook形式で楽手できる方法を探ります。
MacBookAir（2025年1月初売りで購入でウキウキ）
CPU：M3
ユニファイドメモリー：24GB
SSD：1T
バージョン関連
Python==3.11.8
mlx_lm==0.21.1
使用したモデルと学習データ
モデル：llm-jp/llm-jp-3-3.7b
（インストラクションチューニングされていないモデル）
学習データ：llm-jp/magpie-sft-v1.0
コマンドラインでの学習の時と、モデルと学習データを変えました。理由は以下の通りです。
- メモリーのスワップがひどすぎた
- 小さなモデルなのでシンプルなインストラクションチューニングで効果をわかりやすくしたい。
- Qiitaを書くまでの時間を少なくしたい。w
2. GitHubをよく見る
久々にargparseを使うので、復習をしました。そんなレベルの人間です。
まず、ガッツリ、lora.pyを読みました。そのなかで着目したのはmain()関数です。lora.pyを呼んだらこの関数が呼ばれる様になっているのが着目理由です。
ざっと説明すると、設定を読み込んで、run関数に引数を渡しています。
main関数（簡単な解説）
def main():
os.environ["TOKENIZERS_PARALLELISM"] = "true"
parser = build_parser() # コマンドライン引数の処理
args = parser.parse_args()
config = args.config # コマンドライン引数のconfigの処理
args = vars(args)
if config: # コマンドライン引数でconfigを受け取ったら、そのyamlファイルの読み込み
print("Loading configuration file", config)
with open(config, "r") as file:
config = yaml.load(file, yaml_loader)
# Prefer parameters from command-line arguments
for k, v in config.items():
if args.get(k, None) is None:
args[k] = v
# Update defaults for unspecified parameters
for k, v in CONFIG_DEFAULTS.items(): # 文字通り、指定されていないデフォルト設定の読み込み。
if args.get(k, None) is None:
args[k] = v
run(types.SimpleNamespace(**args)) # ここでrun関数に引数を渡す。
3. 作戦
上のmain関数の中にはyamlファイルを読み込むことが書かれています。読み込んだ内容と、指定していない項目はCONFIG_DEFAULTSから読み取って、argsを作っている。
つまりCONFIG_DEFAULTSの内容を何かの変数にいれて、この変数をrunに渡せばいいってことじゃないですか！
決定！これで行こう
mlx_lmのGitHubにはyamlファイルのサンプルがありますので、CONFIG_DEFAULTSに記載されていなくて、サンプルにはある項目を探しましょう。
学習する層を指定して（絞って）あげようかな。
lora_config.yamlから抜粋
# 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"]
rank: 8
scale: 20.0
dropout: 0.0
学習する層を設定するのはlora_parametersのkeysになりそうです。これをアテンション層のQKV層のみを学習して、インストラクションチューニングしましょう。
モデルをプリントすると以下のような出力が出ます。カッコの中に書いてあるself_attenとか、q_projなどが層の名前になります。これを指定する場合はself_atten.q_projとすればいいです。この辺りはhuggingfaceやunslothのlora、PEFTなどの解説を読めばわかるはず。
Model(
(model): LlamaModel(
(embed_tokens): QuantizedEmbedding(99584, 3072, group_size=64, bits=4)
(layers.0): TransformerBlock(
(self_attn): Attention(
(q_proj): QuantizedLinear(input_dims=3072, output_dims=3072, bias=False, group_size=64, bits=4)
(k_proj): QuantizedLinear(input_dims=3072, output_dims=3072, bias=False, group_size=64, bits=4)
(v_proj): QuantizedLinear(input_dims=3072, output_dims=3072, bias=False, group_size=64, bits=4)
(o_proj): QuantizedLinear(input_dims=3072, output_dims=3072, bias=False, group_size=64, bits=4)
(rope): RoPE(128, traditional=False)
)
(mlp): MLP(
(gate_proj): QuantizedLinear(input_dims=3072, output_dims=8192, bias=False, group_size=64, bits=4)
(down_proj): QuantizedLinear(input_dims=8192, output_dims=3072, bias=False, group_size=64, bits=4)
(up_proj): QuantizedLinear(input_dims=3072, output_dims=8192, bias=False, group_size=64, bits=4)
)
(input_layernorm): RMSNorm(3072, eps=1e-05)
(post_attention_layernorm): RMSNorm(3072, eps=1e-05)
)
(layers.1): TransformerBlock(
(self_attn): Attention(
(q_proj): QuantizedLinear(input_dims=3072, output_dims=3072, bias=False, group_size=64, bits=4)
(k_proj): QuantizedLinear(input_dims=3072, output_dims=3072, bias=False, group_size=64, bits=4)
(v_proj): QuantizedLinear(input_dims=3072, output_dims=3072, bias=False, group_size=64, bits=4)
(o_proj): QuantizedLinear(input_dims=3072, output_dims=3072, bias=False, group_size=64, bits=4)
< 以下省略 >
4. パラメータを設定する
4-1. デフォルト値の読み込み
デフォルト値の読み込み
from mlx_lm import lora
training_args = lora.CONFIG_DEFAULTS # デフォルト引数をargsという変数に格納
print(training_args)
# {'model': 'mlx_model',
# 'train': False,
# 'fine_tune_type': 'lora',
# 'data': 'data/',
# 'seed': 0,
# 'num_layers': 16,
# 'batch_size': 4,
# 'iters': 1000,
# 'val_batches': 25,
# 'learning_rate': 1e-05,
# 'steps_per_report': 10,
# 'steps_per_eval': 200,
# 'resume_adapter_file': None,
# 'adapter_path': 'adapters',
# 'save_every': 100,
# 'test': False,
# 'test_batches': 500,
# 'max_seq_length': 2048,
# 'config': None,
# 'grad_checkpoint': False,
# 'lr_schedule': None,
# 'lora_parameters': {'rank': 8, 'alpha': 16, 'dropout': 0.0, 'scale': 10.0}}
確認出来たらこれらの設定を変更します。
4-2. パラメータの設定
args変数にいれたデフォルトののパラメータを更新しちゃいましょう。
argsパラメータ設定
mlx_path = "llm-jp-3-3.7b" # このフォルダに量子化したモデルを保存してあります。
dataset_path = "./magpie" # このフォルダに学習、評価、テストのjsonlを保存してあります。
training_args['model'] = mlx_path# モデルのレポジトリ名orローカルのフォルダパス
training_args['data'] = saved_dataset # データセットのパス
args['train'] = True # 学習モードオン
training_args['max_seq_length'] = 1024 * 2 # 推論トークン数
training_args['iters'] = 200 # 学習するミニバッチ数
training_args['batch_size'] = 2 # 学習率
training_args['learning_rate'] = 3e-04 # 学習率：大き目で効果を出やすく
training_args['steps_per_report'] = 5 # 何ステップごとに学習状態を表示するか
training_args['val_batches'] = 2 # 評価バッチ数
training_args['test_batches'] = 2 # テストバッチ数
training_args['lora_parameters'] = {
'rank': 8,
'alpha': 64, # ここも効果が出やすいように大きくしました
'dropout': 0.2,
'scale': 10.0,
'keys': [ # ここを設定したかった！
"self_attn.q_proj", # modelをプリントして、
"self_attn.k_proj", # 学習したい層を指定しましょう
"self_attn.v_proj"
]}
print(training_args)
# {'model': 'llm-jp-3-3.7b',
# 'train': True,
# 'fine_tune_type': 'lora',
# 'data': './magpie',
# 'seed': 0,
# 'num_layers': 16,
# 'batch_size': 2,
# 'iters': 200,
# 'val_batches': 2,
# 'learning_rate': 0.0003,
# 'steps_per_report': 5,
# 'steps_per_eval': 200,
# 'resume_adapter_file': None,
# 'adapter_path': 'adapters',
# 'save_every': 100,
# 'test': False,
# 'test_batches': 2,
# 'max_seq_length': 2048,
# 'config': None,
# 'grad_checkpoint': False,
# 'lr_schedule': None,
# 'lora_parameters': {'rank': 8,
# 'alpha': 64,
# 'dropout': 0.2,
# 'scale': 10.0,
# 'keys': ['self_attn.q_proj', 'self_attn.k_proj', 'self_attn.v_proj']}}
うっす、設定ができた！
じゃ、学習に移ろう。
【追記】
itersは「学習するミニバッチ数」だと思います。↓
https://github.com/ml-explore/mlx-examples/blob/52c41b5b5abfdd4ee1c35bd362162b1dc7a62138/llms/mlx_lm/tuner/trainer.py#L223
ご意見いただけると嬉しいです。
5. 学習の実行
いよいよです。伊予は愛媛です。はい。つっこんで。
学習の実行！
学習の実行
import types
lora.run(types.SimpleNamespace(**args))
作った辞書をjsonみたいな形式に変換してからlora.pyのrun関数に渡すだけ。
ここはlora.pyのmain関数を参考にしました。
出力
Loading pretrained model
Loading datasets
Training
Trainable parameters: 0.109% (4.129M/3782.913M)
Starting training..., iters: 200
Iter 1: Val loss 2.134, Val took 3.812s
Iter 5: Train loss 2.095, Learning Rate 3.000e-04, It/sec 0.193, Tokens/sec 114.222, Trained Tokens 2958, Peak mem 37.790 GB
Iter 10: Train loss 2.425, Learning Rate 3.000e-04, It/sec 0.142, Tokens/sec 110.436, Trained Tokens 6840, Peak mem 37.790 GB
Iter 15: Train loss 5.140, Learning Rate 3.000e-04, It/sec 0.236, Tokens/sec 112.437, Trained Tokens 9223, Peak mem 37.790 GB
Iter 20: Train loss 3.492, Learning Rate 3.000e-04, It/sec 0.180, Tokens/sec 114.488, Trained Tokens 12403, Peak mem 37.790 GB
Iter 25: Train loss 4.450, Learning Rate 3.000e-04, It/sec 0.273, Tokens/sec 112.811, Trained Tokens 14470, Peak mem 37.790 GB
Iter 30: Train loss 3.355, Learning Rate 3.000e-04, It/sec 0.168, Tokens/sec 112.576, Trained Tokens 17817, Peak mem 37.790 GB
Iter 35: Train loss 2.854, Learning Rate 3.000e-04, It/sec 0.232, Tokens/sec 112.655, Trained Tokens 20245, Peak mem 37.790 GB
Iter 40: Train loss 2.491, Learning Rate 3.000e-04, It/sec 0.166, Tokens/sec 114.074, Trained Tokens 23691, Peak mem 37.790 GB
Iter 45: Train loss 2.639, Learning Rate 3.000e-04, It/sec 0.185, Tokens/sec 111.628, Trained Tokens 26714, Peak mem 37.790 GB
Iter 50: Train loss 2.466, Learning Rate 3.000e-04, It/sec 0.230, Tokens/sec 113.073, Trained Tokens 29174, Peak mem 37.790 GB
Iter 55: Train loss 2.435, Learning Rate 3.000e-04, It/sec 0.159, Tokens/sec 111.421, Trained Tokens 32671, Peak mem 37.790 GB
Iter 60: Train loss 2.116, Learning Rate 3.000e-04, It/sec 0.213, Tokens/sec 110.140, Trained Tokens 35260, Peak mem 37.790 GB
Iter 65: Train loss 2.148, Learning Rate 3.000e-04, It/sec 0.139, Tokens/sec 110.515, Trained Tokens 39223, Peak mem 37.790 GB
Iter 70: Train loss 2.062, Learning Rate 3.000e-04, It/sec 0.200, Tokens/sec 111.092, Trained Tokens 42007, Peak mem 37.790 GB
Iter 75: Train loss 2.005, Learning Rate 3.000e-04, It/sec 0.153, Tokens/sec 110.861, Trained Tokens 45625, Peak mem 37.790 GB
Iter 80: Train loss 2.052, Learning Rate 3.000e-04, It/sec 0.163, Tokens/sec 103.614, Trained Tokens 48795, Peak mem 37.790 GB
Iter 85: Train loss 2.070, Learning Rate 3.000e-04, It/sec 0.180, Tokens/sec 106.875, Trained Tokens 51768, Peak mem 37.790 GB
Iter 90: Train loss 1.909, Learning Rate 3.000e-04, It/sec 0.196, Tokens/sec 111.760, Trained Tokens 54624, Peak mem 37.790 GB
Iter 95: Train loss 2.055, Learning Rate 3.000e-04, It/sec 0.163, Tokens/sec 109.544, Trained Tokens 57983, Peak mem 37.790 GB
Iter 100: Train loss 1.909, Learning Rate 3.000e-04, It/sec 0.121, Tokens/sec 99.251, Trained Tokens 62099, Peak mem 37.790 GB
Iter 100: Saved adapter weights to adapters/adapters.safetensors and adapters/0000100_adapters.safetensors.
Iter 105: Train loss 1.862, Learning Rate 3.000e-04, It/sec 0.149, Tokens/sec 102.148, Trained Tokens 65520, Peak mem 37.790 GB
Iter 110: Train loss 1.913, Learning Rate 3.000e-04, It/sec 0.129, Tokens/sec 102.118, Trained Tokens 69466, Peak mem 37.790 GB
Iter 115: Train loss 1.898, Learning Rate 3.000e-04, It/sec 0.153, Tokens/sec 105.363, Trained Tokens 72913, Peak mem 37.790 GB
Iter 120: Train loss 1.947, Learning Rate 3.000e-04, It/sec 0.169, Tokens/sec 102.597, Trained Tokens 75940, Peak mem 37.790 GB
Iter 125: Train loss 1.866, Learning Rate 3.000e-04, It/sec 0.131, Tokens/sec 100.037, Trained Tokens 79745, Peak mem 37.790 GB
Iter 130: Train loss 1.834, Learning Rate 3.000e-04, It/sec 0.174, Tokens/sec 97.737, Trained Tokens 82560, Peak mem 37.790 GB
Iter 135: Train loss 2.040, Learning Rate 3.000e-04, It/sec 0.312, Tokens/sec 109.667, Trained Tokens 84320, Peak mem 37.790 GB
Iter 140: Train loss 1.890, Learning Rate 3.000e-04, It/sec 0.148, Tokens/sec 102.769, Trained Tokens 87797, Peak mem 37.790 GB
Iter 145: Train loss 1.879, Learning Rate 3.000e-04, It/sec 0.149, Tokens/sec 105.955, Trained Tokens 91364, Peak mem 37.790 GB
Iter 150: Train loss 1.828, Learning Rate 3.000e-04, It/sec 0.182, Tokens/sec 107.092, Trained Tokens 94304, Peak mem 37.790 GB
Iter 155: Train loss 1.806, Learning Rate 3.000e-04, It/sec 0.166, Tokens/sec 107.471, Trained Tokens 97532, Peak mem 37.790 GB
Iter 160: Train loss 1.685, Learning Rate 3.000e-04, It/sec 0.138, Tokens/sec 105.725, Trained Tokens 101368, Peak mem 37.790 GB
Iter 165: Train loss 1.790, Learning Rate 3.000e-04, It/sec 0.162, Tokens/sec 104.677, Trained Tokens 104608, Peak mem 37.790 GB
Iter 170: Train loss 1.727, Learning Rate 3.000e-04, It/sec 0.160, Tokens/sec 102.985, Trained Tokens 107831, Peak mem 37.790 GB
Iter 175: Train loss 1.890, Learning Rate 3.000e-04, It/sec 0.190, Tokens/sec 104.344, Trained Tokens 110579, Peak mem 37.790 GB
Iter 180: Train loss 1.778, Learning Rate 3.000e-04, It/sec 0.168, Tokens/sec 106.778, Trained Tokens 113754, Peak mem 37.790 GB
Iter 185: Train loss 1.753, Learning Rate 3.000e-04, It/sec 0.199, Tokens/sec 107.845, Trained Tokens 116467, Peak mem 37.790 GB
Iter 190: Train loss 1.768, Learning Rate 3.000e-04, It/sec 0.177, Tokens/sec 106.513, Trained Tokens 119483, Peak mem 37.790 GB
Iter 195: Train loss 1.824, Learning Rate 3.000e-04, It/sec 0.171, Tokens/sec 105.739, Trained Tokens 122571, Peak mem 37.790 GB
Iter 200: Val loss 1.602, Val took 3.592s
Iter 200: Train loss 1.830, Learning Rate 3.000e-04, It/sec 1.076, Tokens/sec 499.896, Trained Tokens 124895, Peak mem 37.790 GB
Iter 200: Saved adapter weights to adapters/adapters.safetensors and adapters/0000200_adapters.safetensors.
Saved final weights to adapters/adapters.safetensors.
できたできた！
でもここにたどり着くまで、ずいぶんとGitHubを眺めたのよ。うんうん
argparseの書き方も忘れてて復習したのよ。意外と、俺、頑張り屋（褒めて褒めて）
しっかり学習するなら、パラメの設定やデータの準備をしっかりやしましょうね。
6. 学習結果
じゃ、みていきましょう。
ベースとしたモデルはllm-jp/llm-jp-3-3.7bです。これは事前学習しかされていないので、質問形式で入力してもまともに返答がかえってきません。本来は言葉をいれたらその続きの言葉を出力するモデルなのです。
6-1. 学習「前」の出力確認
まずは学習前（素のllm-jp/llm-jp-3-3.7bの4bit量子化したもの）
response = generate(
model,
tokenizer,
prompt='# 指示\n大規模言語モデルって何？\n# 回答\n',
verbose=True
)
# ==========
#
# # 大規模言語モデルって何？
#
# 大規模言語モデルとは、大量のデータを学習させることで言語の理解を可能にするモデルです。
#
# # 大規模言語モデルって何？
#
# 大規模言語モデルとは、大量のデータを学習させることで言語の理解を可能にするモデルです。
#
# # 大規模言語モデルって何？
#
# < 中略 >
#
# 大規模言語モデルとは、大量の
# ==========
# Prompt: 14 tokens, 113.487 tokens-per-sec
# Generation: 256 tokens, 40.249 tokens-per-sec
# Peak memory: 2.330 GB
本来であれば「大規模言語モデルとは、」と入力すると、この言葉の続きを出力するモデルです。今回は学習の効果がわかりやすくなるように、敢えてインストラクションチューニングモデルに入力するようにプロンプトを入れています。要は同じ言葉を入力しているってこと。
そのため、出力がおかしくなり、同じ言葉の繰り返しが出力されています。
もう一つ試してみましょう。
response = generate(
model,
tokenizer,
prompt='# 指示\n石破茂氏って何をしている人ですか\n# 回答\n',
verbose=True
)
# ==========
#
# 石破茂氏って何をしている人ですか？
# # 回答
#
# # 回答
#
# # 回答
#
#
# < 中略 >
#
#
# # 回答
#
# # 回答
#
#
# ==========
# Prompt: 17 tokens, 137.744 tokens-per-sec
# Generation: 256 tokens, 39.532 tokens-per-sec
# Peak memory: 2.330 GB
こちらも同様の結果になっています。
6-2. 学習「後」の出力確認
では学習したモデルを使って試してみましょう。
まず、モデルの読み込み。
LoRAモデルとベースモデルの読み込み方
from mlx_lm import generate
from mlx_lm.utils import load
mlx_path = "llm-jp-3-3.7b/" # Q4モデルの保存フォルダのパス
model, tokenizer = load(
mlx_path,
adapter_path="./adapters", # LoRAモデルの保存フォルダのパス
)
モデルのロードができたら早速推論です。
推論１
response = generate(
model,
tokenizer,
prompt='# 指示\n大規模言語モデルって何？\n# 回答\n',
verbose=True
)
# ==========
# もちろんです、大規模言語モデル（LMM）は自然言語の理解や生成に役立つ技術です。LMMは、大量のテキストデータを分析し、それらを分析結果から得られる情報を活用して、自然言語の理解や生成を支援します。
#
# LMMは、主に以下の機能を提供します：
#
# 1. **言語モデルの構築**: 大規模言語モデルは、大量のテキストデータを分析し、それらを分析結果から得られる情報を活用して、言語モデルを構築します。これにより、言語の構造や意味をより正確に理解することができます。
#
# 2. **言語の学習**: 大規模言語モデルは、ユーザーが学習したい言語の学習を支援します。これにより、ユーザーは自分の言語をより正確に理解し、新しい言語を学ぶ際の助けとなります。
#
# 3. **自然言語の生成**: 大規模言語モデルは、ユーザーが作成した文章を自動的に生成する機能を提供します。これにより、ユーザーはより効率的に文章を作成することができます。
#
# 4. **対話システムの構築**: 大規模言語モデルは、対話システムの構築に役立ちます。LMMを使用することで、ユーザーが質問を入力すると、自動的に回答を生成することができます。
#
# 5. **自然言語の理解**: 大規模言語モデルは、自然言語の理解を支援します。
# ==========
# Prompt: 14 tokens, 101.316 tokens-per-sec
# Generation: 256 tokens, 36.318 tokens-per-sec
# Peak memory: 2.334 GB
LLMがLMMになってるし。ｗ
ウケるわぁ！
でも、質問に回答する形式に変化しているでしょ？
推論２
response = generate(
model,
tokenizer,
prompt='# 指示\n石破茂氏って何をしている人ですか\n# 回答\n',
verbose=True
)
# ==========
# 石破茂氏は日本の政治家で、自民党の衆議院議員です。現在、自民党の幹事長を務めています。
#
# ### 回答
# 石破茂氏は日本の政治家で、現在自民党の幹事長を務めています。彼は自民党の国会議員であり、特に外交や経済政策に強い関心を持っています。
# ==========
# Prompt: 17 tokens, 133.222 tokens-per-sec
# Generation: 59 tokens, 36.587 tokens-per-sec
# Peak memory: 2.334 GB
はい。チューニングできました。ちょっとこの回答には、「### 回答」とか出てきていて、二度回答しているようにも見えますが、質問に答えるタイプのモデルにファインチューニングできていそうな感じです。ま、圧倒的に学習量が少ないのは事実ですので、時間があれば学習量を増やしてみたいと思います。あ、情報が古いのは学習データが古いからなんです。
7. おわりに
これで、mlx_lmを使ってLLMを学習することができるようになりました。
いきなり13Bのモデルだとメモリーのスワップがすごかったので、今回は3.7Bのモデルに変更しました。きっと性能はそこそこなんだと思いますが、llm-jp-3-3.7bは、ほんとに素晴らしいベースモデルです。
ま、今回は学習する手立てを作ることが目標ですので、今回はこれで良し。
やっぱり、unslothのGradient Accumulationの機能が欲しいなぁ。
mlx_lmでは、まだやりたいことがあります。それは途中のcheckpoint（一時保存）からの学習の継続や、モデルのマージなどです。
またやり方がわかったらQiitaに投稿したいと思います。
ほんじゃまた