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

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 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.

## introduction program synthesis (or code generation) is a longstanding problem explored since the early days of computer science @xcite . recently, instruction tuning of code large language models (llms) has been used to improve many coding tasks @xcite @xcite , such as text-to-code generation @xcite , code completion @xcite , and data science engineering @xcite . a typical instruction tuning flow involves two steps @xcite tion dataset of instruction-output pairs, where the instruction reflects human intents in natural language and the output includes target code snippets that correspond to the intent; and (ii) supervised fine-tuning of pre-trained llm on the instruction dataset. in the realm of code instruction tuning, most recent works have been focusing on curating high-quality instruction datasets. for example, code evol-instruct @xcite uses chatgpt to obtain complex synthetic code instructions with heuristic prompts, while oss-instruct @xcite prompts chatgpt to generate new coding problems by drawing inspiration from open source code snippets. since existing works focus on the data perspectives of instruction tuning, they all follow the standard sft, leaving room for exploring advanced training schemes. we argue that prior works largely overlook the possibility of improving code instruction tuning by advancing existing training schemes. figure 1 depicts supervised fine-tuning (sft), which directly uses the pre-trained weights and architecture for fine-tuning. the model is dense here because all parameters are activated to predict the next token (assuming it is a decoder-only llm). in contrast to fine-tuning a dense model, following the scaling laws @xcite (i.e., more parameters, better performance), sparse upcycling @xcite is proposed to efficiently upgrade the model size by upcycling a dense llm to a sparsely activated @xcite model. an moe model is efficient because its prediction of the next token only invokes a subset of parameters (i.e., experts) and thus is sparsely activated. for example, @xcite , compared to a dense 7b model, uses approximately 8× parameters and 2× computation, i.e., only 2 out of 8 experts are dynamically selected to compute the next token. however, there are two key limitations when using sparse upcycling in instruction tuning: (i) slow scaling: it is reported that sparse upcycling improves the performance of dense models marginally with limited training steps, requiring orders of magnitude of extra compute to achieve decent improvement @xcite ; (ii) inference cost: although moe is more efficient than directly scaling up the size of dense llms, moe is still expensive, especially at inference, as it introduces significantly more parameters (i.e., memory) and computes during inference, compared to its dense counterparts. in this paper, we propose x ft: by simply merging upcycled moe models, we push the performance limit of instruction-tuned code llms. while vanilla sparse upcycling fails to improve instruction tuning efficiently @xcite , x ft addresses this challenge by isolating one expert as the shared expert among all the other experts in each moe layer, inspired by deepseek-moe @xcite and mocle @xcite . x ft also proposes a novel routing weight normalization strategy to eliminate scale mismatch between the upcycled moe layer with the shared expert and the original dense layer, which will otherwise lead to performance degradation @xcite . after the upcycled moe model finishes its sft phase, motivated by model soups @xcite , x ft uses a learnable model merging mechanism to output a dense model by merging all the expert networks in the upcycled moe, i.e., the final dense model is of the same model structure and size as the original pre-trained model, achieving similar performance without paying extra inference cost as the sparse upcycling. with only 1.3b parameters, x ft achieves 67.1 pass@1 on humaneval and 64.6 pass@1 on hu-maneval+, which is the new state-of-the-art for tiny code llms (<3b). compared with sft, x ft achieves 13% improvement on humaneval+. surprisingly, our learnable merging mechanism can preserve or even further boost the performance of the upcycled moe with only around 1 /8× parameters! we conclude our contribution as follows: ## x ft we describe the details of x ft in this section. there are two steps in our framework: upcycling (section 3.1) and merging (section 3.2). during upcycling, we construct an mixture-of-experts (moe) model from the pre-trained dense model, namely moe ds , which is then fine-tuned on coding instruc- layer norm attention layer norm mlp layer norm attention layer norm layer norm attention layer norm mlp (merged) original dense block upcycled moe block merged dense block router shared mlp mlp 2 mlp e mlp 1 normalized weighted sum moe upcycling mlp layers merging with learned weights instruction tuning … before training post training tion data. for merging, we propose a learnable model merging method to convert the instructiontuned moe ds back to a normal dense model by merging each moe layer into an ffn layer through weight averaging while directly copying other remaining layers. consequently, we can obtain x ft ds that has the same model architecture and size as the original pre-trained dense model, which eliminates all the additional inference overhead brought by the original sparse upcycling, while preserving or even improving the performance of moe ds . our framework is illustrated in figure 2 . ## conclusion this paper introduces x ft to unlock the power of code instruction tuning by simply merging upcycled moe. similar to sft, x ft starts with a dense llm and outputs a fine-tuned dense llm with the exact size and model structure. yet, x ft improves sft by upcycling the pre-trained dense llm to an moe model for fine-tuning, after which we compile the moe model back to an efficient dense llm with a learnable merging mechanism. as such, we unleash the performance limit of instruction tuning without any additional inference overhead. using the same training dataset, x ft outperforms sft on a variety of benchmarks, including humaneval(+), mbpp(+), multipl-e, and ds-1000, from 2% to 13%. by applying x ft to deepseek-coder-base 1.3b, we create the new state-of-the-art tiny code llm (<3b). the final dense llm produced by x ft preserves or even outperforms the full upcycled moe which uses 8× parameters as much as our final dense llm. x ft is fully orthogonal to the existing instruction tuners such as evol-instruct and oss-instruct, opening a new dimension for code instruction tuning.