LLaMA-265M
👋 Very nice to meet you here~
❤️ This repo contains the model LLaMA-265M. This model is trained from scratch with FP32 precision. We firstly train the model through wikipedia dataset with 1 epoch and then through 10% of C4 dataset (10 data shards among 1024 data shards) with 1 epoch. This is NOT fine-tuned by instruction pairs, so it may not be good enough to act like a chatbot. The model size is only 265M, which is very convenient for deployment and research usage.
📢 This series also includes a MoE version, see 🤗this repo.
1. 🚀QuickStart
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
model_dir = "JuncaiL/llama-265m"
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_dir, trust_remote_code=True)
model.eval()
model.to("cuda:0")
input_text = "Beijing is a famous city"
inputs = tokenizer(input_text, return_tensors="pt",return_token_type_ids=False)
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))
# Beijing is a famous city in China.
# The city is famous for its beaches, the most famous of which is the most famous and famous of the Beijing. It is also the home for the famous Beijing Opera
2. 📑Checkpoint Details and Evaluation
Model Parameter
| Model | #Experts | #Activated Experts | #Params | # Activated Params | Flops(T) per sample （se q=2048） | Model Weights |
|---|---|---|---|---|---|---|
| 265M | - | - | 265M | 265M | 0.48 | 🤗 llama-265m |
| 8 $\times$ 265M MoE | 8 | 2 | 970M | 332M | 0.76 | 🤗 llama-8x265m-moe |
| llama-7b | - | - | 7B | 7B | 25.29 |
Model Evaluation
We use the "Average number of tokens verified" $N$ ( see reference link ） as the metric to evaluate these models. This metric demonstrates that giving the same input to the small speculative model and llama-7b, counting from the first predicted tokens, how many successive tokens in the output sentence of the small speculative model are the same as the output sentence of the llama-7b.
- Average number of tokens verified
| Dataset | 8 $\times$ 265M MoE | GPT without MoE |
|---|---|---|
| tatsu-lab/alpaca | 3.2362 | 3.0334 |
| alespalla/chatbot_instruction_prompts | 3.2031 | 3.0823 |
| web_questions | 2.7201 | 2.5541 |
| MohamedRashad/ChatGPT-prompts | 3.0954 | 2.9768 |
Supposed that the small speculative model can have a hit rate $p$ for the next token when giving the same input. Then we have
$$ 1p + 2p^2 + 3p^3 + ... = N $$
We can get the hit rate as follow.
$$ p = 1 + \frac{1-\sqrt{1+4N}}{2N}$$
- Hit Rate
| Dataset | 8 $\times$ 265M MoE | GPT without MoE |
|---|---|---|
| tatsu-lab/alpaca | 0.578 | 0.567 |
| alespalla/chatbot_instruction_prompts | 0.576 | 0.570 |
| web_questions | 0.550 | 0.540 |
| MohamedRashad/ChatGPT-prompts | 0.571 | 0.565 |
Acknowledgment
- My implementation of MoE structure is based on the repo
https://huggingface.co/llama-moe/LLaMA-MoE-v1-3_5B-2_8 - My inspiration for Speculative Inference comes from the paper "SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification" (link) . I am very appreciative of the help and suggestions from the SpecInfer group. ❤️
Citation
@misc{specmoe-2024,
title={SpecMoE: Building A Speculative MoE Model To Accelerate Inference},
author={Juncai Liu},
year={2024},
month={March},
url={https://github.com/JuncaiL/SpecMoE/}
}
Contact
If you have any interest or question about this project, please feel free to contact me.
liujc19@mails.tsinghua.edu.cn （before June 30, 2024） or liujc19@tsinghua.org.cn （After June 30, 2024）