進化的モデルマージでAIモデルをパワーアップ！mergekit-evolve（GoogleColab版）

進化的なモデルマージとは、複数の言語モデルを組み合わせて、特定の能力や特性を持つ新しいモデルを作成する手法です。従来の手法では、どのモデルをどのように組み合わせるかは手探りで進める必要がありましたが、進化的なモデルマージでは、目標とする特性を指定することで、最適な組み合わせを自動的に探索できます。

この記事では、進化的なモデルマージを誰でも手軽に試せるツール「mergekit-evolve」の使い方を、初心者にもわかりやすく解説します。基本的なインストールから、モデルの評価、そして最終的なモデルの作成まで、順を追って説明していきます。コード例も豊富に用意し、実際に手を動かしながら理解を深められるように工夫しました。

準備：必要な環境を整えよう
- mergekitのインストール
タスクの定義：モデルに何をさせたいか決めよう
進化的なモデルマージの設定
マージの実行と最終モデルの取得
Google Colabノートブック
参考サイト
- 関連

準備：必要な環境を整えよう

進化的なモデルマージを行うには、GPUを搭載した環境が必要です。7B程度のモデルであれば、24GBのVRAMがあれば十分です。

mergekitのインストール

まずは、mergekitをインストールしましょう。以下のコマンドをターミナルで実行します。

!pip install --upgrade pip

!rm -rf /opt/conda/lib/python3.10/site-packages/aiohttp*
!rm -rf /opt/conda/lib/python3.10/site-packages/multidict*

# 作業ディレクトリに移動
!mkdir /content/workspace
%cd /content/workspace
# mergekitのリポジトリをクローン
!git clone https://github.com/arcee-ai/mergekit.git
# クローンしたディレクトリに移動
%cd /content/workspace/mergekit
# mergekitをインストールし、進化的なモデルマージ(evolve)とvllmバックエンドを使用するための追加パッケージをインストール
!pip install -e .[evolve,vllm]
#!pip install --no-cache-dir -e .[evolve,vllm]
#!pip install --force-reinstall --no-deps aiohttp==3.9.1
#!pip install --force-reinstall --no-deps multidict==6.0.4
#!pip install --force-reinstall -e .[evolve,vllm]

username = 'MakiAi'
token = 'HF_TOKEN'
license = "apache-2.0"

!pip install -qU huggingface_hub

import yaml

from huggingface_hub import ModelCard, ModelCardData, HfApi

try:
    from google.colab import userdata
    HF_TOKEN = userdata.get(token)
    WANDB_API_KEY = userdata.get('WANDB_API_KEY')
except:
    from kaggle_secrets import UserSecretsClient
    user_secrets = UserSecretsClient()
    HF_TOKEN = user_secrets.get_secret(token)
    WANDB_API_KEY = user_secrets.get_secret('WANDB_API_KEY')

from huggingface_hub import login
login(token=HF_TOKEN)

タスクの定義：モデルに何をさせたいか決めよう

進化的なモデルマージでは、モデルの評価基準となるタスクを定義する必要があります。ここでは、空間認識能力を評価するspartqa-mchoiceと、プロンプト形式の応答能力を評価するalpaca-gpt4の2つのタスクを定義します。

spartqa-mchoiceデータセットの準備

import datasets
# Hugging Face Datasetsからspartqa-mchoiceデータセットをロード。["train"]は訓練データのみを取得するという意味
ds = datasets.load_dataset("metaeval/spartqa-mchoice")["train"]
# データをシャッフルし、ランダムに1000個のサンプルを選択
ds_p = ds.shuffle(seed=9163).select(range(1000))
# 選択したサンプルをHugging Face Hubにアップロード(private=Trueは非公開でアップロードするという意味)
ds_p.push_to_hub(f"{username}/spartqa-train-1k", private=True, token=HF_TOKEN)

spartqa-mchoiceタスクの定義

# タスク名を指定
# `/workspace/eval_tasks/spartqa_1k_train.yaml`ファイルを作成し、以下の内容を記述します。
import os

# タスクの定義を文字列として作成
task = f"""
task: spartqa_train
# データセットのパスを指定(Hugging Face Hubにアップロードしたデータセットのパス)
dataset_path: {username}/spartqa-train-1k
output_type: multiple_choice  # 出力形式を多肢選択式に指定
training_split: train
validation_split: train
test_split: train
doc_to_text: !function preprocess_spartqa.doc_to_text
doc_to_choice: [ 'A', 'B', 'C', 'D' ]
doc_to_target: "{{answer}}"
metric_list:
  - metric: acc
    aggregation: mean
    higher_is_better: true
metadata:
  version: 1.0

"""

# 保存先のディレクトリを作成
directory_path = "/workspace/eval_tasks"
os.makedirs(directory_path, exist_ok=True)  # exist_ok=Trueで、ディレクトリが存在してもエラーにならないようにする

# ファイルパスを指定
file_path = os.path.join(directory_path, "spartqa_1k_train.yaml")

# ファイルを書き込みモードで開く
with open(file_path, "w") as f:
    # 文字列をファイルに書き込む
    f.write(task)

# データセットの各サンプルを、モデルが理解できる形式のテキストに変換する関数
# `/workspace/eval_tasks/preprocess_spartqa.py`ファイルを作成し、以下の内容を記述します。

# doc_to_text関数の定義を文字列として作成
doc_to_text = '''def doc_to_text(doc) -> str:
  answer_chunks = []
  for idx, answer in enumerate(doc["candidate_answers"]):
    letter = "ABCD"[idx]
    answer_chunks.append(f"{letter}. {answer}")
  answers = "\\n".join(answer_chunks)
  return f"Context:\\n{doc['story']}\\n\\nQuestion: {doc['question']}\\n{answers}\\nAnswer:"
'''

doc_to_text_file_path = os.path.join(directory_path, "preprocess_spartqa.py")

# ファイルを書き込みモードで開く
with open(doc_to_text_file_path, "w") as f:
    # 文字列をファイルに書き込む
    f.write(doc_to_text)

!cat $doc_to_text_file_path

alpaca-gpt4データセットの準備とタスク定義

# `spartqa-mchoice`と同様の手順で、`alpaca-gpt4`データセットの準備とタスク定義を行います。
# Hugging Face Datasetsからalpaca-gpt4データセットをロード。["train"]は訓練データのみを取得するという意味

ds = datasets.load_dataset("vicgalle/alpaca-gpt4")["train"]
df = ds.to_pandas() # pandasのデータフレームに変換

no_input = df[df.input.map(len) < 1]
examples = no_input.sample(n=500, replace=False, random_state=749)
ds_p = datasets.Dataset.from_pandas(examples)
ds_p.push_to_hub(f"{username}/alpaca-gpt4-500", private=True, token=HF_TOKEN) # 選択したサンプルをHugging Face Hubにアップロード(private=Trueは非公開でアップロードするという意味)

# タスク名を指定
# `/workspace/eval_tasks/alpaca_prompt_format.yaml`ファイルを作成し、以下の内容を記述します。

# doc_to_text関数の定義を文字列として作成
prompt_format = """
task: alpaca_prompt_format
dataset_path: MakiAi/alpaca-gpt4-500
output_type: multiple_choice
training_split: train
validation_split: train
test_split: train
doc_to_text: |
  ### Instruction:
  {instruction}
  ### Response:
  {output}
doc_to_choice:
  - "</s>" # replace with your model's EOS token if it is different
  # and now some incorrect options
  - "<|im_end|>"
  - "<|im_start|>"
  - "### Instruction:"
  - "USER:"
doc_to_target: 0
metric_list:
  - metric: acc
    aggregation: mean
    higher_is_better: true
metadata:
  version: 1.0
"""

doc_to_text_file_path = os.path.join(directory_path, "alpaca_prompt_format.yaml")

# ファイルを書き込みモードで開く
with open(doc_to_text_file_path, "w") as f:
    # 文字列をファイルに書き込む
    f.write(prompt_format)

!cat $doc_to_text_file_path

進化的なモデルマージの設定

次に、どのモデルを組み合わせて、どのタスクで評価するのかを定義する設定ファイルを作成します。

genome:
    # 組み合わせるモデルを指定(ここでは3つのモデルを指定)
    models:  
      - NousResearch/Hermes-2-Pro-Mistral-7B  
      - PocketDoc/Dans-AdventurousWinds-Mk2-7b
      - HuggingFaceH4/zephyr-7b-beta
    # マージ方法を指定 (ここではtask_arithmeticという方法を使用)
    merge_method: task_arithmetic  
    base_model: mistralai/Mistral-7B-v0.1
    layer_granularity: 8 # sane default
    allow_negative_weights: true # useful with task_arithmetic
# タスク名を指定
tasks:
  - name: alpaca_prompt_format  
    # タスクの重みを指定 (このタスクの重要度を0.4とする)
    weight: 0.4  
  # タスク名を指定
  - name: spartqa_train  
    # タスクの重みを指定 (このタスクの重要度を0.6とする)
    weight: 0.6

# マージ設定を文字列として作成
merge_config = """
genome:
    # 組み合わせるモデルを指定(ここでは3つのモデルを指定)
    models:
      - NousResearch/Hermes-2-Pro-Mistral-7B
      - PocketDoc/Dans-AdventurousWinds-Mk2-7b
      - HuggingFaceH4/zephyr-7b-beta
    # マージ方法を指定 (ここではtask_arithmeticという方法を使用)
    merge_method: task_arithmetic
    base_model: mistralai/Mistral-7B-v0.1
    layer_granularity: 8 # sane default
    allow_negative_weights: true # useful with task_arithmetic
# タスク名を指定
tasks:
  - name: alpaca_prompt_format
    # タスクの重みを指定 (このタスクの重要度を0.4とする)
    weight: 0.4
  # タスク名を指定
  - name: spartqa_train
    # タスクの重みを指定 (このタスクの重要度を0.6とする)
    weight: 0.6
"""

# ファイルパスを指定 (作業ディレクトリ直下に保存する場合)
config_file_path = "./evol_merge_config.yml"

# ファイルを書き込みモードで開く
with open(config_file_path, "w") as f:
    # 文字列をファイルに書き込む
    f.write(merge_config)

マージの実行と最終モデルの取得

いよいよ、進化的なモデルマージを実行します。

# 進化的なモデルマージを実行
!mergekit-evolve ./evol_merge_config.yml \  
    # マージ結果を保存するディレクトリを指定 \
    --storage-path /workspace/evol_merge_storage \  
    # タスク定義ファイルを保存しているディレクトリを指定 \
    --task-search-path /workspace/eval_tasks \  
    # vllmバックエンドを使用 (vllmはモデルの推論を高速化するためのライブラリ)
    --vllm \  
    # モデルをメモリにロード (高速化のため)
    --in-memory \  
    # GPUを使用してマージ (高速化のため)
    --merge-cuda \  
    # Weights & Biasesにログを送信 (実験の進捗状況や結果を記録・可視化できるツール)
    --wandb

!wandb login $WANDB_API_KEY

!cat /workspace/eval_tasks/alpaca_prompt_format.yaml

# 進化的なモデルマージを実行
!mergekit-evolve ./evol_merge_config.yml \
    --storage-path /workspace/evol_merge_storage \
    --task-search-path /workspace/eval_tasks \
    --vllm \
    --in-memory \
    --merge-cuda \
    --wandb

# 最適なモデルの設定を読み込んで、最終的なモデルを作成 (--cudaはGPUを使用するオプション)
# マージが完了すると、最適なモデルの設定が`/workspace/evol_merge_storage/best_config.yaml`に保存されます。
!mergekit-yaml /workspace/evol_merge_storage/best_config.yaml --cuda /workspace/final_merge

これで、進化的なモデルマージによって作成された、あなたの目的に合った新しいモデルが完成しました！