Mergoo: Efficiently Merge and Train Multiple LLM Experts

Introduction

Mergoo is a powerful Python library developed by Leeroo-AI, aimed at streamlining the process of combining multiple Large Language Model (LLM) experts and efficiently training the resulting merged model. This library is particularly useful for integrating diverse knowledge bases from different LLM experts, whether they are generic or specialized for specific domains. Mergoo offers robust support for various merging techniques, including Mixture-of-Experts (MoE), Mixture-of-Adapters (MoA), and Layer-wise merging, providing flexibility in how models are combined.

Key features of Mergoo include:

Multiple Merging Methods: Supports Mixture-of-Experts, Mixture-of-Adapters, and Layer-wise merging.
Flexible Layer Merging: Allows for flexible merging configurations for each layer.
Broad Base Model Support: Compatible with popular models like Llama (including LLaMa3), Mistral, Phi3, and BERT.
Trainer Integration: Works seamlessly with Hugging Face's Trainer, SFTrainer, and PEFT.
Device Compatibility: Supports CPU, MPS, and GPU environments.
Training Choices: Offers options to train only the Router of MoE layers or perform full fine-tuning of the merged LLM.

Installation

Getting started with Mergoo is straightforward. You can install it using pip:

pip install mergoo

For the latest unstable version directly from GitHub:

pip install git+https://github.com/Leeroo-AI/mergoo

Alternatively, you can install it from the source:

git clone https://github.com/Leeroo-AI/mergoo
cd mergoo
pip install -e .

Examples

Mergoo simplifies the configuration and merging process. Here are examples for setting up configurations for different merging scenarios.

Fully Fine-tuned Experts

This configuration demonstrates merging fully fine-tuned LLM experts, such as combining math and code-specific Mistral models.

config = {
    "model_type": "mistral",
    "num_experts_per_tok": 2,
    "experts": [
        {"expert_name": "base_expert", "model_id": "mistralai/Mistral-7B-v0.1"},
        {"expert_name": "expert_1", "model_id": "meta-math/MetaMath-Mistral-7B"},
        {"expert_name": "expert_2", "model_id": "ajibawa-2023/Code-Mistral-7B"}
    ],
    "router_layers": ["gate_proj", "up_proj", "down_proj"]
}

Mixture of Adapters (MoE on LoRA)

Mergoo can also build a routing layer on top of LoRAs, creating a Mixture of Adapters. This is ideal for merging LoRA fine-tuned LLM experts.

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"},
        {"expert_name": "adapter_4", "model_id": "predibase/customer_support_payments"}
    ],
}

After configuring, you can merge and save the checkpoint:

import torch
from mergoo.compose_experts import ComposeExperts

model_id = "data/mistral_lora_moe"
expertmerger = ComposeExperts(config, torch_dtype=torch.float16)
expertmerger.compose()
expertmerger.save_checkpoint(model_id)

Then, load and fine-tune the merged model using Hugging Face Trainer:

from transformers import Trainer
from mergoo.models.modeling_mistral import MistralForCausalLM

model = MistralForCausalLM.from_pretrained("data/mistral_lora_moe") 
# NOTE: 'gate' / router layers are untrained hence weight loading warning would appeare for them

trainer = Trainer( ... )
trainer.train()

Why Use Mergoo?

Mergoo addresses the growing need to combine specialized LLMs effectively. By providing a flexible and efficient framework for merging, it allows developers and researchers to:

Leverage Diverse Expertise: Integrate knowledge from multiple domain-specific or task-specific LLMs into a single, more capable model.
Improve Model Performance: Potentially enhance performance on complex tasks by routing inputs to the most relevant expert or combining their strengths.
Reduce Training Costs: Efficiently train merged models, focusing on router layers or specific components, rather than retraining entire large models from scratch.
Simplify Complex Workflows: Abstract away the complexities of implementing Mixture-of-Experts or Mixture-of-Adapters architectures.
Stay Current: Support for the latest LLMs like Llama3 and integration with Hugging Face ecosystem ensures compatibility with modern AI practices.