XFT: Unlocking the Power of Code Instruction Tuning by Simply Merging Upcycled Mixture-of-Experts

TL;DR

XFT tackles the limits of code instruction tuning by upcycling a pre-trained dense code LLM into a sparse MoE with a shared expert and a routing normalization scheme, followed by a learnable merging phase that compiles the MoE back into a dense model of the same architecture and size. This two-step process delivers substantial gains over standard SFT and prior MoE baselines on Python code generation (HumanEval/HumanEval+), multilingual coding (MultiPL-E), and data science tasks (DS-1000) with only dense-model compute during inference. The approach achieves state-of-the-art results for tiny code LLMs (<3B) at 67.1 pass@1 on HumanEval and 64.6 on HumanEval+ for a 1.3B base, while also showing strong generalization to larger models and general instruction tasks (MMLU). The findings indicate that carefully designed upcycling and a learnable merging mechanism can unlock the power of code instruction tuning without incurring extra inference overhead, and the method is orthogonal to existing instruction-tuning strategies.

Abstract

We introduce XFT, a simple yet powerful training scheme, by simply merging upcycled Mixture-of-Experts (MoE) to unleash the performance limit of instruction-tuned code Large Language Models (LLMs). While vanilla sparse upcycling fails to improve instruction tuning, XFT introduces a shared expert mechanism with a novel routing weight normalization strategy into sparse upcycling, which significantly boosts instruction tuning. After fine-tuning the upcycled MoE model, XFT introduces a learnable model merging mechanism to compile the upcycled MoE model back to a dense model, achieving upcycled MoE-level performance with only dense-model compute. By applying XFT to a 1.3B model, we create a new state-of-the-art tiny code LLM (<3B) with 67.1 and 64.6 pass@1 on HumanEval and HumanEval+ respectively. With the same data and model architecture, XFT improves supervised fine-tuning (SFT) by 13% on HumanEval+, along with consistent improvements from 2% to 13% on MBPP+, MultiPL-E, and DS-1000, demonstrating its generalizability. XFT is fully orthogonal to existing techniques such as Evol-Instruct and OSS-Instruct, opening a new dimension for improving code instruction tuning. Codes are available at https://github.com/ise-uiuc/xft.

Figures (4)

• Figure 1: Overview of SFT, sparse upcycling, and ${\mathcal{X}}$FT.
• Figure 2: Overview of ${\mathcal{X}}$FT.
• Figure 3: Training loss curve of EWA with the constant schedule setting and the linear schedule setting.
• Figure 4: Proportion of tokens assigned to each expert on different programming languages from MultiPL-E (including Python) for layers 0, 11, and 23. The shared expert $\hbox{FFN}_1$ is excluded from the chart because all the tokens are always assigned to it. The gray vertical line $\frac{1}{7}$ is the proportion expected with the uniform sampling.