Multi-task Code LLMs: Data Mix or Model Merge?

This repository contains the implementation for the paper "Multi-task Code LLMs: Data Mix or Model Merge?" accepted by the third international workshop on large language models for code.

Overview

This research explores two fundamental approaches for creating multi-task code large language models:

1. Data Mixing (DM): Training models on mixed multi-task datasets
2. Model Merging (MM): Training task-specific models and merging them post-hoc

We provide comprehensive implementations, 28 pre-trained model checkpoints, and evaluation frameworks to reproduce our results.

• Model Checkpoints
• Repository Structure

Model Checkpoints

All 28 chekpoints from the experiments reported in Table 3 of the paper can be accessed from Google Drive using the links in the file checkpoints.csv

Each models is zipped with the following naming convention:

MODEL_NAME_SIZE_TRAINING_TASK.zip

For example

Deepseek_Coder_1_3B_SFT_Code_Gen.zip

contains the DeepseekCoder 1.3B model which was fine-tuned for the Code Generation task.

Repository Structure

├── llama_factory/          # SFT training framework (based on LLaMA-Factory)
│   ├── my_yaml/           # Training configuration files
│   │   ├── 2b/            # 2B parameter model configs
│   │   ├── 8b/            # 8B parameter model configs  
│   │   └── 20b/           # 20B parameter model configs
│   └── job_sh/            # Training job scripts
├── mergekit/              # Model merging toolkit and scripts
│   ├── job_sh/            # Merging job scripts
│   └── my_yaml/           # Merging configuration files
├── eval_hm/               # Evaluation framework
│   ├── job_sh/            # Evaluation job scripts
│   ├── eval_codesum.py    # Code summarization evaluation
│   └── json_to_markdowns.py  # Result processing
└── README.md              # This file

Quick Start

Prerequisites

• Python 3.10+
• CUDA-capable GPU
• HuggingFace account with access tokens
• Weights & Biases account (optional)

Environment Setup

1. Clone the repository:

git clone <repository-url>
cd Model_Merging_Data_Mixture

2. Install LLaMA-Factory:

cd llama_factory
pip install -e .
# OR install requirements
pip install -r requirements.txt

3. Install MergeKit:

cd mergekit
pip install -e .

4. Set up environment variables:

export HF_HOME="/path/to/your/huggingface/cache"
export HF_TOKEN="your_huggingface_token_here"
export TOKENIZERS_PARALLELISM="false"

# Optional: for wandb logging
wandb login your_wandb_token_here

Training Models (SFT)

Configuration Files

Training configurations are organized in llama_factory/my_yaml/ by model size:

• 2B models: 2b/ct/, 2b/kod/, 2b/gs_dm/
• 8B models: 8b/kod/, 8b/gs_dm/
• 20B models: 20b/

Example Training Commands

1. Single-task training (Code Summarization):

cd llama_factory
llamafactory-cli train my_yaml/2b/ct/dscoder_full_sftct.yaml

2. Single-task training (Code Generation):

llamafactory-cli train my_yaml/2b/kod/dscoder_full_sft4o.yaml

3. Multi-task training (Data Mixing):

llamafactory-cli train my_yaml/2b/gs_dm/dscoder_full_sftdm.yaml

Using Job Scripts

For batch job submission:

cd llama_factory/job_sh
# Edit paths in the script before running
bash run_sft_kod_qwc2515.sh

Important: Update all file paths in job scripts to match your environment before running.

Model Merging

Merging Methods

We implement several state-of-the-art merging techniques:

• Linear: Simple weighted averaging
• TIES: Task Interference Elimination via Sparse merging
• DARE: Drop And REscale for efficient merging
• DELLA: Depth-Localized LoRA for efficient merging

Configuration Files

Merging configurations are in mergekit/my_yaml/ organized by method and model combinations.

Example Merging Commands

1. TIES Merging:

cd mergekit
mergekit-yaml my_yaml/merge_ties_config.yaml ./output/merged_model

2. DARE Merging:

mergekit-yaml my_yaml/merge_dare_config.yaml ./output/merged_model

Using Merging Scripts

cd mergekit/job_sh
# Example: TIES merging for Qwen2.5-Coder-7B models
bash merge_ties_qwc7.sh

Evaluation

Supported Benchmarks

1. Code Summarization: Evaluate on code-to-docstring generation
2. Code Generation: Evaluate on programming problem solving (HumanEval, MBPP)
3. Code Summarization: Evaluate on cross-language code Summarization

Running Evaluations

1. Code Summarization:

cd eval_hm
python eval_codesum.py \
    --model /path/to/your/model \
    --data /path/to/test/data.jsonl \
    --output results.json

2. Code Generation (HumanEval):

cd eval_hm/job_sh/g_eval
bash evalplus_sft_model_name.sh

3. Using evaluation job scripts:

cd eval_hm
bash sub_s_e.sh  # Submit code summarization evaluation jobs  
bash sub_g_e.sh  # Submit code generation evaluation jobs

Processing Results

Convert summarization evaluation results to markdown format:

python json_to_markdowns.py \
    --input results/ \
    --output reports/ \
    --prefix experiment_

Model Checkpoint Breakdown

Our 28 pre-trained models include:

1. Single-task SFT Checkpoints (8 models):

   • 4 SFT checkpoints for code generation (one for each base model)
   • 4 SFT checkpoints for code summarization (one for each base model)

2. Multi-task Data Mixing Checkpoints (4 models):

   • 4 data mixture SFT checkpoints (one for each base model)

3. Model Merging Checkpoints (16 models):

   • 4 base model combinations × 4 merging methods (TIES, DARE, DELLA, Linear)
   • Each merging method applied to all 4 base model pairs

Download Models

Download all 28 pre-trained models from Google Drive:

Download the checkpoints from Google Drive using the links in the file checkpoints.csv

Important: Please modify the model paths in configuration files and scripts according to your download location. Update all references to model paths in:

• Job scripts (llama_factory/job_sh/, mergekit/job_sh/, eval_hm/job_sh/)
• Evaluation scripts

Customization

Adding New Tasks

1. Create dataset configuration in llama_factory/data/
2. Add training YAML in appropriate my_yaml/ subdirectory
3. Create corresponding job script in job_sh/

Adding New Merging Methods

1. Implement in MergeKit (if not already available)
2. Create configuration YAML in mergekit/my_yaml/
3. Add job script in mergekit/job_sh/

Custom Evaluation

1. Add evaluation script in eval_hm/
2. Create job script in eval_hm/job_sh/
3. Update result processing in json_to_markdowns.py

Reproducing Paper Results

Full Experimental Pipeline

1. Train single-task models:

# Code Summarization models
bash llama_factory/job_sh/ct/run_sft_ct_*.sh

# Code Generation models  
bash llama_factory/job_sh/kod/run_sft_kod_*.sh

2. Train data mixing models:

bash llama_factory/job_sh/gs_dm/run_sft_dm_*.sh

3. Merge single-task models:

bash mergekit/job_sh/merge_ties_*.sh
bash mergekit/job_sh/merge_dare_*.sh
bash mergekit/job_sh/merge_della_*.sh

4. Evaluate all models:

bash eval_hm/sub_s_e.sh  # Code summarization
bash eval_hm/sub_g_e.sh  # Code generation

Expected Results

Our experiments show that:

• Data Mixing generally outperforms model merging for 2B models
• Larger models benefit more from model merging strategies

Citation

If you use this code or our models, please cite our paper:

@article{zhu2026multi,
  title={Multi-task Code LLMs: Data Mix or Model Merge?},
  author={Zhu, Mingzhi and Sobolev, Boris and Krishna, Rahul and Pavuluri, Raju and Patterson, Stacy and Merler, Michele},
  journal={arXiv preprint arXiv:2601.21115},
  year={2026}
}

Contributing

This repository is released for research purposes. For questions or issues, please open a GitHub issue.

License

This project is licensed under the Apache License 2.0. See LICENSE for details.