PaperlayerView exact revision on arXiv

Source-linked AI summary

LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training

Tong Zhu, Xiaoye Qu, Daize Dong, Jiacheng Ruan, Jingqi Tong, Conghui He, Yu Cheng

arXiv:2406.16554v1cs.CL

TL;DR

The paper addresses the difficulty of building scalable, efficient MoE models without training them from scratch. It constructs LLaMA-MoE by splitting existing LLaMA-2 FFNs into experts and continually pre-training the transformed model, achieving strong performance with limited activated parameters.

  • Problem

    Training large-scale MoE models from scratch is data-hungry and unstable, while dense-model scaling incurs substantial computational costs.

  • Method

    The paper constructs experts by splitting FFNs in pre-trained LLaMA-2 and uses continual pre-training with additional gate networks and studied data-sampling strategies.

  • Results

    LLaMA-MoE-3.5B significantly outperforms open-source models with similar activation parameters, while LLaMA-MoE-3.0B achieves similar performance to Open-LLaMA-3B with fewer activated parameters.

  • Takeaways & Limitations

    Existing dense decoder-style LLMs can be transformed into effective MoE models through expert construction and continual pre-training under limited training budgets.

Abstract

from arXiv · show

Mixture-of-Experts (MoE) has gained increasing popularity as a promising framework for scaling up large language models (LLMs). However, training MoE from scratch in a large-scale setting still suffers from data-hungry and instability problems. Motivated by this limit, we investigate building MoE models from existing dense large language models. Specifically, based on the well-known LLaMA-2 7B model, we obtain an MoE model by: (1) Expert Construction, which partitions the parameters of original Feed-Forward Networks (FFNs) into multiple experts; (2) Continual Pre-training, which further trains the transformed MoE model and additional gate networks. In this paper, we comprehensively explore different methods for expert construction and various data sampling strategies for continual pre-training. After these stages, our LLaMA-MoE models could maintain language abilities and route the input tokens to specific experts with part of the parameters activated. Empirically, by training 200B tokens, LLaMA-MoE-3.5B models significantly outperform dense models that contain similar activation parameters. The source codes and models are available at https://github.com/pjlab-sys4nlp/llama-moe .

1 Introduction

The paper builds sparse MoE models from existing dense LLMs to reduce training costs while addressing expert construction and language-ability recovery.

  • Training MoEs from scratch requires a significant overall budget, motivating construction from existing dense LLMs.
  • The paper targets full MoE models based on LLaMA, with an MoE block in every transformer layer.
  • The central challenges are constructing experts from decoder-only SwiGLU FFNs and recovering performance after changing dense networks to sparse ones.
  • The study compares four expert-construction methods and dynamic or static data-sampling strategies during continual pre-training.
  • With static domain proportions matched to activated parameters, LLaMA-MoE models converge to a decent level using 200B tokens.
  • The framework splits FFNs into experts, continually pre-trains the transformed model, and evaluates the resulting models across varied tasks.

2 Related Work

Prior work establishes sparse MoE as a way to activate only part of a model, while LLaMA-MoE adapts this idea by decomposing decoder-only LLaMA FFNs into smaller experts.

  • Dense models process every input with all parameters, whereas sparse MoE activates a subset of parameters for each input.
  • MoE layers use multiple expert networks and a gating network to route inputs through selected experts.
  • One prior construction line splits FFN parameters while preserving the total parameter count, including importance-based and activation-pattern approaches.
  • Sparse upcycling instead creates a larger sparse MoE by copying the original dense MLP into identical experts.
  • LLaMA-MoE follows the parameter-splitting line for SwiGLU-based decoder-only models.

3 Preliminary

A standard sparse MoE layer routes each input embedding through a fixed small number of experts selected by a gating network.

  • A standard MoE layer contains N expert networks and a gating network that activates the top-k experts for each input.
  • In LLaMA-MoE, token-level noisy top-k gating with load balancing determines which experts accept each token.

4 Methodology

LLaMA-MoE partitions LLaMA-2-7B FFNs into experts, routes tokens sparsely, and continually pre-trains the reorganized model using alternative construction and data-sampling strategies.

  • Framework: LLaMA-MoE is built by partitioning LLaMA-2-7B FFNs into experts, routing each token to top-k experts, and continually pre-training the transformed model.
  • Expert construction: Each LLaMA FFN uses SwiGLU with up, gate, and down projection weights.
  • Expert construction: Experts are implemented as feed-forward layers using selected intermediate-neuron index sets.
  • Expert construction: Expert construction is divided into Neuron-Independent and Neuron-Sharing methods according to whether intermediate neurons are shared.
  • Expert construction: Neuron-Independent construction uniformly partitions all intermediate-neuron indices into n equal-sized sets.
  • Expert construction: IndependentRandom randomly partitions the intermediate-neuron indices, while IndependentClustering uses balanced k-means on Wup row vectors.
  • Expert construction: Neuron-Sharing uses importance vectors to select neurons, with SharingInner selecting high-importance neurons and SharingInter isolating neurons shared by most experts.
  • Expert construction: After partitioning, output rescaling compensates for the smaller number of activated expert parameters.

5 Experiments

The experiments use SlimPajama and evaluate LLaMA-MoE across multiple downstream tasks, while tracking ARC-c and HellaSwag during training. They also compare expert-construction methods, with IndependentRandom reported as strongest.

  • Training and evaluation: SlimPajama contains 627B tokens from seven domains, including CommonCrawl, C4, Github, Wikipedia, Books, arXiv, and StackExchange.
  • Main evaluation: LLaMA-MoE-3.5B significantly outperforms publicly available models of comparable size on most downstream tasks.
  • Training analysis: ARC-c and HellaSwag are used to track model performance during pre-training, with HellaSwag typically providing smoother progress.
  • Training and evaluation: The evaluation covers zero-, five-, ten-, and twenty-five-shot accuracy across ARC, LAMBADA, LogiQA, PIQA, SciQ, WinoGrande, HellaSwag, and MMLU.
  • Construction analysis: IndependentRandom obtains the best result among four expert-construction methods.

5.3 Experiment Settings

The models are initialized from LLaMA-2 7B and trained with continual pre-training under a fixed large-scale configuration. After 200B tokens, LLaMA-MoE-3.5B outperforms comparable-activation dense models.

  • Training setup: Each LLaMA-MoE variant is expected to train on 200B tokens, using 112 A100 GPUs, a 15M-token global batch, and 4096-token context length.
  • Main results: LLaMA-MoE-3.5B (4/16) exceeds Sheared-LLaMA by 1.3 average points, while LLaMA-MoE-3.0B performs comparably with Open-LLaMA-3B-v2.
  • Training progress: LLaMA-MoE-3.0B and LLaMA-MoE-3.5B converge to training losses of about 1.95 and 1.90, respectively.
  • Training progress: ARC-c and HellaSwag performance grows during training, with more fluctuations on ARC-c than on HellaSwag.

5.5 Expert Construction

The study compares expert-construction strategies and finds that IndependentRandom performs best within the token budget. Expert-output re-scaling also substantially improves MoE performance.

  • Expert construction: IndependentRandom achieves the best average score among the four expert-construction methods.
  • Expert construction: SharingInner shares approximately 50% of neurons between each expert pair, limiting its attainable performance.
  • Training duration: Models may require at least 15–20B tokens before expert-construction methods can be properly compared.
  • Output re-scaling: Expert-output re-scaling provides significantly better performance for MoE models.
  • Sampling comparison: StaticSheared achieves the best performance among four data-sampling strategies, despite not having the lowest training loss.
  • Sampling comparison: StaticSheared and StaticLLaMA keep sampling weights fixed, whereas DynamicUniform and DynamicLLaMA vary domain importance during training.

5.6 Data Sampling Weights

StaticSheared gives the strongest downstream performance among the tested sampling strategies, even though its training loss is higher. Fluency filtering improves results, while advertisement filtering hurts downstream performance.

  • Sampling strategies: StaticSheared surpasses the other sampling methods within the token budget, while dynamic sampling weights perform worse than static weights.
  • Sampling strategies: StaticSheared has higher training loss than other methods, indicating that continual pre-training loss may be less relevant to downstream performance.
  • Sampling strategies: DynamicUniform reduces loss quickly but exhibits instability and many fluctuations.
  • Data filtering: Fluency filtering produces lower training loss and improves average score, whereas advertisement filtering lowers downstream-task performance.
  • Data filtering: The final model uses a fluency-filtered dataset created by removing part of the original data rather than introducing new datasets.
  • Data filtering: The filtering comparison uses IndependentRandom with StaticSheared sampling weights as its baseline.

5.8 LLaMA-MoE vs. Training from Scratch

Building MoE models from dense models offers significant advantages over training from scratch, while routing analyses reveal layer- and domain-specific expert preferences.

  • Constructing MoE models from dense models has significant advantages and can reduce computing resources compared with training MoE from scratch.
  • Deep layers exhibit stronger routing preferences than shallow layers, suggesting greater specialization in later layers.
  • Different experts develop distinct domain preferences in deeper layers, while some experts are shared across domains.
  • CommonCrawl and C4 have similar expert preferences, whereas GitHub aligns more closely with arXiv and StackExchange.
  • HellaSwag and ARC-c most closely match CommonCrawl and C4, while GSM-8K is more similar to arXiv in expert selection.

5.10 Instruction Tuning

Instruction tuning improves LLaMA-MoE performance over the dense model on several evaluation tasks, with additional alignment gains over a Sheared-LLaMA chat model.

  • 48.29 vs. 43.69 on ARC-C and 75.10 vs. 71.70 on HellaSwag: instructed LLaMA-MoE-3.5B (4/16) outperforms the dense model.
  • 48.95 vs. 47.41 on Open LLM Leaderboard 2 tasks: instructed LLaMA-MoE-3.5B (4/16) surpasses the dense model overall.
  • 0.66 scores on MT-Bench: LLaMA-MoE-3.5B (2/8) outperforms Sheared LLaMA-2.7B in alignment abilities.

6 Conclusion

The report builds LLaMA-MoE models from pretrained LLaMA 2 by exploring expert construction and continual-training strategies under limited budgets, with strong performance at similar activation sizes.

  • LLaMA-MoE-3.0B and LLaMA-MoE-3.5B are built from pretrained LLaMA 2 using explored expert construction and continual-training strategies.
  • LLaMA-MoE-3.5B significantly outperforms open-source models with similar activation parameters, including Sheared-LLaMA-2.7B and Open-LLaMA-3.0B.
  • LLaMA-MoE-3.0B achieves similar performance to Open-LLaMA-3B with fewer activated parameters.
  • Optimized static data-sampling weights improve results, and filtering low-fluency texts provides additional performance gains.
  • LLaMA-MoE models exhibit expert specialization, with experts showing domain preferences that can be compared across pretraining and downstream datasets.
Loading 2406.16554v1…