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

📢 A SMALLER AFFORDABLE MoE MODEL FOR EVERYONE!!

🤗 Model Weights | 📃 Technical Report | 🚀 Quick Start
⚙️ Installation Guide | 🚧 Expert Construction | 🚅 Continual Pre-training | 💎 Evaluation

LLaMA-MoE is a series of open-sourced Mixture-of-Expert (MoE) models based on LLaMA and SlimPajama. We build LLaMA-MoE with the following two steps:

1. Partition LLaMA's FFNs into sparse experts and insert top-K gate for each layer of experts.
2. Continually pre-train the initialized MoE model with an optimized data sampling weights from Sheared LLaMA and filtered datasets from SlimPajama.

1. Lightweight Models: The number of activated model parameters is only 3.0~3.5B, which is friendly for deployment and research usage.
2. Multiple Expert Construction Methods:
   1. Neuron-Independent: Random, Clustering, Co-activation Graph, Gradient (Zhang et al., 2022, Zuo et al., 2022)
   2. Neuron-Sharing: Inner, Inter (residual)
3. Multiple MoE Gating Strategies:
   1. TopK Noisy Gate (Shazeer et al., 2017)
   2. Switch Gating (Fedus et al., 2022)
4. Fast Continual Pre-training:
   1. FlashAttention-v2 integrated (Dao, 2023)
   2. Fast streaming dataset loading
5. Abundant Monitor Items:
   1. Gate load, gate importance
   2. Loss on steps, loss on tokens, balance loss
   3. TGS (tokens/GPU/second), MFU (model FLOPs utilization)
   4. Other visualization utilities
6. Dynamic Weight Sampling:
   1. Self-defined static sampling weights
   2. Sheared LLaMA's dynamic batch loading (Xia et al., 2023)

# python>=3.10

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_dir = "llama-moe/LLaMA-MoE-v1-3_5B-2_8"
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.bfloat16, trust_remote_code=True)
model.eval()
model.to("cuda:0")

input_text = "Suzhou is famous of"
inputs = tokenizer(input_text, return_tensors="pt")
inputs = inputs.to("cuda:0")

pred = model.generate(**inputs, max_length=50, temperature=0.0)
print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))
# Suzhou is famous of its beautiful gardens. The most famous one is the Humble Administrator's Garden. It is a classical Chinese garden with a history of more than 600 years. The garden is divided into three

1. Prepare conda environment: conda create -n smoe python=3.11 (If your environment name is not smoe, you may need to change environment in launching scripts)
2. Add correct environment variables in ~/.bashrc (gcc is set to newer version for installing flash-attn). e.g.:

   export PATH=/mnt/petrelfs/share/cuda-11.8/bin:$PATH
   export LD_LIBRARY_PATH=/mnt/petrelfs/share/cuda-11.8/lib64:$LD_LIBRARY_PATH
   export PATH=/mnt/petrelfs/share/gcc-10.1.0/bin:$PATH
   export LD_LIBRARY_PATH=/mnt/petrelfs/share/gcc-10.1.0/lib64:$LD_LIBRARY_PATH

3. Take the variables into effect: source ~/.bashrc
4. Install PyTorch (CUDA-11.8): pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
5. Install dependencies: pip install -r requirements.txt
6. Install flash-attn: pip install flash-attn==2.0.1 --no-build-isolation. You may need to follow the flash-attn installation instructions to avoid some errors.
7. Install the latest Git: conda install git
8. Clone the repo: git clone [email protected]:pjlab-sys4nlp/llama-moe.git (If you don't setup the ssh key to GitHub, you may not able to clone through ssh. Check the docs about it.)
9. Change current directory: cd llama-moe
10. Install smoe in editable mode: pip install -e .[dev]
11. Setup pre-commit hooks: pre-commit install

• Neuron-Independent
  • Independent_Random: bash ./scripts/expert_construction/split/run_split_random.sh
  • Independent_Clustering: bash ./scripts/expert_construction/split/run_split_clustering.sh
• Neuron-Sharing
  • Sharing_Inner: bash ./scripts/expert_construction/split/run_split_gradient.sh
  • Sharing_Inter: bash ./scripts/expert_construction/split/run_split_gradient_residual.sh

For more information, please refer to Expert Construction docs.

Download SlimPajama into /path_to_data and put data from different domains into separate folders:

• /path_to_data/en_arxiv
• /path_to_data/en_book
• /path_to_data/en_c4
• /path_to_data/en_cc
• /path_to_data/en_stack
• /path_to_data/en_wikipedia
• /path_to_data/github

Each file should be end with *.jsonl and each line looks like:

{"id": "id-info", "content": "raw text to be tokenized"}

Run the following command to tokenize the data in each folder:

python -m smoe.utils.tokenize \
  -f jsonl \
  -t /path_to_tokenizer \
  -i /path_to_data/en_arxiv \
  -o /path_to_data_tokenized/en_arxiv

• NOTICE: Please create logs/ folder manually: mkdir -p logs
• To run the continual pre-training, please check the CPT docs.

• For evalution on Natural Questions (NQ), please refer to opencompass.
• For other tasks, please refer to lm-eval-harness.

@article{llama-moe-2023,
  title={LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training},
  author={LLaMA-MoE Team},
  journal={arXiv preprint arXiv:},
  url={https://arxiv.org/abs/},
  year={2023}
}

LLaMA-MoE Team w/ ❤️