Back in March 2024, Sakana.ai released a paper on Evolutionary Model Merging and the subsequent model and eval results of this game-changing merge method. Unfortunately for the community, they never released the algorithm behind their results.

Since this release, we've been fully focused on developing this groundbreaking technique to share with the community. We're now excited to announce the launch of this state-of-the-art functionality in MergeKit, our toolkit for merging Large Language Models (LLMs).

MergeKit lets people target specific competencies or qualities in their merges. Without it, model merging is an extremely manual exploratory process–trying dozens of merges, manually evaluating them, and trying to come up with a mental framework that explains how the merging parameters are related to the performance of the final model. With MergeKit, we can specify what qualities we want a model to have, and optimization will take care of it for us.

How to get started with MergeKit

I've created a tutorial to help you get started. You can find additional resources on GitHub: https://github.com/arcee-ai/mergekit

Hardware Requirements

To run MergeKit, you need at least one GPU. You need to be able to inference a model in FP16. If you're working with models in the 7B size range, 24GB of VRAM will do just fine. If you're a big spender then you can use a Ray cluster with however many GPUs you want. For this demo I'm using a RunPod instance with a single NVIDIA A100.

Installing

First let's set up our environment with an installation of mergekit. We need to use the feature flag, and I'm using vllm as well because it's faster.

Defining Tasks

To optimize a merge recipe we need to first decide what exactly to optimize. MergeKit uses EleutherAI's language model evaluation harness as a backend, so in theory all of the benchmarks supported by lm-eval can be used. Since I'm not an evil little gnome I'm going to define some custom tasks instead of directly optimizing against Open LLM Leaderboard scores.

Let's say that we want some spatial awareness in our model. The spartqa-mchoice dataset is a set of synthetic question-and-answer pairs involving the arrangement of objects that aims to test the spatial reasoning capabilities of language models. Let's take a random sampling of their training split and use that for one part of our scoring.

Now we need to define a lm eval task that scores against this data. This can be done by writing a YAML file (and any necessary helper code). For more details on how to do this look at the New Task Guide.

One common problem with merges is that the result often doesn't conform to any one particular prompting style. When manually creating merge recipes it's fairly easy to get the behavior you want by varying weights across layers, but since we're letting an algorithm optimize things let's make a silly little task for it instead. Alpaca is a very common standard and all it really needs from the model is to correctly output an EOS token after a completed response.

First, let's put together another tiny set of data for evaluating our metric with. I'll use a few hundred prompts from vicgalle/alpaca-gpt4

And now the actual task definition is quite simple:

Of course, there are more robust ways to evaluate but the multiple choice setup is nice in that it evaluates really quickly. Experiment at will!

Writing an Evolutionary Merge Config

We now have all the parts in place needed to actually define the merge we want to optimize. MergeKit takes a YAML configuration file that defines what models we want to include, what merge method to use, and what tasks to optimize against.

For this example, I'm going to throw three models into the soup:

• Hermes 2 Pro Mistral 7B, because it's a generally good model
• Dan's Adventurous Winds Mk2 7B, because it's a really fun model but answers to no prompt format
• Zephyr 7B beta for its quality instruction following

Most of the methods implemented by MergeKit can be used. I arbitrarily chose Task Arithmetic arbitrarily and I made the spartqa_train task slightly more important than alpaca_prompt_format purely as an example.

Running the Merge

Now we finally have all the pieces set up to actually run Merge Kit. Here's the command I used:

This will kick off the process and start merging and evaluating models. If you used WandB option then metrics on the evaluated models will be reported to Weights & Biases. This can be useful to obsessively refresh while you should be doing other things.

By default, MergeKit will keep going until it has evaluated over 100 merges (or you stop it with CTRL+C). You can increase this limit by passing the max-fevalsargument. Once the script has terminated, the mergekit configuration for the best-scoring merge will be written to /workspace/evol_merge_storage/best_config.yaml. You can get your final model by running it through mergekit-yaml</span>, like so:

We are thrilled this novel merging technique is now available to everyone through MergeKit. We will also be integrating evolutionary model merging into our Arcee models.

Let us know how you're using MergeKit and happy merging!