{"id":13754113,"url":"https://github.com/pjlab-sys4nlp/llama-moe","last_synced_at":"2025-05-09T22:30:57.730Z","repository":{"id":214003487,"uuid":"670026629","full_name":"pjlab-sys4nlp/llama-moe","owner":"pjlab-sys4nlp","description":"⛷️ LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training (EMNLP 2024)","archived":false,"fork":false,"pushed_at":"2024-12-06T04:47:07.000Z","size":1774,"stargazers_count":959,"open_issues_count":6,"forks_count":56,"subscribers_count":8,"default_branch":"main","last_synced_at":"2025-04-30T03:39:36.236Z","etag":null,"topics":["continual-pre-training","expert-partition","llama","llm","mixture-of-experts","moe"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2406.16554","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/pjlab-sys4nlp.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-07-24T06:15:51.000Z","updated_at":"2025-04-28T04:37:43.000Z","dependencies_parsed_at":"2023-12-25T04:19:31.657Z","dependency_job_id":"2090a98c-a999-43ce-8ee4-e8f4b2040d3c","html_url":"https://github.com/pjlab-sys4nlp/llama-moe","commit_stats":{"total_commits":142,"total_committers":5,"mean_commits":28.4,"dds":0.471830985915493,"last_synced_commit":"b17aff436cce58e2fcd7327789c7fadafe15d19e"},"previous_names":["pjlab-sys4nlp/llama-moe"],"tags_count":4,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pjlab-sys4nlp%2Fllama-moe","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pjlab-sys4nlp%2Fllama-moe/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pjlab-sys4nlp%2Fllama-moe/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pjlab-sys4nlp%2Fllama-moe/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/pjlab-sys4nlp","download_url":"https://codeload.github.com/pjlab-sys4nlp/llama-moe/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":253335377,"owners_count":21892663,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["continual-pre-training","expert-partition","llama","llm","mixture-of-experts","moe"],"created_at":"2024-08-03T09:01:40.681Z","updated_at":"2025-05-09T22:30:57.396Z","avatar_url":"https://github.com/pjlab-sys4nlp.png","language":"Python","funding_links":[],"categories":["A01_文本生成_文本对话","Uncategorized","NLP","🚀 Open Models"],"sub_categories":["大语言对话模型及数据","Uncategorized","3. Pretraining"],"readme":"\u003cdiv align=\"center\"\u003e\n  \u003ch1\u003eLLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training\u003c/h1\u003e\n  \u003cimg src=\"docs/imgs/title-favicon.png\" width=\"200\" alt=\"LLaMA-MoE favicon\" style=\"border-radius: 5%;\"\u003e\u003cbr /\u003e\n  \u003cspan style=\"color:red\"\u003e📢 \u003cstrong\u003e\u003ci\u003eA SMALLER AFFORDABLE MoE MODEL FOR EVERYONE!!\u003c/i\u003e\u003c/strong\u003e\u003c/span\u003e\n  \u003cdiv\u003e\n    \u003ca href=\"https://huggingface.co/llama-moe\" target=\"_blank\"\u003e🤗 Model Weights\u003c/a\u003e | \u003ca href=\"#quick-start\"\u003e🚀 Quick Start\u003c/a\u003e | \u003ca href=\"#installation\"\u003e⚙️ Installation Guide\u003c/a\u003e | \u003ca href=\"#expert-construction\"\u003e🚧 Expert Construction\u003c/a\u003e | \u003ca href=\"#continual-pretraining\"\u003e🚅 Continual Pre-training\u003c/a\u003e | \u003ca href=\"#evaluation\"\u003e💎 Evaluation\u003c/a\u003e | \u003ca href=\"#sft\"\u003e💬 Supervised Fine-Tuning (SFT)\u003c/a\u003e\n  \u003c/div\u003e\n  \u003ca href=\"docs/LLaMA_MoE.pdf\" target=\"_blank\"\u003e\u003cstrong\u003e📃 Technical Report\u003c/strong\u003e\u003c/a\u003e\n\u003c/div\u003e\n\n\u003ch2 id=\"llama-moe\"\u003e🎉 Introduction\u003c/h2\u003e\n\nLLaMA-MoE is a series of open-sourced Mixture-of-Expert (MoE) models based on [LLaMA](https://github.com/facebookresearch/llama) and [SlimPajama](https://www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama).\nWe build LLaMA-MoE with the following two steps:\n1. Partition LLaMA's FFNs into sparse experts and insert top-K gate for each layer of experts.\n2. Continually pre-train the initialized MoE model with an optimized data sampling weights from [Sheared LLaMA](https://arxiv.org/abs/2310.06694) and filtered datasets from [SlimPajama](https://www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama).\n\n![MoE Routing](./docs/imgs/MoE-Routing.gif)\n\n\u003ch2 id=\"features\"\u003e🔥 Features\u003c/h2\u003e\n\n1. **Lightweight Models**: The number of activated model parameters is only 3.0~3.5B, which is friendly for deployment and research usage.\n2. **Multiple Expert Construction Methods**:\n   1. Neuron-Independent: Random, Clustering, Co-activation Graph, Gradient ([Zhang et al., 2022](http://arxiv.org/abs/2110.01786), [Zuo et al., 2022](http://arxiv.org/abs/2204.07675))\n   2. Neuron-Sharing: Inner, Inter (residual)\n3. **Multiple MoE Gating Strategies**:\n   1. TopK Noisy Gate ([Shazeer et al., 2017](http://arxiv.org/abs/1701.06538))\n   2. Switch Gating ([Fedus et al., 2022](http://arxiv.org/abs/2101.03961))\n4. **Fast Continual Pre-training**:\n   1. FlashAttention-v2 integrated ([Dao, 2023](https://github.com/Dao-AILab/flash-attention))\n   2. Fast streaming dataset loading\n5. **Abundant Monitor Items**:\n   1. Gate load, gate importance\n   2. Loss on steps, loss on tokens, balance loss\n   3. TGS (tokens/GPU/second), MFU (model FLOPs utilization)\n   4. Other visualization utilities\n6. **Dynamic Weight Sampling**:\n   1. Self-defined static sampling weights\n   2. Sheared LLaMA's dynamic batch loading ([Xia et al., 2023](http://arxiv.org/abs/2310.06694))\n\n\n\u003ch2 id=\"quick-start\"\u003e🚀 QuickStart\u003c/h2\u003e\n\n```python\n# python\u003e=3.10\n\nimport torch\nfrom transformers import AutoTokenizer, AutoModelForCausalLM\n\nmodel_dir = \"llama-moe/LLaMA-MoE-v1-3_5B-2_8\"\ntokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)\nmodel = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.bfloat16, trust_remote_code=True)\nmodel.eval()\nmodel.to(\"cuda:0\")\n\ninput_text = \"Suzhou is famous of\"\ninputs = tokenizer(input_text, return_tensors=\"pt\")\ninputs = inputs.to(\"cuda:0\")\n\npred = model.generate(**inputs, max_length=50, temperature=0.0)\nprint(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))\n# 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\n```\n\n\u003ch2 id=\"installation\"\u003e⚙️ Installation\u003c/h2\u003e\n\n1. 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)\n2. Add correct environment variables in `~/.bashrc` (`gcc` is set to newer version for installing `flash-attn`). e.g.:\n    ```bash\n    export PATH=/mnt/petrelfs/share/cuda-11.8/bin:$PATH\n    export LD_LIBRARY_PATH=/mnt/petrelfs/share/cuda-11.8/lib64:$LD_LIBRARY_PATH\n    export PATH=/mnt/petrelfs/share/gcc-10.1.0/bin:$PATH\n    export LD_LIBRARY_PATH=/mnt/petrelfs/share/gcc-10.1.0/lib64:$LD_LIBRARY_PATH\n    ```\n3. Take the variables into effect: `source ~/.bashrc`\n4. Install PyTorch (CUDA-11.8): `pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118`\n5. Install dependencies: `pip install -r requirements.txt`\n6. Install `flash-attn`: `pip install flash-attn==2.0.1 --no-build-isolation`. You may need to follow the [flash-attn installation instructions](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#installation-and-features) to avoid some errors.\n7. Install the latest Git: `conda install git`\n8. Clone the repo: `git clone git@github.com: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](https://docs.github.com/en/authentication/connecting-to-github-with-ssh/adding-a-new-ssh-key-to-your-github-account) about it.)\n9. Change current directory: `cd llama-moe`\n10. Install `smoe` in [editable mode](https://pip.pypa.io/en/stable/cli/pip_install/#cmdoption-e): `pip install -e .[dev]`\n11. Setup `pre-commit` hooks: `pre-commit install`\n\n\u003ch2 id=\"performance\"\u003e📊 Model Performance\u003c/h2\u003e\n\n| Model                     | \\#Activated Experts | \\#Experts | \\#Activated Params |                         Foundation Model                          |                              SFT Model                               |\n| :------------------------ | :-----------------: | :-------: | :----------------: | :---------------------------------------------------------------: | :------------------------------------------------------------------: |\n| **LLaMA-MoE-3.0B**        |          2          |    16     |        3.0B        | [🤗 base](https://huggingface.co/llama-moe/LLaMA-MoE-v1-3_0B-2_16) | [🤗 SFT](https://huggingface.co/llama-moe/LLaMA-MoE-v1-3_0B-2_16-sft) |\n| **LLaMA-MoE-3.5B (4/16)** |          4          |    16     |        3.5B        | [🤗 base](https://huggingface.co/llama-moe/LLaMA-MoE-v1-3_5B-4_16) | [🤗 SFT](https://huggingface.co/llama-moe/LLaMA-MoE-v1-3_5B-4_16-sft) |\n| **LLaMA-MoE-3.5B (2/8)**  |          2          |     8     |        3.5B        | [🤗 base](https://huggingface.co/llama-moe/LLaMA-MoE-v1-3_5B-2_8)  | [🤗 SFT](https://huggingface.co/llama-moe/LLaMA-MoE-v1-3_5B-2_8-sft)  |\n\n- Foundation models\n\n| Model                                                                                 | Average  |   SciQ   |   PIQA   | WinoGrande |  ARC-e   | ARC-c (25) | HellaSwag (10) |  LogiQA  | BoolQ (32) | LAMBADA  | NQ (32)  | MMLU (5) |\n| :------------------------------------------------------------------------------------ | :------: | :------: | :------: | :--------: | :------: | :--------: | :------------: | :------: | :--------: | :------: | :------: | :------: |\n| [OPT-2.7B](https://huggingface.co/facebook/opt-2.7b)                                  |   50.3   |   78.9   |   74.8   |    60.8    |   54.4   |    34.0    |      61.4      |   25.8   |    63.3    |   63.6   |   10.7   |   25.8   |\n| [Pythia-2.8B](https://huggingface.co/EleutherAI/pythia-2.8b)                          |   51.5   |   83.2   |   73.6   |    59.6    |   58.8   |    36.7    |      60.7      |   28.1   |    65.9    |   64.6   |   8.7    |   26.8   |\n| [INCITE-BASE-3B](https://huggingface.co/togethercomputer/RedPajama-INCITE-Base-3B-v1) |   53.7   |   85.6   |   73.9   |    63.5    |   61.7   |    40.3    |      64.7      |   27.5   |    65.8    |   65.4   |   15.2   |   27.2   |\n| [Open-LLaMA-3B-v2](https://huggingface.co/openlm-research/open_llama_3b_v2)           |   55.6   |   88.0   |   77.9   |    63.1    |   63.3   |    40.1    |      71.4      |   28.1   |    69.2    |   67.4   |   16.0   |   26.8   |\n| [Sheared-LLaMA-2.7B](https://huggingface.co/princeton-nlp/Sheared-LLaMA-2.7B)         |   56.4   |   87.5   |   76.9   |    65.0    |   63.3   |    41.6    |      71.0      |   28.3   |    73.6    |   68.3   |   17.6   | **27.3** |\n| **LLaMA-MoE-3.0B**                                                                    |   55.5   |   84.2   |   77.5   |    63.6    |   60.2   |    40.9    |      70.8      | **30.6** |    71.9    |   66.6   |   17.0   |   26.8   |\n| **LLaMA-MoE-3.5B (4/16)**                                                             | **57.7** |   87.6   | **77.9** |    65.5    | **65.6** |  **44.2**  |    **73.3**    |   29.7   |  **75.0**  | **69.5** | **20.3** |   26.8   |\n| **LLaMA-MoE-3.5B (2/8)**                                                              |   57.6   | **88.4** |   77.6   |  **66.7**  |   65.3   |    43.1    |    **73.3**    |   29.6   |    73.9    |   69.4   |   19.8   |   27.0   |\n\n- SFT models\n\n| Model                                  | MMLU  | ARC-c | HellaSeag | TruthfulQA | MT-Bench |\n| :------------------------------------- | :---: | :---: | :-------: | :--------: | :------: |\n| Sheared LLaMA-2.7B ShareGPT            | 28.41 | 41.04 |   71.21   |   47.65    |   3.79   |\n| Sheared LLaMA-2.7B Deita6K (Our Impl.) | 25.24 | 43.69 |   71.70   |   49.00    |   4.06   |\n| LLaMA-MoE-v1-3.0B (2/16)               | 23.61 | 43.43 |   72.28   |   44.24    |   4.15   |\n| LLaMA-MoE-v1-3.5B (4/16)               | 26.49 | 48.29 |   75.10   |   45.91    |   4.60   |\n| LLaMA-MoE-v1-3.5B (2/8)                | 25.53 | 45.99 |   74.95   |   44.39    |   4.72   |\n\n\u003ch2 id=\"expert-construction\"\u003e🚧 Expert Construction\u003c/h2\u003e\n\n- Neuron-Independent\n  - Independent\u003csub\u003eRandom\u003c/sub\u003e: `bash ./scripts/expert_construction/split/run_split_random.sh`\n  - Independent\u003csub\u003eClustering\u003c/sub\u003e: `bash ./scripts/expert_construction/split/run_split_clustering.sh`\n- Neuron-Sharing\n  - Sharing\u003csub\u003eInner\u003c/sub\u003e: `bash ./scripts/expert_construction/split/run_split_gradient.sh`\n  - Sharing\u003csub\u003eInter\u003c/sub\u003e: `bash ./scripts/expert_construction/split/run_split_gradient_residual.sh`\n\nFor more information, please refer to [Expert Construction docs](docs/expert_construction/README.md).\n\n\u003ch2 id=\"continual-pretraining\"\u003e🚅 Continual Pre-training\u003c/h2\u003e\n\n\n### Tokenization\n\nDownload [SlimPajama](https://www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama) into `/path_to_data` and put data from different domains into separate folders:\n  - `/path_to_data/en_arxiv`\n  - `/path_to_data/en_book`\n  - `/path_to_data/en_c4`\n  - `/path_to_data/en_cc`\n  - `/path_to_data/en_stack`\n  - `/path_to_data/en_wikipedia`\n  - `/path_to_data/github`\n\nEach file should be end with `*.jsonl` and each line looks like:\n```\n{\"id\": \"id-info\", \"content\": \"raw text to be tokenized\"}\n```\n\nRun the following command to tokenize the data in each folder:\n\n```bash\npython -m smoe.utils.tokenize \\\n  -f jsonl \\\n  -t /path_to_tokenizer \\\n  -i /path_to_data/en_arxiv \\\n  -o /path_to_data_tokenized/en_arxiv\n```\n\n### Continual Pre-training (CPT)\n\n- **NOTICE:** Please create `logs/` folder manually: `mkdir -p logs`\n- To run the continual pre-training, please check the [CPT docs](docs/continual_pretraining/README.md).\n\n\u003ch2 id=\"evaluation\"\u003e💎 Evaluation\u003c/h2\u003e\n\n- For evalution on Natural Questions (NQ), please refer to [opencompass](https://github.com/Spico197/opencompass/tree/main).\n- For other tasks, please refer to [lm-eval-harness](https://github.com/spico197/smoe-eval).\n\n\u003ch2 id=\"sft\"\u003e💬 Supervised Fine-Tuning (SFT)\u003c/h2\u003e\n\nWe provide simple examples of SFT to build chatbots.\nPlease refer to [SFT docs](/mnt/petrelfs/zhutong/smoe/docs/supervised_fine_tuning/SFT.md) and `/mnt/petrelfs/zhutong/smoe/scripts/sft` for more details.\n\n\u003ch2 id=\"citation\"\u003e📑 Citation\u003c/h2\u003e\n\n```bibtex\n@article{llama-moe,\n  title={LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training},\n  author={Tong Zhu and Xiaoye Qu and Daize Dong and Jiacheng Ruan and Jingqi Tong and Conghui He and Yu Cheng},\n  journal={arXiv preprint arXiv:2406.16554},\n  year={2024},\n  url={https://arxiv.org/abs/2406.16554},\n}\n```\n\n\u003chr\u003e\n\u003cp align=\"center\"\u003eLLaMA-MoE Team w/ ❤️\u003c/p\u003e\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpjlab-sys4nlp%2Fllama-moe","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fpjlab-sys4nlp%2Fllama-moe","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpjlab-sys4nlp%2Fllama-moe/lists"}