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

heading: 1 Introduction
text: Program synthesis (or code generation) is a longstanding problem explored since the early days of computer science (Manna and Waldinger, 1971). Recently, instruction tuning of code Large Language Models (LLMs) has been used to improve many coding tasks (Chaudhary, 2023; Luo et al., 2023; Wei et al., 2023), such as text-to-code generation (Chen et al., 2021; Austin et al., 2021), code completion (Cassano et al., 2022), and data science engineering (Lai et al., 2022).
A typical instruction tuning flow involves two steps (Zhang et al., 2023): (i) curating an instruc-
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 (Luo et al., 2023) uses ChatGPT to obtain complex synthetic code instructions with heuristic prompts, while OSSINSTRUCT (Wei et al., 2023) 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 (Kaplan et al., 2020) (i.e., more parameters, better performance), sparse upcycling (Komatsuzaki et al., 2023) is proposed to efficiently upgrade the model size by upcycling a dense LLM to a sparsely activated Mixture-of-Experts (MoE)
12941
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, Mixtral-8x7B (Jiang et al., 2024), 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 (Komatsuzaki et al., 2023); (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 XFT: 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 (Komatsuzaki et al., 2023), XFT addresses this challenge by isolating one expert as the shared expert among all the other experts in each MoE layer, inspired by DeepSeekMoE (Dai et al., 2024) and MoCLE (Gou et al., 2024). XFT 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 (Wu et al., 2022). After the upcycled MoE model finishes its SFT phase, motivated by Model Soups (Wortsman et al., 2022), XFT 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, XFT achieves 67.1 pass@1 on HumanEval and 64.6 pass@1 on HumanEval+, which is the new state-of-the-art for tiny code LLMs (<3B). Compared with SFT, XFT 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:
• Dimension: We open a new dimension of improving instruction tuning of code LLMs by advancing its training scheme, using enhanced sparse upcycling and learnable model merging mechanism, which neither changes the final model structure nor requires more training data. • Technique: We present XFT, a new training scheme for code instruction tuning. XFT involves two steps: upcycling and merging. A pretrained dense LLM is first upcycled into an MoE with the shared expert setting and then fine-tuned on the instruction dataset. We propose a novel routing weight normalization strategy to avoid the performance degradation caused by the scale mismatch problem. In addition, we introduce the first learnable mechanism for merging the upcycled MoE into a dense model, eliminating additional inference overhead while preserving or even improving the upcycled MoE performance. • Results: With only 1.3B parameters, XFT achieves 67.1 pass@1 on HumanEval and 64.6 pass@1 on HumanEval+, which is the new stateof-the-art for tiny code LLMs (<3B). Compared with normal SFT, XFT achieves a significant 13% improvement on HumanEval+! XFT also achieves consistent improvements from 2% to 13% on MBPP, MultiPL-E, and DS-1000 over SFT, demonstrating its generalizability.

heading: 2 Related Work

heading: 2.1 Mixture-of-Experts
text: Mixture-of-Experts (MoE) can efficiently scale up model sizes with only sub-linear increases in computation (Shazeer et al., 2017). Compared with the standard Transformer, MoE replaces each FeedForward Network (FFN) layer with an MoE layer, which uses N (i.e., multiple) expert networks that are structurally equivalent to the original FFN layer and uses a router that directs each input token to K out of N expert networks. Formally, for the l-th MoE layer, output hidden state hlt of the t-th input token is computed as follows (Dai et al., 2024):
hlt = N∑
i=1
(gi,tFFNi(ult)) + u l t
gi,t = { si,t si,t ∈ Topk(st,K) 0 otherwise
st = {si,t | 1 ≤ i ≤ N} si,t = Softmaxi(ult T eli)
(1)
where gi,t refers to the gate value for the i-th expert given the t-th token, FFNi(·) refers to the i-th expert, ult refers to the hidden states of the t-th token which is the input of the l-th MoE layer, si,t refers to the affinity score between the i-th expert and the t-th token, Topk(S,K) refers to a function extracting K highest scores out of S, and eli refers to the centroid of the i-th expert in the l-th MoE layer. By definition, each token will only be computed in the top K experts among all the N experts and such sparsity assures the efficiency of MoE.
Recently, many works have been proposed to scale model sizes with MoE architecture (Lepikhin et al., 2020; Du et al., 2022; Fedus et al., 2022; Jiang et al., 2024; Xue et al., 2024). While most MoE models are trained from scratch, sparse upcycling (Komatsuzaki et al., 2023) is proposed to initialize MoE models based on pre-trained dense models, which can efficiently reduce the computational costs of training MoE models, compared with training MoE models from scratch. Specifically, sparse upcycling constructs a new MoE model by initializing each expert of each MoE layer as a copy of the original FFN layer in the dense model, while directly copying the remaining layers from the dense model to the new MoE model.

heading: 2.2 Instruction Tuning
text: Instruction tuning is designed to improve the instruction-following ability of LLMs by finetuning them on the instruction datasets in a supervised fashion (Wei et al., 2022). The quality of the instruction dataset is significant for the effectiveness of instruction tuning and researchers have proposed multiple methods to improve data quality. For example, SELF-INSTRUCT (Wang et al., 2023) synthesizes high-quality instruction data by prompting a foundation LLM with carefully designed prompts. To improve SELF-INSTRUCT, Evol-Instruct (Xu et al., 2023) enhances the complexity and diversity of the instruction dataset by prompting ChatGPT with heuristic prompts. OSSINSTRUCT (Wei et al., 2023) queries ChatGPT to generate instruction-output pairs by getting inspiration from real-world code snippets.
Recently, some parameter-efficient fine-tuning techniques have been proposed to use MoE for better instruction tuning. For example, LoRAMoE (Dou et al., 2023) and MoCLE (Gou et al., 2024) propose MoE-like modules that are constructed with Low-Rank Adaptations (LoRA) to improve instruction tuning, while PESC (Wu et al.,
2024) proposes to integrate adapters into MoE that are upcycled from dense models. Unlike these works, XFT focuses on full-parameter fine-tuning, which is proven generally stronger than parameterefficient fine-tuning (Chen et al., 2022).

heading: 2.3 Weight Averaging
text: Weight averaging is a commonly used technique to improve the performance of deep learning models. For example, Model Soups (Wortsman et al., 2022) averages the weights of multiple models that are initialized from the same pre-trained model but finetuned with different hyperparameter configurations to improve the accuracy and robustness of the fine-tuned model. However, only a few works have been proposed to merge experts of an MoE layer to a normal FFN layer with weight averaging to reduce both parameter and computation overhead of inference. For example, OneS (Xue et al., 2022) proposes several simple weight averaging methods to merge expert networks of a BERT-based MoE model. Closely related to our work, Experts Weights Averaging (EWA) (Huang et al., 2023) proposes to convert an MoE model to a dense model with two steps: (i) During MoE training, EWA conducts weighted averaging of all the expert weights after each weight update of MoE, which is based on a manually-crafted hyperparameter β; (ii) After training, EWA converts each MoE layer into an FFN layer by uniformly averaging the experts.
Different from all the aforementioned existing works, XFT is the first work proposing a learnable mechanism to merge expert networks in the upcycled MoE model. Furthermore, while the training scheme of EWA is deeply coupled to a specific MoE architecture, XFT can be easily adapted to different MoE architectures by only adjusting the final merging process. In addition, unlike EWA, XFT does not introduce any hyperparameters into the training of large MoE models, significantly reducing the computational resources for hyperparameter searching. Our empirical results in Section 4 also demonstrate the clear advantage of XFT.
3 XFT We describe the details of XFT 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 MoEDS, which is then fine-tuned on coding instruc-
tion data. For merging, we propose a learnable model merging method to convert the instructiontuned MoEDS 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 XFTDS 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 MoEDS. Our framework is illustrated in Figure 2.

heading: 3.1 Upcycling
text: Inspired by sparse upcycling (Komatsuzaki et al., 2023), we convert the pre-trained dense LLM to a new MoE by initializing each expert of each MoE layer as a copy of the original FFN layer in the dense model, while directly copying the remaining layers from the dense model to the new MoE model. However, the performance gain brought by sparse upcycling is negligible with a limited training budget (Komatsuzaki et al., 2023) – which is exactly the situation we are facing during instruction tuning. Intuitively, it is because each expert in the upcycled MoE model is trained on fewer instruction data than the original dense model does because traditional routers used in sparse upcycling will assign different tokens to different experts and thus reduce the amount of data each expert is trained on (Gou et al., 2024). Consequently, inspired by DeepSeekMoE (Dai et al., 2024) and MoCLE (Gou et al., 2024), XFT introduces the shared expert setting into sparse upcycling to tackle this challenge. We further propose a novel routing weight normalization strategy for XFT to avoid the potential perfor-
mance degradation caused by the scale mismatch problem (Wu et al., 2022).

heading: 3.1.1 Shared Expert for Upcycling
text: During upcycling, we isolate one shared expert among all the other normal experts in each MoE layer, where the shared expert will be deterministically assigned to handle all the tokens while other normal experts are assigned by the router. By doing so, the upcycled MoE model can achieve a clear performance boost in instruction tuning, where the shared expert learns general knowledge across the whole instruction dataset while other normal experts learn specific knowledge among different instructions assigned by the router. Formally, the output hidden state hlt of the l-th MoE layer when processing the t-th token can be expressed as:
hlt = N∑
i=1
(gi,tFFNi(ult)) + u l t
gi,t =    1− stmax i = 1 Softmaxi(si,t) · stmax si,t ∈ StK 0 otherwise
StK = Topk({si,t | 1 ≤ i ≤ N},K − 1) stmax = max({si,t | 1 ≤ i ≤ N})
si,t = { −∞ i = 1 Softmaxi(ult
T eli) i ≥ 2 (2)
where N refers to the total number of experts, K refers to the number of activated experts, gi,t refers to the gate value for the i-th expert given the t-th token, FFNi(·) refers to the i-th expert, ult refers to the output hidden state of the l-th attention layer
given the t-th token (which is also the input of the l-th MoE layer), si,t refers to the affinity score between the i-th expert and the t-th token, stmax refers to the maximum affinity score among all the experts besides the shared expert, Topk(S,K) refers to a function extracting K highest scores out of S, StK refers to a set of K − 1 highest affinity scores among all the experts besides the shared expert, and eli refers to the centroid of the i-th expert in the l-th MoE layer.
FFN1 is chosen as the shared expert in each MoE layer and each token will be assigned to top K experts including one shared expert and K − 1 other normal experts. Compared with the original sparse upcycling, there are two major differences:
• Weighted Shared Expert. Following MoCLE (Gou et al., 2024), with the token-to-expert affinity score si,t, we get the maximum affinity score stmax and use its complement 1 − stmax as the routing weight of the shared expert.
• Routing Weight Normalization. Although the shared expert setting is also used in recent works (Dai et al., 2024; Gou et al., 2024), we cannot directly follow their routing strategy because they cannot handle a scale mismatch problem that is unique for sparse upcycling. The scale mismatch problem is that differences between the scale of the output of the upcycled MoE layer and that of the original FFN layer can cause performance degradation (Wu et al., 2022). To handle this problem, we need to ensure the sum of gi,t equals 1, so that the output of the MoE layer matches that of the FFN layer in scale. To do so, we normalize the affinity scores of top K − 1 normal experts with Softmax and scale their sum to stmax to make sure that the sum of the gi,t of top K experts, including one shared expert and K − 1 normal experts, equals 1.

heading: 3.2 Merging
text: We propose a learnable model merging method to convert the large MoE model, namely MoEDS, back to a dense model XFTDS. By doing so, we expect XFTDS to keep the boosted performance gained during upcycling while keeping its model size the same as the original dense model size to avoid any additional inference overhead. Inspired by Model Soups (Wortsman et al., 2022), we choose to merge MoEDS by learning the mixing coefficients that can be used to average the parameters of all experts in each MoE layer to obtain a normal FFN layer,
while directly copying other remaining layers. Formally speaking, given the weights of N experts at the l-th layer W l1,W l 2, · · · ,W lN , the process of merging each MoE layer to an FFN layer can be stated below:
W l =
N∑
i=1
αliW l i (3)
where W l denotes the merged parameter of all N experts and αli denotes the learnable mixing coefficient of expert W li . Mixing coefficients α is parameterized as the output of a softmax, ensuring that αli is positive and ∑N i=1 α l i = 1. Given input x, we denote the output of a neural network with parameters θ as f(x; θ). For loss L and instruction dataset {(xi, yi)}mi=1, such mixing coefficients α of all the L layers can be learned via:
argmin α
m∑
j=1
L(f(xj ; θo, (W l)1:L), yi) (4)
where θo refers to all the remaining layers of MoEDS other than MoE layers.
While the learning process defined in Eq. (4) is the most intuitive way of learning α, our preliminary experiment shows that, due to the shared expert setting, it tends to increase the mixing coefficient of the shared expert at each layer as much as possible to decrease the loss. It is not helpful because, although the shared expert has learned general knowledge across the whole instruction dataset and needs a relatively large mixing coefficient, we still need to keep the scale of the mixing coefficient of other normal experts at a certain level to keep the specific knowledge learned by other normal experts in the merged parameter W l.
To solve this issue, we introduce a shared expert rate λ to fix the mixing coefficient of the shared expert and learn the mixing coefficients of the remaining normal experts which sums to 1 − λ in each layer. By doing so, we can easily control the scale of the mixing coefficient of the shared expert, while still being able to learn the optimal layer-wise mixing coefficients of other normal experts. Let’s say W l1 is the shared expert of the l-th layer, then Eq. (3) and Eq. (4) can be reformulated as follows:
W lλ = λW l 1 +
N∑
i=2
αliW l i (5)
argmin α
m∑
j=1
L(f(xj ; θo, (W lλ)1:L), yi) (6)
In practice, we uniformly initialize the mixing coefficients α of all the normal experts as 1−λ N−1 , which is then trained on the same instruction dataset used during upcycling.

heading: 4 Main Evaluation

heading: 4.1 Experimental Setup
text: Training. DeepSeek-Coder-Base 1.3B (Guo et al., 2024) is used as our main base code LLM. evol-codealpaca-v1, an open-source Evol-Instruct (Luo et al., 2023) dataset containing 110K samples, is used as our code instruction dataset. MoEDS, our MoE model upcycled from the base model, is implemented following LlamaMoE (LLaMA-MoE Team, 2023). It is constructed with 8 experts in one MoE layer and the top 6 experts1 are activated for each token, including one shared expert. As such, we denote the model size of MoEDS as 8×1.3B. Other training settings are detailed in Appendix A.1. We further obtain XFTDS by using the learned mixing coefficients to compile MoE layers inside MoEDS to normal FFN layers. Note that XFTDS is the final instructiontuned code LLM we output, while MoEDS is only an intermediate product of our XFT framework.
Baselines. To study the effectiveness of XFT, we build a baseline model, namely SFTDS, by directly performing SFT for DeepSeek-Coder-Base 1.3B on evol-codealpaca-v1. To compare XFT with EWA (Huang et al., 2023), we also implement a baseline EWADS and instruction-tune it using the same hyperparameter setting as SFTDS, which is described in Appendix A.1. More implementation details of EWADS are described in Appendix A.2. Furthermore, we incorporate multiple tiny open-source LLMs (<3B) as our baselines, including DeepSeek-Coder-Base 1.3B, DeepSeek-CoderInstruct 1.3B (Guo et al., 2024), Phi-2 2.7B, and STABLE-CODE 3B (Pinnaparaju et al., 2024).

heading: 4.2 Python Text-to-Code Generation
text: HumanEval (Chen et al., 2021) and MBPP (Austin et al., 2021) benchmarks are the two most widelyused collections of Python code generation tasks. We further employ HumanEval+ and MBPP+, which use more tests automatically generated by EvalPlus (Liu et al., 2023) for more rigorous evaluation. We leave the detailed description of HumanEval(+) and MBPP(+) in Appendix A.3.
16 is the best-performing number of activated experts per our HumanEval+ experiments using top {2, 4, 6} experts.
Table 1 shows the pass@1 results of different LLMs. XFT achieves 67.1 pass@1 on HumanEval and 64.6 pass@1 on HumanEval+, which makes it the new state-of-the-art tiny code LLM (<3B). We can also observe that XFTDS has a clear improvement over the SFTDS on both benchmarks, with 13% and 2% improvement on HumanEval+ and MBPP+ respectively. In contrast, EWADS not only underperforms XFTDS on both benchmarks, but also fails to improve SFTDS on MBPP(+). Surprisingly, XFTDS even surpasses MoEDS on HumanEval and HumanEval+, despite only using around 1/8× parameters and around 1/6× computations, which showcases the effectiveness of our simple learnable merging technique. More comprehensive experiments in Appendix A.4 demonstrate the statistical significance of the improvements brought by XFT. Furthermore, while XFT will inevitably introduce training overhead, our experiment in Appendix A.5 shows that XFT still significantly outperforms SFT using the same training budget, demonstrating the ability of XFT to unlock the power of code instruction tuning.

heading: 4.3 Multilingual Code Generation
text: We use MultiPL-E (Cassano et al., 2022), a multiprogramming benchmark that supports 18 programming languages in addition to Python, to evaluate the multilingual ability and generalizability of XFT. Following previous work (Wei et al., 2023), we choose six representative programming languages in our evaluation for their distinct language features: Java, JavaScript, C++, PHP, Swift, and Rust. Table 2 shows that, among all 1.3B models, XFTDS achieves the best average multilingual performance and performs the best on five (out of six) programming languages and overall largely improves SFTDS which uses standard SFT. Notably, the overall performance of EWADS is on par with SFTDS, indicating that EWADS fails to improve SFT on multilingual coding. Appendix A.6 further studies whether each expert in MoEDS specializes differently in different programming languages.

heading: 4.4 Code Generation for Data Science
text: The DS-1000 dataset (Lai et al., 2022) is a collection of 1000 realistic data science coding problems ranging from seven popular data science libraries in Python, including Matplotlib (plt), NumPy (np), Pandas (pd), SciPy (scp), Scikit-Learn (sk), PyTorch (py), and TensorFlow (tf). We evaluate XFT on DS-1000 to understand its effectiveness for prac-
tical data science engineering. We follow the evaluation setting of prior works (Guo et al., 2024; Wei et al., 2023). Table 3 shows that XFTDS achieves the best overall performance among all the evaluated 1.3B models. Specifically, XFTDS consistently surpasses SFTDS among all the seven studied libraries and outperforms EWADS in general.

heading: 5 Ablation Study

heading: 5.1 Effect of Shared Expert with Routing Weight Normalization
text: We demonstrate the importance of the shared expert of XFT by comparing its performance with the original sparse upcycling (Komatsuzaki et al., 2023) baseline that does not employ any shared expert. As shown in Table 4, the performance of the original sparse upcycling (with the "- Shared Expert" label) drops greatly compared with MoEDS. Notably, the sparse upcycling model performs even worse than SFTDS on HumanEval+, indicating its ineffectiveness for instruction tuning.
While the shared expert setting is also employed
in most recent works (Dai et al., 2024; Gou et al., 2024), their routing strategy will cause performance degradation due to the scale mismatch problem, which is handled by the routing weight normalization design in XFT. To demonstrate its importance, we conduct an ablation experiment by excluding it from XFT. Table 4 shows that, after removing routing weight normalization, the performance substantially decreases, despite still performing better than the original sparse upcycling that does not use the shared expert setting.

heading: 5.2 Effect of Merging Strategy
text: In this section, we demonstrate the effectiveness of our learnable merging technique by comparing it with (1) directly merging experts with initialized mixing coefficients, and (2) the learnable merging technique without the shared expert rate setting, which is the same setting as the learned soup in Model Soups (Wortsman et al., 2022) and is described in Eq. (3) and Eq. (4). Specifically, we initialize the learnable mixing coefficient of the
shared expert as 0.75 and that of the other 7 normal experts as 128 for a fair comparison. As shown in Table 5, trained mixing coefficients outperform the initialized mixing coefficients for merging. Furthermore, removing the shared rate setting will degrade the performance of XFTDS on both HumanEval and HumanEval+, demonstrating its importance. An ablation study on the shared expert rate in Appendix A.7 further shows that (1) XFTDS consistently outperforms SFTDS regardless of their shared expert rate, and (2) both the general knowledge learned in the shared expert and the specific knowledge learned in other experts are important and integral for better performance after merging.

heading: 5.3 Effect of Code LLM Choice
text: To show that the effectiveness of XFT is not dependent on any specific code LLMs, we apply XFT to STABLE-CODE 3B (Pinnaparaju et al., 2024), whose architecture is different from DeepSeekCoder-Base 1.3B (Guo et al., 2024), to study whether XFT can still improve the performance of this new model. The training settings are detailed in Appendix A.8. As shown in Table 6, XFTSTABLE significantly improves SFTSTABLE by 10% on HumanEval and 11% on HumanEval+ respectively. Furthermore, XFTSTABLE consistently boosts the performance of MoESTABLE while only using 1/4× parameters and 1/2× inference computations. These results show that the effectiveness of XFT is generalizable across different code LLMs.

heading: 6 Discussion
text: 6.1 Scaling up XFT to 7B Scale We scale up XFT to 7B-level code LLMs by applying it to DeepSeek-Coder-Base 6.7B (Guo et al., 2024). The training settings are detailed in Ap-
pendix A.9. As shown in Table 7, XFTDS-6.7B significantly improves SFTDS-6.7B by 6% on HumanEval and 9% on HumanEval+ respectively. Moreover, XFTDS-6.7B further boosts the performance of MoEDS-6.7B with only 1/8× parameters and 1/2× computations during inference! These promising results demonstrate the consistent effectiveness of XFT on 7B-level code LLMs.

heading: 6.2 Generalizability for General Tasks
text: To demonstrate that XFT can also improve the performance of LLMs on general tasks across different domains, we apply XFT to general instruction tuning. We use TinyLlama 1.1B (Zhang et al., 2024) as the base model and use evol-instruct-70k (Xu et al., 2023) as the training dataset for general instruction tuning. Following existing work (Zhang et al., 2024), we use MMLU (Hendrycks et al., 2021) with the 5-shot setting as our evaluation benchmark to evaluate the general performance of instruction-tuned LLMs. More training settings are detailed in Appendix A.10. As shown in Table 8, XFTTL improves SFTTL by 5% on MMLU in general, demonstrating the generalizable effectiveness of XFT for general instruction tuning.

heading: 6.3 Preliminary Theoretical Explanation
text: We provide a preliminary theoretical explanation of XFT by considering a simplified variant of it. Let’s start by analyzing the two major steps of XFT:
• Step 1: Upcycling. According to the scaling laws (Kaplan et al., 2020), the upcycled MoE model performs better than the normal SFT dense model due to more trainable parameters.
• Step 2: Merging. We consider a simplified variant of XFT, where the upcycled MoE model (e.g., MoEDS) can be viewed as the ensembling of two dense models and the merged dense model (e.g., XFTDS) can be viewed as the merging of the same two dense models. More details are included in Appendix A.11. As such, we can directly apply the theoretical analyzing pro-
cess in Section 4 of Model Soups (Wortsman et al., 2022) to analyze the performance difference between the upcycled MoE model and the merged dense model, which is initially designed to analyze the performance difference between model ensembling and model merging. According to the analysis (Wortsman et al., 2022), the convexity of the loss can help the merged dense model achieve a similar expected loss as that of the upcycled MoE model.
Overall, our preliminary theoretical explanation shows that (1) the Upcycling step improves the performance with more trainable parameters, and (2) the Merging step can provably maintain the performance of the aforementioned simplified MoE model with only dense-model compute.

heading: 7 Conclusion
text: This paper introduces XFT to unlock the power of code instruction tuning by simply merging upcycled MoE. Similar to SFT, XFT starts with a dense LLM and outputs a fine-tuned dense LLM with the exact size and model structure. Yet, XFT 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, XFT outperforms SFT on a variety of benchmarks, including HumanEval(+), MBPP(+), MultiPL-E, and DS-1000, from 2% to 13%. By applying XFT to DeepSeek-Coder-Base 1.3B, we create the new state-of-the-art tiny code LLM (<3B). The final dense LLM produced by XFT preserves or even outperforms the full upcycled MoE which uses 8× parameters as much as our final dense LLM. XFT is fully orthogonal to the existing instruction tuners such as Evol-Instruct and OSS-INSTRUCT, opening a new dimension for code instruction tuning.

heading: Limitations
text: To balance the general knowledge in the shared expert and the specific knowledge in other normal experts, we introduce a hyperparameter λ in the merging process of XFT, which might slightly increase the efforts for hyperparameter search. It would be interesting to explore other hyperparameter-free techniques to tackle this challenge. Furthermore, while we have provided a preliminary theoretical explanation for the strong empirical performance of XFT, it would be interesting to provide a complete theoretical explanation in the future.

heading: Acknowledgement
text: We extend our special thanks to Terry Yue Zhuo for his assistance with the scale-up experiments on DeepSeek-Coder-Base 6.7B (§6.1) after our submission. His contributions are good enough to merit authorship; however, due to the policy of ACL 2024, post-submission authorship changes are not permitted. As a result, we have included him in the author list of our arXiv version. We also thank Sea AI Lab and Dr. Qian Liu for their valuable feedback and computing resource assistance. We appreciate all the reviewers for their insightful comments. This work was partially supported by NSF grant CCF-2131943, as well as Kwai Inc.

heading: A.1 Training Settings for
text: DeepSeek-Coder-Base 1.3B
We use a batch size of 64 and a learning rate of 5e-5 with a linear scheduler to fine-tune MoEDS for 4 epochs with 500 warmup steps, following the implementation of previous work (Wei et al., 2023). We further use a batch size of 64, a shared expert rate λ of 0.75, and a learning rate of 1e-5 with a linear schedule to fine-tune the learnable mixing coefficients for experts in the instruction-tuned MoEDS on the same instruction-tuning dataset for 1 epoch with 125 warmup steps. Detailedly, we use Softmax to keep the sum of the mixing coefficients of the other 7 normal experts as 0.25. For SFTDS
and EWADS, we use the same hyperparameter setting as XFT, where the batch size is 64 and the learning rate is 5e-5 with a linear scheduler. Because XFT is trained for 4 epochs during upcycling and 1 epoch during merging, for a fair comparison, we train SFTDS and EWADS for 5 (= 4 + 1) epochs with 625 warmup steps.
A.2 Implementation details of EWA
Because EWA (Huang et al., 2023) does not release their implementation, we implement EWA by ourselves, including the constant schedule setting and the linear schedule setting. We use a share rate β of 0.3, following the original setting of EWA. While EWA with the constant schedule setting achieves reasonable performance in our evaluation, the training loss of EWA with the linear schedule setting is unstable during instruction tuning, as shown in Figure 3, and thus cannot achieve reasonable performance. As a result, we report the results of EWA with the constant schedule setting in Section 4.

heading: A.3 Details of HumanEval(+) and MBPP(+)
text: In these benchmarks, each task consists of a task description in English, which is sent to LLMs as the prompt, and LLMs are expected to generate the corresponding code to satisfy the requirements in the description. While these benchmarks provide a handful of test cases to validate the correctness of the generated code, these tests are often insufficient for more rigorous evaluation. As such, HumanEval+ and MBPP+ proposed by EvalPlus (Liu et al., 2023) are usually used to evaluate the correctness of the generated code, which provides 80×/35× more tests compared with the original benchmarks.

heading: A.4 Statistical Significance Analysis
text: In our main experiments, we follow prior works (Wei et al., 2023; Lozhkov et al., 2024) to conduct experiments on HumanEval(+) using greedy decoding. To demonstrate the statistical significance of our improvements, we change our setting from greedy decoding to sampling. In detail, to conduct one experiment on HumanEval(+), the model will sample one solution for each problem in HumanEval(+) with top p = 0.95 and temperature = 0.8, which is the same setting used in prior works (Liu et al., 2023; Chen et al., 2021).
Following prior work (Liu et al., 2023), we repeat this experiment 200 times for three techniques: XFTDS, EWADS, and SFTDS. EWADS is included because it is the best-performing baseline in our main experiment. We first compute their average pass@1 performance in these 200 experiments. As is shown in Table 9, XFTDS outperforms both EWADS and SFTDS.
Furthermore, we use the Wilcoxon signed-rank test (Wilcoxon, 1945; Dror et al., 2018), a widely used statistical test, to check if the improvements brought by XFT are statistically significant. As shown in Table 10, the p-values for both XFTDS vs. EWADS and XFTDS vs. SFTDS are much smaller than both 0.0025 (the significance level recommended for NLP work (Søgaard et al., 2014)) and 0.05 (the most common significance level), demonstrating the statistical significance of the improvements brought by XFT.

heading: A.5 Training Overhead Analysis
text: Compared with SFT, XFT will inevitably introduce additional overhead in the training process because
XFT needs to fine-tune the upcycled MoE model, which contains more parameters than the original dense model and thus requires more computation. In contrast, the normal SFT technique only needs to fine-tune the original dense model. To better understand the effect of such overhead, we conduct an experiment using the same training budget (i.e., the same GPU hours) instead of the same training steps for the normal SFT baseline. As shown in Table 11, although sharing the same training budget as XFTDS, the performance of SFTDS is still significantly worse than that of XFTDS, demonstrating the ability of XFT to unlock the power of code instruction tuning using the same training budget.

heading: A.6 Expert Specialization Analysis
text: Inspired by recent works (Jiang et al., 2024; Xue et al., 2024), we analyze whether each expert in MoEDS has different specializations in different programming languages by visualizing the routing decision of the tokens from different programming languages in the MultiPL-E benchmark (including Python). We collect the routing decision for the MultiPL-E benchmark when conducting experiments in Section 4.3. For Python, we collect the routing decision by running HumanEval experiment following the same setting used in Section 4.3. Following the analysis setting of recent work (Jiang et al., 2024), we get the visualization results from layers 0, 11, and 23 in MoEDS, where layer 0 and layer 23 are the first and the last layers of MoEDS. As shown in Figure 4, we do not observe any obvious patterns in the assignment of experts based on the type of programming languages, which is in line with the findings reported by recent works (Jiang et al., 2024; Xue et al., 2024).

heading: A.7 Effect of Shared Expert Rate
text: We further study the effect of the shared expert rate λ on the performance of the final merged dense model. We evenly choose five shared expert rates, including 0.00, 0.25, 0.50, 0.75, and 1.00, to perform the learnable merging process and evaluate each merged dense model accordingly. Note that 0.75 is the default shared expert rate used in our main experiments. If the shared expert rate is 0.00, it means that the shared expert is ignored when constructing the merged dense model from the upcycled MoE model; if the shared expert rate is 1.00, it means that the final dense model is built by simply extracting the shared expert from the upcycled MoE model. As shown in Table 12, there are mainly three interesting observations:
• The performance of the final merged dense model improves gradually when the shared expert rate grows from 0.00 to 0.75, indicating that general knowledge learned by the shared expert is important for better performance.
• The performance of the final merged dense model drops significantly when the shared expert rate grows from 0.75 to 1.00, showing that specific knowledge learned by other experts is also integral and ignoring them will lead to a
significant performance drop. • All the final merged dense models consistently
outperform the normal SFT baseline regardless of their shared expert rate, further demonstrating the effectiveness of XFT.

heading: A.8 Training Settings for STABLE-CODE 3B
text: We use evol-codealpaca-v1 as the training dataset. Since STABLE-CODE 3B is the base model, we upcycle a new MoE model from the base model, namely MoESTABLE. We construct MoESTABLE with 4 experts in one MoE layer, where the top 2 experts are activated for each token, including one shared expert. Consequently, the size of
MoESTABLE can be described as 4×3B. We use a batch size of 64 and a learning rate of 5e-5 with a linear scheduler to fine-tune MoESTABLE for 4 epochs with 500 warmup steps. Similar to XFTDS, we obtain XFTSTABLE by learning mixing coefficients to merge MoE layers inside MoESTABLE as normal FFN layers, which is fine-tuned with a batch size of 64, a shared expert rate λ of 0.85, and a learning rate of 1e-5 with a linear schedule for 1 epoch with 125 warmup steps. Our baseline model, namely SFTSTABLE, is fine-tuned for 5 (= 4 + 1) epochs with a batch size of 64, a learning rate of 5e-5, and 625 warmup steps for a fair comparison.

heading: A.9 Training Settings for
text: DeepSeek-Coder-Base 6.7B
We use evol-codealpaca-v1 as the training dataset. We upcycle a new MoE model from DeepSeek-Coder-Base 6.7B, namely MoEDS-6.7B. We construct MoEDS-6.7B with 8 experts in one MoE layer, where the top 2 experts are activated for each token, including one shared expert. As such, MoEDS-6.7B includes 8×6.7B parameters. We use a batch size of 64 and a linear scheduler to finetune MoEDS-6.7B for 4 epochs with 500 warmup steps. We choose the best-performing learning rate from {2e − 5, 5e − 5} for MoEDS-6.7B. Because the FFN weights of MoEDS-6.7B are too large to fit in our GPU memory, during our merging step, we realize that one part of computation in the training of XFTDS-6.7B has to be moved to CPUs, which significantly slows down the training speed. Consequently, we use a batch size of 16, a shared expert rate λ of 0.75, a constant learning rate of 1e-4, and 400 training steps in merging to obtain XFTDS-6.7B. Our baseline model, namely SFTDS-6.7B, is fine-tuned for 5 epochs with a batch size of 64 and 625 warmup steps for a fair comparison. We also choose the best-performing learning rate from {2e− 5, 5e− 5} for SFTDS-6.7B.

heading: A.10 Training Settings for TinyLlama 1.1B
text: Using TinyLlama 1.1B as the base model, we upcycle a new MoE model, namely MoETL, from the pre-trained dense model. Following the setting of MoEDS, we construct MoETL with 8 experts in one MoE layer, where the top 6 experts are activated for each token, including one shared expert. As such, the number of parameters for MoETL can be written as 8×1.1B. We use a batch size of 64 and a learning rate of 5e-5 with a linear scheduler to finetune MoETL for 4 epochs with 240 warmup steps.
To obtain XFTTL, we learn mixing coefficients to merge MoE layers inside MoETL by fine-tuning them with a batch size of 64, a shared expert rate λ of 0.85, and a learning rate of 2e-5 with a linear schedule for 1 epoch with 60 warmup steps. For a fair comparison, we fine-tune a baseline model SFTTL for 5 (= 4 + 1) epochs with a batch size of 64, a learning rate of 5e-5, and 300 warmup steps.

heading: A.11 Details of Preliminary Theoretical Explanation
text: We consider a simplified variant of XFT as follows: • The original dense model is a one-layer trans-
former model, which contains one attention layer connected with one feed-forward network (FFN) layer. As such, the upcycled MoE model is also a one-layer transformer model, containing one attention layer connected with an MoE layer.
• The upcycled MoE model only has two experts (e1 and e2), both of which are always selected for processing the input tokens.
• The router in the MoE model assigns constant weights to each expert, regardless of the input token. Consequently, the output of the MoE layer for the t-th token ht can be represented as (1 − α)e1(ut) + αe2(ut), where 1 − α is the router weight assigned to e1, α is the router weight assigned to e2, and ut is the input of the MoE layer for the t-th token.
• We simplify the process of merging the MoE model back to a dense model as Weα = (1 − α)We1 + αWe2 , where We refers to the weight of e and eα refers to the weight of the FFN in the merged dense model.
In this simplified scenario, if we denote f(x; θ) as the output of the model θ for the input x, the output of this simplified MoE model for input token x can be represented as f(x; θMoE). Interestingly, if we define two new dense models θ1 and θ2, where θ1 and θ2 both use the attention layer of this MoE model as their attention layer while using e1 and e2 as their FFN layer respectively, f(x; θMoE) can be represented as (1−α)f(x; θ1)+αf(x; θ2). Consequently, the computation process of this simplified MoE model can be viewed as ensembling the outputs of two dense models θ1 and θ2. Meanwhile, the process of merging the upcycled MoE model back to a dense model in this simplified scenario can be represented as θα = (1−α)θ1+αθ2, which can be viewed as the model merging of the same two dense models θ1 and θ2.

heading: Future Work
text: Furthermore, while we have provided a preliminary theoretical explanation for the strong empirical performance of XFT, it would be interesting to provide a complete theoretical explanation in the future