2024年5月25日土曜日

M2 mac mini 上で mlx を使って LoRA Fine-Tuning をしてみる

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

準備

  • pyenv install 3.10.12

プロジェクトセットアップ

  • git clone https://github.com/ml-explore/mlx-examples.git
  • cd mlx-examples/lora
  • pyenv local 3.10.12
  • pip install -r requirements.txt

モデルのコンバート

LLaMA ベースの日本語モデル elyza/ELYZA-japanese-Llama-2-7b-instruct を更に safetensors 化したモデルを使います

  • python convert.py --hf-path 4bit/ELYZA-japanese-Llama-2-7b-instruct -q --q-bits 4

コンバートしたモデルは mlx_model 配下に作成されます

  • ls -1 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"

# JSONデータをPandas DataFrameに読み込む
df = pd.read_json(input_json_file)

# ターゲットをclosed_QAとopen_QAに絞る
df = df[(df.category == "closed_qa") | (df.category == "open_qa")].reset_index(
    drop=True
)

# テキスト長が短い学習データTOP1300をデータセットとして扱う
# 学習に時間がかかる場合もしくはメモリを使いすぎる場合は ascending=True にして短い順にするか 1300 の数を小さくすること
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とtest用のデータを100件ずつ取り出し、残りをtrainに分割
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
)


# ヘルパー関数:データフレームを新しい形式でJSON Linesファイルに変換
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")


# 各データセットを対応するJSON Linesファイルに変換
df_to_jsonl(train_df, "./ojousama/train.jsonl")
df_to_jsonl(valid_df, "./ojousama/valid.jsonl")
df_to_jsonl(test_df, "./ojousama/test.jsonl")
  • python data_convert.py

で各種 json ファイルができれば OK です

ls -1 ojousama
databricks-dolly-15k-ja-ojousama.json
test.jsonl
train.jsonl
valid.jsonl

学習させるデータ量について

完結に言うと

  • 文章をたくさん学習させたほうがいいがメモリが足らなくなる
  • 文章が短いと短い文章しか生成してくれなくなるがメモリは少なくて済む

という感じなのでデータセットを作成する際に以下のように文章の長さと文章数を調整すると良いです

# 質問+解答の長さを length というフィールドに登録
df["length"] = df.instruction.str.len() + df.output.str.len()
# 文章が長い順にソート
# df = df.sort_values(by="length", ascending=False)
# 質問+解答の長さが200文字以上400文字以下のみ抽出、この条件で734件で14GBほどメモリ消費
filterd_df = (df["length"] >= 200) & (df["length"] <= 400)
df = df.loc[filterd_df]
# あとはここを調整、1300以下ならそのままで OK
df = df.head(1300)

Fine-Tuning

  • python lora.py --model mlx_model --data ojousama --train --iters 600

iters で学習回数を指定できます
その他学習時のパラメータは python lora.py --help で確認できますが TrainingArguments で指定可能なパラメータすべてを指定できるわけではないようです

結果は以下の通りでだいたい5時間ほどかかりました (学習させるデータを短い版にすると1時間ほどで完了しました)
adapters.npz という numpy ベースのファイルができれば OK です

動作確認: 推論

  • python lora.py --model mlx_model --adapter-file adapters.npz --prompt "USER:日本の首都は? ASSISTANT:"

結果

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 上でも機械学習がやりやすくなるのかなと思いました

参考サイト

0 件のコメント:

コメントを投稿