JuncaiL/llama-265m

LLaMA-265M

💻 Code

👋 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+\mathrm{.}\mathrm{.}\mathrm{.}=N1p\; +\; 2p^2\; +\; 3p^3\; +\; ...\; =\; N$$

We can get the hit rate as follow.

$$p=1+\frac{1-\sqrt{1+4N}}{2N}p\; =\; 1\; +\; \backslash frac\{1-\backslash 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

1. My implementation of MoE structure is based on the repo https://huggingface.co/llama-moe/LLaMA-MoE-v1-3_5B-2_8
2. 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）