Zenn
🥷

MLXを使ってMacのローカル環境で”ござる”loraを作成してみる

2024/03/31に公開
macOS
lora
LLM
mlx
tech

はじめに

この記事は、mlxを用いてMacのローカル環境でloraファイルを作る方法をまとめたものです。
私の別の記事でllama.cppを用いたlora作成を試しましたが、GPUを効率的に使用できるためか作成速度が3倍程度早くなりました。

この記事で紹介する内容

  • huggingFace上のモデルファイルをMLX用にコンバートする
  • ござるデータセットをLora用に事前加工する
  • Loraを構築する
  • 構築したLoraを使って推論を行う

利用するデータセット

ござるデータセットを使用します。

https://huggingface.co/datasets/bbz662bbz/databricks-dolly-15k-ja-gozaru/tree/main

利用するモデル

CyberAgent社のcalm2-chatを利用します。

https://huggingface.co/cyberagent/calm2-7b-chat

環境構築

前提条件

以下の環境構築は済んでいる前提とします。

  • venvの導入

mlx-examplesをクローン

mlx-examplesをgitからクローンします。
適当なフォルダで以下を実行してください。

git clone https://github.com/ml-explore/mlx-examples.git

loraディレクトリでvevn環境を構築する。

loraディレクトリに移る

cd mlx-examples/lora

仮想環境を構築する

以下のコマンドを1行ずつ実行してください。
仮想環境を構築後、フォルダ内のrequirements.txtの内容に従って仮想環境を構築しています。

python -m venv venv
. ./venv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt

モデルの準備

以下のコマンドを実行します。
calm2を量子化した上でmlx形式に変更しています。

python convert.py --hf-path cyberagent/calm2-7b-chat -q

実行後、mlx_modelというディレクトリが生成され、その中にモデルデータが保存されます。

学習用データの用意

事前知識：学習用データについて

loraファイルの作成には以下の3ファイルを指定のディレクトリに入れておく必要があります。サンプルデータがdataディレクトリに入っているのでそちらを見るとイメージが湧きやすいと思います。

  • test.jsonl
  • train.jsonl
  • valid.jsonl

それぞれのデータは以下のjson形式で記載されている必要があります。
USER: ASSISTANT:の部分はモデルに合わせてください。ここではcalm2-chatの形式に合わせています。

{"text": "USER:アメリカにはいくつの州があるのでしょうか？ ASSISTANT:50でござる。"}

学習用データの作成

ござるデータセットをtrain.jsonlなどの形式に加工します。
今回は"gozaru"というディレクトリを作成し、その中にhuggingFaceのdatabricks-dolly-15k-ja-gozaru.jsonをダウンロードした状態で以下を実行しました。

import pandas as pd
from sklearn.model_selection import train_test_split

# 元のサンプルデータのファイルパス
input_json_file = './gozaru/databricks-dolly-15k-ja-gozaru.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をデータセットとして扱う
df["length"] = df.instruction.str.len() + df.output.str.len()
df = df.sort_values(by="length")
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, "./gozaru/train.jsonl")
df_to_jsonl(valid_df, "./gozaru/valid.jsonl")
df_to_jsonl(test_df, "./gozaru/test.jsonl")

学習用データセット作成時の注意点

学習データがあまりにも長いと学習中にスタックする事象が発生します。
そのため、この記事ではテキスト長を短いものに限定して学習対象としました。

loraの構築

以下のコマンドを実行してください。
mlx_mldelのディレクトリのファイルを対象に、gozaruディレクトリのデータセットを用いて学習を行っています。

python lora.py --model mlx_model --data gozaru --train --iters 600

学習結果はadapters.npzという形式で保存されます。

学習速度とLossの低下状況

ログそのまま

Iter 1: Val loss 4.280, Val took 24.563s
Iter 10: Train loss 3.974, It/sec 0.759, Tokens/sec 85.948
Iter 20: Train loss 3.139, It/sec 0.666, Tokens/sec 73.791
Iter 30: Train loss 2.405, It/sec 0.668, Tokens/sec 75.600
Iter 40: Train loss 2.187, It/sec 0.642, Tokens/sec 74.842
Iter 50: Train loss 2.036, It/sec 0.663, Tokens/sec 77.176
Iter 60: Train loss 2.264, It/sec 0.609, Tokens/sec 74.810
Iter 70: Train loss 2.202, It/sec 0.667, Tokens/sec 74.153
Iter 80: Train loss 2.067, It/sec 0.570, Tokens/sec 64.564
Iter 90: Train loss 2.067, It/sec 0.678, Tokens/sec 77.798
Iter 100: Train loss 1.971, It/sec 0.675, Tokens/sec 75.792
Iter 100: Saved adapter weights to adapters.npz.
Iter 110: Train loss 2.081, It/sec 0.634, Tokens/sec 76.433
Iter 120: Train loss 2.137, It/sec 0.689, Tokens/sec 79.014
Iter 130: Train loss 2.146, It/sec 0.638, Tokens/sec 73.037
Iter 140: Train loss 1.871, It/sec 0.647, Tokens/sec 70.466
Iter 150: Train loss 2.192, It/sec 0.634, Tokens/sec 75.312
Iter 160: Train loss 2.076, It/sec 0.641, Tokens/sec 74.046
Iter 170: Train loss 1.895, It/sec 0.663, Tokens/sec 74.628
Iter 180: Train loss 2.184, It/sec 0.620, Tokens/sec 70.821
Iter 190: Train loss 1.993, It/sec 0.685, Tokens/sec 75.099
Iter 200: Train loss 2.001, It/sec 0.713, Tokens/sec 78.022
Iter 200: Val loss 1.945, Val took 25.359s
Iter 200: Saved adapter weights to adapters.npz.
Iter 210: Train loss 2.044, It/sec 0.685, Tokens/sec 74.006
Iter 220: Train loss 2.110, It/sec 0.666, Tokens/sec 73.499
Iter 230: Train loss 2.036, It/sec 0.642, Tokens/sec 75.719
Iter 240: Train loss 1.818, It/sec 0.659, Tokens/sec 76.298
Iter 250: Train loss 2.047, It/sec 0.669, Tokens/sec 76.333
Iter 260: Train loss 1.810, It/sec 0.719, Tokens/sec 81.283
Iter 270: Train loss 1.883, It/sec 0.654, Tokens/sec 75.320
Iter 280: Train loss 1.880, It/sec 0.631, Tokens/sec 76.220
Iter 290: Train loss 1.851, It/sec 0.624, Tokens/sec 75.879
Iter 300: Train loss 1.940, It/sec 0.669, Tokens/sec 75.097
Iter 300: Saved adapter weights to adapters.npz.
Iter 310: Train loss 1.965, It/sec 0.648, Tokens/sec 75.416
Iter 320: Train loss 1.805, It/sec 0.634, Tokens/sec 69.200
Iter 330: Train loss 2.072, It/sec 0.688, Tokens/sec 78.941
Iter 340: Train loss 1.955, It/sec 0.652, Tokens/sec 73.846
Iter 350: Train loss 1.985, It/sec 0.646, Tokens/sec 78.429
Iter 360: Train loss 1.929, It/sec 0.670, Tokens/sec 76.650
Iter 370: Train loss 1.906, It/sec 0.663, Tokens/sec 75.363
Iter 380: Train loss 1.997, It/sec 0.670, Tokens/sec 73.378
Iter 390: Train loss 1.865, It/sec 0.688, Tokens/sec 76.438
Iter 400: Train loss 1.925, It/sec 0.691, Tokens/sec 74.740
Iter 400: Val loss 1.908, Val took 25.114s
Iter 400: Saved adapter weights to adapters.npz.
Iter 410: Train loss 2.014, It/sec 0.657, Tokens/sec 75.481
Iter 420: Train loss 1.837, It/sec 0.645, Tokens/sec 72.056
Iter 430: Train loss 1.790, It/sec 0.668, Tokens/sec 76.596
Iter 440: Train loss 1.769, It/sec 0.698, Tokens/sec 76.131
Iter 450: Train loss 1.740, It/sec 0.609, Tokens/sec 67.691
Iter 460: Train loss 1.944, It/sec 0.649, Tokens/sec 75.149
Iter 470: Train loss 1.737, It/sec 0.661, Tokens/sec 79.137
Iter 480: Train loss 1.800, It/sec 0.622, Tokens/sec 71.625
Iter 490: Train loss 1.742, It/sec 0.590, Tokens/sec 68.466
Iter 500: Train loss 1.766, It/sec 0.572, Tokens/sec 65.863
Iter 500: Saved adapter weights to adapters.npz.
Iter 510: Train loss 1.884, It/sec 0.648, Tokens/sec 72.992
Iter 520: Train loss 1.759, It/sec 0.609, Tokens/sec 69.372
Iter 530: Train loss 1.835, It/sec 0.644, Tokens/sec 76.073
Iter 540: Train loss 1.913, It/sec 0.677, Tokens/sec 74.091
Iter 550: Train loss 1.543, It/sec 0.614, Tokens/sec 72.283
Iter 560: Train loss 1.597, It/sec 0.660, Tokens/sec 75.033
Iter 570: Train loss 1.884, It/sec 0.679, Tokens/sec 75.031
Iter 580: Train loss 1.734, It/sec 0.671, Tokens/sec 78.293
Iter 590: Train loss 1.745, It/sec 0.661, Tokens/sec 73.682
Iter 600: Train loss 1.921, It/sec 0.675, Tokens/sec 77.199
Iter 600: Val loss 1.902, Val took 25.761s

loraを用いて推論する

以下のコマンドで推論を実行します。

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

推論結果

プロンプト

日本の首都は?

アウトプット

ASSISTANT:日本首都は東京でござる。

まとめ

いかがだったでしょうか？

今回は話題のLLMでLoraを作る方法をまとめました。
Macのスペック持て余している方は是非今回の手順で使ってみてください！

私のTwitterではLLMに限らず、AIを活用した業務改善情報の発信をしておりますのでご興味のある方は是非フォローをお願いします。
https://twitter.com/Linus_lab

Discussion