M2 mac mini 上で mlx を使って LoRA Fine-Tuning をしてみる
概要
前回の Google Colab 上で動作させたコードを mlx 化して M2 mac mini 上で動作させてみました
環境
- macOS 14.5
- Python 3.10.12
- mlx 0.13.1
- mlx-example 69700d84
準備
プロジェクトセットアップ
モデルのコンバート
LLaMA ベースの日本語モデル elyza/ELYZA-japanese-Llama-2-7b-instruct を更に safetensors 化したモデルを使います
コンバートしたモデルは mlx_model 配下に作成されます
config.json
model.safetensors
model.safetensors.index.json
special_tokens_map.json
tokenizer.json
tokenizer_config.json
データセットのコンバート
データセットのフォーマットが少し独自なので変換する必要があります
ここから databricks-dolly-15k-ja-ojousama.json をダウンロードし ojousama ディレクトリ配下に保存します
こちらのスクリプトを拝借して test.json, valid.json, test.json を作成します
- pip install pandas scikit-learn
- vim data_convert.py
import json
import pandas as pd
from sklearn.model_selection import train_test_split
input_json_file = "./ojousama/databricks-dolly-15k-ja-ojousama.json"
df = pd.read_json(input_json_file)
df = df[(df.category == "closed_qa") | (df.category == "open_qa")].reset_index(
drop=True
)
df["length"] = df.instruction.str.len() + df.output.str.len()
df = df.sort_values(by="length", ascending=False)
df = df.head(1300)
df = df.sample(frac=1).reset_index(drop=True)
valid_df, remaining_df = train_test_split(df, test_size=len(df) - 100, random_state=42)
test_df, train_df = train_test_split(
remaining_df, test_size=len(remaining_df) - 100, random_state=42
)
def df_to_jsonl(df, file_name):
with open(file_name, "w", encoding="utf-8") as file:
for _, row in df.iterrows():
formatted_data = {
"text": f"USER:{row['instruction']} ASSISTANT:{row['output']}"
}
file.write(json.dumps(formatted_data, ensure_ascii=False) + "\n")
df_to_jsonl(train_df, "./ojousama/train.jsonl")
df_to_jsonl(valid_df, "./ojousama/valid.jsonl")
df_to_jsonl(test_df, "./ojousama/test.jsonl")
で各種 json ファイルができれば OK です
ls -1 ojousama
databricks-dolly-15k-ja-ojousama.json
test.jsonl
train.jsonl
valid.jsonl
学習させるデータ量について
完結に言うと
- 文章をたくさん学習させたほうがいいがメモリが足らなくなる
- 文章が短いと短い文章しか生成してくれなくなるがメモリは少なくて済む
という感じなのでデータセットを作成する際に以下のように文章の長さと文章数を調整すると良いです
df["length"] = df.instruction.str.len() + df.output.str.len()
filterd_df = (df["length"] >= 200) & (df["length"] <= 400)
df = df.loc[filterd_df]
df = df.head(1300)
Fine-Tuning
iters で学習回数を指定できます
その他学習時のパラメータは python lora.py --help で確認できますが TrainingArguments で指定可能なパラメータすべてを指定できるわけではないようです
結果は以下の通りでだいたい5時間ほどかかりました (学習させるデータを短い版にすると1時間ほどで完了しました)
adapters.npz という numpy ベースのファイルができれば OK です
動作確認: 推論
結果
Loading pretrained model
Total parameters 1055.199M
Trainable parameters 2.097M
Loading datasets
Generating
USER:日本の首都は? ASSISTANT:ワタクシ思いますの、 東京は、日本の首都ですわ
ξ゚⊿゚)ξ
==========
使用したモデル
使用したデータセット
最後に
M2 mac mini 上で LLaMA ベースの日本語モデルを LoRA Fine-Tuning してみました
スクラッチではないので細かいところまで手は届きませんが性能的にはこれで十分かなと思います
また mlx に対応しておりフルに GPU を使って学習してくれるので Apple Sillicon 上でもそれなりのスピードが出せます
やはり M2 上では mlx が必須のようです
過去にいろいろな webui 系のツールを試してみましたがそれらも mlx に対応してくれると Apple Sillicon 上でも機械学習がやりやすくなるのかなと思いました
参考サイト