Mergoo: LLMの効率的なマージと微調整のためのライブラリ (MoE, Mixture of Adapters)

 

  • 「Merggoo」というLLMマージ用の新しいライブラリができたそうで、開発者の方がRedditでもPRポストを上げておられます。
  • もしかするとMegekitに代わるものになるのかもしれません。MoEのルーティングの学習や、LoRA Adapterの混合など。

面白そうなので後で試せればと思いますが、概要だけ先に見ておこうと思います。以下、開発者ブログの和訳です。

Mergoo: 効率的にマージし、微調整する (MoE, Mixture of Adapters)

huggingface.co

私たちは最近、複数の LLM エキスパートを簡単に統合し、統合された LLM を効率的にトレーニングするためのライブラリである mergoo を開発しました。このライブラリでは、さまざまな汎用またはドメインベースの LLM エキスパートの知識を効率的に統合できます。

🚀 Mergooでは次のことが可能です。

  • 複数のオープンソースLLMを簡単に統合できます
  • さまざまなマージ方法を適用：エキスパートの混合、アダプタの混合、レイヤーごとのマージが可能です
  • ゼロから始めることなくMoEを効率的にトレーニングできます
  • HuggingFece🤗モデルやトレーナーと互換性があります
導入

mergoo は、汎用的なものからドメイン固有のものまで、さまざまな LLM エキスパートの知識を統合するための信頼性が高く透過的なパイプラインを構築するように設計されています。エキスパートの混合、アダプタの混合、レイヤーごとのマージなどのさまざまな統合手法が組み込まれており、LLM ビルダーに柔軟性を提供します。マージされた LLM は、SFTrainer、PEFT、Trainer などの Hugging Face 🤗 トレーナーを使用して、特定のユースケースに合わせてさらに微調整できます。
次のセクションでは、MoE を使用して完全に微調整された LLM からマージされた LLM を構築する方法と、LoRA で微調整されたエキスパートからアダプタの混合 LLM を作成する方法を示す 2 つの例について説明します。

完全に微調整済みのLLMの混合

Branch-Train-Mixに続いて、ドメイン固有の LLM エキスパートは、Mixture-of-Expert (MoE) レイヤーのエキスパートとしてフィードフォワードパラメータをまとめ、残りのパラメータを平均化することで統合できます。MoE レイヤーは、後で微調整してトークン レベルのルーティングを学習できます。

例として、次のドメイン固有の LLM エキスパートを統合します。

  • ベースモデル: metal-llama/Llama-2-7b-hf
  • コードエキスパート: codellama/CodeLlama-7b-Python-hf
  • WikiChat エキスパート: stanford-oval/Flame-2-7b-WikiChat-fused

マージするための設定を指定します:

config = \
{
    "model_type": "llama",
    "num_experts_per_tok": 2,
    "experts":[
        {
            "expert_name" : "base_expert",
            "model_id" : "meta-llama/Llama-2-7b-hf"
        },
        {
            "expert_name" : "expert_1",
            "model_id" : "codellama/CodeLlama-7b-Python-hf"
        },
        {
            "expert_name" : "expert_2",
            "model_id" : "stanford-oval/Llama-2-7b-WikiChat-fused"
        }
    ],
    "router_layers":[
        "gate_proj",
        "up_proj",
        "down_proj"
    ],
}

次に、マージされたエキスパートのチェックポイントを構築して保存します。

import torch
from mergoo.compose_experts import ComposeExperts
model_id = "mergoo_llama_code_wikichat"
expertmerger = ComposeExperts(config, torch_dtype=torch.float16)
expertmerger.compose()
expertmerger.save_checkpoint(model_id)

以下では、マージされた LLM のチェックポイントが読み込まれ、Python コード命令データセットでさらに微調整されます。

from mergoo.models.modeling_llama import LlamaForCausalLM
import torch
import datasets
import random
from trl import SFTTrainer
from transformers import TrainingArguments

# load the composed checkkpoint
model = LlamaForCausalLM.from_pretrained(
    "mergoo_llama_code_wikichat", 
    device_map="auto", 
    torch_dtype=torch.bfloat16,
)# 'gate' / router layers are untrained hence loaded warning would appeare for them

# load the train dataset
dataset = datasets.load_dataset("iamtarun/python_code_instructions_18k_alpaca")['train']
dataset = dataset['prompt']
random.shuffle(dataset)
train_dataset =  datasets.Dataset.from_dict(dict(prompt=dataset[:-1000]))
eval_dataset = datasets.Dataset.from_dict(dict(prompt=dataset[-1000:]))

# specify training arguments
trainer_args = TrainingArguments(
    output_dir= "checkpoints/llama_moe",
    per_device_train_batch_size = 1,
    per_device_eval_batch_size = 1, 
    learning_rate= 1e-5,
    save_total_limit=1,
    num_train_epochs=1,
    eval_steps= 5000,
    logging_strategy="steps",
    logging_steps= 25,
    gradient_accumulation_steps=4,
    bf16=True
)

trainer = SFTTrainer(
    model,
    args= trainer_args,
    train_dataset= train_dataset,
    eval_dataset= eval_dataset,
    dataset_text_field="prompt",
)

# start training
trainer.train()

その後、コードを Huggingface Hub にプッシュできます (ぜひプッシュしてください):

model.push_to_hub("mergoo_llama_code_wikichat_trained")

mergooはMistralとBertベースのエキスパートもサポートしています。

アダプターの混合(Mixture of Adapters)

mergoo は、複数のアダプタ (LoRA) を統合された MoE スタイルのアーキテクチャに統合することを容易にします。これは、微調整された LoRA の上にゲーティング レイヤーとルーティング レイヤーを適用することで実現されます。

アダプタ混合 LLM を構築するには:

  • 同じベースモデルで微調整されたアダプタ（LoRA）のプールを収集する
  • Mergoo を適用して MoE スタイルのマージされたエキスパートを作成する
  • 下流のタスクでマージされたエキスパートを微調整する

たとえば、顧客サポート領域では次の専門家を統合できます。

  • predibase/customer_support
  • predibase/customer_support_accounts
  • predibase/customer_support_orders

次のように構成を指定して、マージされたチェックポイントを構築します。

import torch
from mergoo.compose_experts import ComposeExperts

model_id = "mergoo_customer_suppoer_moe_lora"
config = {
    "model_type": "mistral",
    "num_experts_per_tok": 2,
    "base_model": "mistralai/Mistral-7B-v0.1",
    "experts": [
        {
            "expert_name": "adapter_1", 
            "model_id": "predibase/customer_support"
        },
        {
            "expert_name": "adapter_2", 
            "model_id": "predibase/customer_support_accounts"
        },
        {
            "expert_name": "adapter_3", 
            "model_id": "predibase/customer_support_orders"
        }
    ],
}

expertmerger = ComposeExperts(config, torch_dtype=torch.bfloat16)
expertmerger.compose()
expertmerger.save_checkpoint(model_id)

注：LoRAでマージ候補のエキスパートを微調整する場合、expert_nameはadapter_で始まる。

「完全に微調整された LLM の混合」セクションで定義されているように、マージされたエキスパートはさらに微調整できます。

結び

以上のように、mergooでは、複数の専門家の知識を確実かつ透過的に統合することができます。mergooはmixture-of-expert、mixture-of-adapters (MoE-LoRA)、layer-wise mergingを含むいくつかの統合技術をサポートします。マージされたLLMは、信頼できるエキスパートを提供するために、下流のタスク上でさらに微調整することができます。

 

github.com