Large Language Models (LLMs) have transformed today’s industries with their ability to understand and generate useful responses tailored to logic, code, maths, and what not. However, their utility often comes with challenges like high computational costs and limited adaptability. MergeKit, a model merging toolkit, allows developers to merge the strengths of multiple LLMs into a single, optimized model without the need for additional training or ensembling overhead. By working directly in the weight space of models, MergeKit makes it possible to combine specialized models, find trade-offs between models, and even create new capabilities.
This article will guide you through the process of merging LLMs and building a custom model tailored to your needs using MergeKit. In this guide, we will show the step-by-step process of merging pre-trained LLMs specialized in code generation into one single powerful code-generation LLM.
Key Features of MergeKit
Before we dive further, here is a quick glimpse of the features offered by MergeKit:
- It supports well-known models like Llama, Mistral, GPT-NeoX, StableLM, and more.
- You can merge using different merge methods, depending on your preference.
- GPU is preferred for smoother and faster execution, but one can also use a CPU to merge.
- Find and balance trade-offs between different models by merging them with measurable weights.
- Interpolated gradients for parameter values (inspired by Gryphe’s BlockMerge_Gradient script).
- Mixture of Experts merging
Prerequisites
The minimum system requirements for merging models is:
- GPUs: A100 80GB or H100 (for smooth execution).
- Disk Space: 200 GB
- RAM: At least 64 GB.
- Jupyter Notebook installed.
Note: The prerequisites for this are highly variable across use cases. A high-end configuration could be used for a large-scale deployment.
Step-by-step process to merge LLMs using MergeKit on cloud GPU
For the purpose of this tutorial, we’ll use a GPU-powered Virtual Machine by NodeShift since it provides high compute Virtual Machines at a very affordable cost on a scale that meets GDPR, SOC2, and ISO27001 requirements. Also, it offers an intuitive and user-friendly interface, making it easier for beginners to get started with Cloud deployments. However, feel free to use any cloud provider of your choice and follow the same steps for the rest of the tutorial.
Step 1: Setting up a NodeShift Account
Visit app.nodeshift.com and create an account by filling in basic details, or continue signing up with your Google/GitHub account.
If you already have an account, login straight to your dashboard.
Step 2: Create a GPU Node
After accessing your account, you should see a dashboard (see image), now:
- Navigate to the menu on the left side.
- Click on the GPU Nodes option.
- Click on Start to start creating your very first GPU node.
These GPU nodes are GPU-powered virtual machines by NodeShift. These nodes are highly customizable and let you control different environmental configurations for GPUs ranging from H100s to A100s, CPUs, RAM, and storage, according to your needs.
Step 3: Selecting configuration for GPU (model, region, storage)
- For this tutorial, we’ll be using H100 SXM GPU, however, you can choose any GPU of your choice as per your needs.
- Similarly, we’ll opt for 200GB storage by sliding the bar. You can also select the region where you want your GPU to reside from the available ones.
Step 4: Choose GPU Configuration and Authentication method
- After selecting your required configuration options, you’ll see the available VMs in your region and according to (or very close to) your configuration. In our case, we’ll choose a 1x H100 SXM GPU node with 128 vCPUs/129GB RAM/200 GB SSD.
2. Next, you’ll need to select an authentication method. Two methods are available: Password and SSH Key. We recommend using SSH keys, as they are a more secure option. To create one, head over to our official documentation.
Step 5: Choose an Image
The final step would be to choose an image for the VM, which in our case is Jupyter Notebook, where we’ll deploy and run the inference of our merged model.
That’s it! You are now ready to deploy the node. Finalize the configuration summary, and if it looks good, click Create to deploy the node.
Step 6: Connect to active Compute Node using SSH
- As soon as you create the node, it will be deployed in a few seconds or a minute. Once deployed, you will see a status Running in green, meaning that our Compute node is ready to use!
- Once your GPU shows this status, navigate to the three dots on the right and click on Connect with SSH, which will open a Jupyter Notebook session on a new tab where we can run our model.
Step 7: Setting up Python Notebook
Start by creating a .ipynb notebook by clicking on Python 3 (ipykernel).
Next, If you want to check the GPU details, run the following command in the Jupyter Notebook cell:
!nvidia-smi
Output:
Step 8: Install MergeKit & its dependencies
- First, it’s important to upgrade any older versions of
pip to the latest one to avoid Python version errors.
!python3 -m pip install --upgrade pip
Output:
2. Clone the MergeKit GitHub repository and install its libraries.
!cd mergekit && pip install -q -e
Output:
Step 9: Write the merge configuration in YAML file
While selecting the trained language models that you want to merge for your final LLM, make sure to choose models with mixing architectures. To achieve this, select one base model and other pre-trained models to merge.
For the purpose of this tutorial, we’ll choose:
Secondly, we’ll use the TIES (Trim, Elect Sign, & Merge) algorithm to merge our models. TIES is a merging algorithm that consolidates multiple task-specific models into a single, efficient multi-task model.
When merging multiple fine-tuned models for the same task, models merged using TIES show more robustness and improved performance than those merged using other techniques, such as Averaging, Fisher, Task Arithmetic, and even ensembling.
Here’s the YAML configuration defining the above use-case:
import yaml
MODEL_NAME = "NodeShift-7B-v0.1"
yaml_config = """
models:
- model: deepseek-ai/deepseek-coder-6.7b-instruct
parameters:
density: [1, 0.7, 0.1] # density gradient
weight: 1.0
- model: m-a-p/OpenCodeInterpreter-DS-6.7B
parameters:
density: 0.5
weight: [0, 0.3, 0.7, 1] # weight gradient
merge_method: ties
base_model: deepseek-ai/deepseek-coder-6.7b-base
parameters:
normalize: true
int8_mask: true
dtype: float16
"""
# Save config as yaml file
with open('config.yaml', 'w', encoding="utf-8") as f:
f.write(yaml_config)
Run this code block to save the configuration file.
Step 10: Merge the models & upload to Hugging Face
- Now, merge the models using MergeKit based on the parameters defined in the configuration file. Finally, we are saving the merged model files in a folder named “merge”.
# Merge models
!mergekit-yaml config.yaml merge --copy-tokenizer --allow-crimes --out-shard-size 1B --lazy-unpickle
Output:
2. Create an empty repository in your Hugging Face account to push the new model to Hugging Face.
3. Upload the newly merged model to Hugging Face.
Make sure to replace:
<HF_USERNAME> – with your Hugging Face username<HF_TOKEN> – with your Hugging Face account token with WRITE access<MODEL_NAME> – name of the model that you’ve defined above<FOLDER_NAME> – with the folder where all your merged model files are saved (e.g., “merge”).
from huggingface_hub import HfApi
username = <HF_USERNAME>
# Defined in the secrets tab in Google Colab
api = HfApi(token=<HF_TOKEN>)
api.upload_folder(
repo_id=f"{HF_USERNAME}/{MODEL_NAME}",
folder_path="<FOLDER_NAME>",
)
Output:
Step 11: Perform model inference to check responses
Finally, we’ll run the new model using the base model tokenizer and pass it a prompt to check its responses.
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-coder-6.7b-base")
model = AutoModelForCausalLM.from_pretrained("halfacupoftea/NodeShift_7b_v0.1", torch_dtype=torch.bfloat16,
device_map="auto")
input_text = "write a function in TypeScript for sorting trending posts for a social media platform based on number of likes and comments."
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_length=1024)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
Output:
Below is the output obtained after running the above inference:
You can plug and play with different prompts and enhance your output in a similar way.
Conclusion
Merging large language models with tools like MergeKit opens a new frontier in AI/ML development for creating high-performing and highly specialized models for specific use-cases without the traditional overhead of ensembling or retraining. By thoroughly following this guide, developers can quickly merge and deploy models with new capabilities while maintaining low inference costs. We have merged the model using H100 GPU by NodeShift, which offers robust computing power, ensuring seamless execution and allowing even resource-intensive merge operations to be completed efficiently.
