mergekit-evolve による 進化的モデルマージ

以下の記事が面白かったので、簡単にまとめました。

・Evolutionary Model Merging For All

1. 進化的モデルマージ

「Sakana.ai」は約1か月前、「進化的モデルマージ」に関する論文を発表し、大きな話題を呼びました。

「進化的モデルマージ」を使用すると、マージで特定のコンピテンシーや資質をターゲットにすることができます。これがないと、モデルマージは手動の探索プロセスになります。数十回のマージを試し、それらを手動で評価し、マージパラメータが最終モデルの性能にどのように関連するか頭の中で考え出そうとすることになります。「進化的モデルマージ」を使用すると、どのような性質を持たせたいかを指定でき、最適化がそれを処理します。

「mergekit」では、この「進化的モデルマージ」を利用できます。

2. ハードウェア要件

7Bモデル の場合は 24GB のVRAMで十分です。

3. インストール

インストール手順は、次のとおりです。

cd /workspace
git clone https://github.com/arcee-ai/mergekit.git
cd mergekit
pip install -e .[evolve,vllm]

4. タスクの定義

マージレシピを最適化するには、まず何を最適化するかを決定する必要があります。「mergekit-evolve」は 「EleutherAI's language model evaluation harness」をバックエンドとして使用するため、理論上は「lm-eval」でサポートする全てのベンチマークを使用できます。

(1) spartqa-mchoiceデータセットの準備。
spartqa-mchoiceデータセットは、言語モデルの空間推論の機能をテストするための質問と回答のペアです。これをスコアリングに使用します。

ds = datasets.load_dataset("metaeval/spartqa-mchoice")["train"]
ds_p = ds.shuffle(seed=9163).select(range(1000))
ds_p.push_to_hub("my-hf-username/spartqa-train-1k", private=True)

(2) タスクの定義。
データセットに対してスコア付けするタスクを定義するため、YAMLおよび必要なヘルパーコードを作成します。

mkdir /workspace/eval_tasks

・/workspace/eval_tasks/spartqa_1k_train.yaml

task: spartqa_train
dataset_path: my-hf-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

・/workspace/eval_tasks/preprocess_spartqa.py

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:"

詳しくは、「New Task Guide」を参照してください。

(3) vicgalle/alpaca-gpt4データセットの準備。
vicgalle/alpaca-gpt4は一般的なプロンプト書式のデータセットであり、完了した応答後にEOSトークンを正しく出力することを覚えさせます。

マージでの一般的な問題の1つは、結果が特定のプロンプト書式に準拠しないことがよくあることです。マージレシピを手動で作成する場合、レイヤー間で重みを変更することで目的の動作を得るのは簡単ですが、アルゴリズムに最適化を任せているので、代わりにそのための小さなタスクを作成します。

ds = datasets.load_dataset("vicgalle/alpaca-gpt4")["train"]
df = ds.to_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("my-hf-username/alpaca-gpt4-500", private=True)

(4) タスクの定義。

・/workspace/eval_tasks/alpaca_prompt_format.yaml

task: alpaca_prompt_format
dataset_path: my-hf-username/alpaca-gpt4-500
output_type: multiple_choice
training_split: train
validation_split: train
test_split: train
doc_to_text: "### Instruction:\n{instruction}\n### Response:\n{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

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

どのモデルを含めるのか、どのマージ方法を使用するのか、どのタスクに対して最適化するのかをYAMLファイルで定義します。

今回は、次の3つのモデルを含めます。

・Hermes 2 Pro Mistral 7B (全体的に良いモデル)
・Dan's Adventurous Winds Mk2 7B (楽しいモデルだが、プロンプト書式なし)
・Zephyr 7B beta (質の高いプロンプト書式)

「mergekit」で実装されているほとんどの手法が使用できます。タスクは任意に重み付けできます。

genome:
    models:
      - NousResearch/Hermes-2-Pro-Mistral-7B
      - PocketDoc/Dans-AdventurousWinds-Mk2-7b
      - HuggingFaceH4/zephyr-7b-beta
    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
    weight: 0.4
  - name: spartqa_train
    weight: 0.6

6. マージの実行

(1) マージの実行。
マージの実行コマンドは、次のとおりです。

mergekit-evolve ./evol_merge_config.yml \
    --storage-path /workspace/evol_merge_storage \
    --task-search-path /workspace/eval_tasks \
    --vllm \
    --in-memory \
    --merge-cuda \
    --wandb

デフォルトでは、「mergekit-evolve」は100以上のマージを評価するか、CTRL+Cで停止するまで続行します。--max-fevalsで、この制限を増やすことができます。スクリプトが終了すると、最高スコアのマージの設定が /workspace/evol_merge_storage/best_config.yaml に出力されまれます。

(2) 最終モデルの取得。
次のように、「mergekit-yaml」を通じて実行することで、最終モデルを取得できます。

mergekit-yaml /workspace/evol_merge_storage/best_config.yaml --cuda /workspace/final_merge

Arcee Blog