Aurora: Activating chinese chat capability for Mistral-8x7B sparse Mixture-of-Experts through Instruction-Tuning
Rongsheng Wang, Haoming Chen, Ruizhe Zhou, Yaofei Duan, Kunyan Cai, Han Ma, Jiaxi Cui, Jian Li, Patrick Cheong-Iao Pang, Yapeng Wang, Tao Tan☨
☨Corresponding author
Please follow our Github: https://github.com/WangRongsheng/Aurora
Overview
Existing research has demonstrated that refining large language models (LLMs) through the utilization of machine-generated instruction-following data empowers these models to exhibit impressive zero-shot capabilities for novel tasks, without requiring human-authored instructions. In this paper, we systematically investigate, preprocess, and integrate three Chinese instruction-following datasets with the aim of enhancing the Chinese conversational capabilities of Mixtral-8x7B sparse Mixture-of-Experts model. Through instruction fine-tuning on this carefully processed dataset, we successfully construct the Mixtral-8x7B sparse Mixture-of-Experts model named "Aurora." To assess the performance of Aurora, we utilize three widely recognized benchmark tests: C-Eval, MMLU, and CMMLU. Empirical studies validate the effectiveness of instruction fine-tuning applied to Mixtral-8x7B sparse Mixture-of-Experts model. This work is pioneering in the execution of instruction fine-tuning on a sparse expert-mixed model, marking a significant breakthrough in enhancing the capabilities of this model architecture.
Evaluation
It is known that LLM evaluation remains a significant challenge. We use three public benchmarks in our study.
Next are some references we gave you about GPU memory usage during the training and inference stage. Please note that we did all inference and training on a single GPU.
|Stage|GPU Memory Usage| |:-|:-| |Training|~43 GiB| |Inference|~25 GiB|
Quick-Use
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
from threading import Thread
from peft import PeftModel
import time
model_name_or_path = "mistralai/Mixtral-8x7B-Instruct-v0.1" # download weights from https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1
lora_weights = "wangrongsheng/Aurora-Mixtral-8x7B" # download weights from https://modelscope.cn/models/wangrongsheng/Aurora-Mixtral-8x7B
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
model0 = AutoModelForCausalLM.from_pretrained(model_name_or_path, load_in_4bit=True, device_map="auto", torch_dtype=torch.bfloat16)
model = PeftModel.from_pretrained(
model0,
lora_weights,
)
class StopOnTokens(StoppingCriteria):
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
stop_ids = [0,]
for stop_id in stop_ids:
if input_ids[0][-1] == stop_id:
return True
return False
def convert_history_to_text(history):
text = ""
if len(history) > 1:
text = "<s> " + "".join(
[
"".join(
[
f"[INST]{item[0]}[/INST] {item[1]} ",
]
)
for item in history[:-1]
]
) + "</s> "
text += "".join(
[
"".join(
[
f"[INST]{history[-1][0]}[/INST]",
]
)
]
)
return text
def predict(message, history):
history_transformer_format = history + [[message, ""]]
stop = StopOnTokens()
messages = convert_history_to_text(history_transformer_format)
model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
generate_kwargs = dict(
model_inputs,
streamer=streamer,
max_new_tokens=4096,
do_sample=True,
top_p=0.95,
top_k=1000,
temperature=1.0,
num_beams=1,
pad_token_id=tokenizer.eos_token_id,
stopping_criteria=StoppingCriteriaList([stop])
)
t = Thread(target=model.generate, kwargs=generate_kwargs)
t.start()
partial_message = ""
t1 = time.time()
count = 0
for new_token in streamer:
if new_token != '<':
partial_message += new_token
count += 1
yield partial_message
t2 = time.time()
speed = count/(t2-t1)
print("inference speed: %f tok/s" % speed)
gr.ChatInterface(predict,chatbot=gr.Chatbot(height=600,),title="MoE").queue().launch()
Easy-to-Use
1. Clone and Set up
https://github.com/WangRongsheng/Aurora.git
cd Aurora
pip install -r requirements.txt
2. Download Model
Base Model: |Model|Download| |:-|:-| |Mixtral-8x7B-Instruct-v0.1|[HuggingFace] [HuggingFace-mirror] [ModelScope]|
LoRA Model: |Model|Download| |:-|:-| |Aurora|[HuggingFace]|
The huge model parameters are not convenient for you to manage your task, so we provide LoRA weights, which will be merged with the base model before inference. You don't have to worry about it.
3. Inference
Web:
CUDA_VISIBLE_DEVICES=0 python src/web_demo.py \
--model_name_or_path ./Mixtral-8x7B-Instruct-v0.1 \
--checkpoint_dir Aurora \
--finetuning_type lora \
--quantization_bit 4 \
--template mistral
Then you can visit: http://127.0.0.1:7860/
CLI:
CUDA_VISIBLE_DEVICES=0 python src/cli_demo.py \
--model_name_or_path ./Mixtral-8x7B-Instruct-v0.1 \
--checkpoint_dir Aurora \
--finetuning_type lora \
--quantization_bit 4 \
--template mistral
API:
CUDA_VISIBLE_DEVICES=0 python src/api_demo.py \
--model_name_or_path ./Mixtral-8x7B-Instruct-v0.1 \
--checkpoint_dir Aurora \
--finetuning_type lora \
--quantization_bit 4 \
--template mistral
If you need to load weights for specific checkpoints, you can set them up like this: --checkpoint_dir Aurora/checkpoint-5000
.
Train
If you have a single GPU and its GPU memory size is larger than 48GB, you can train your own models.
Train your MoE model
CUDA_VISIBLE_DEVICES=5 python src/train_bash.py \
--stage sft \
--model_name_or_path ./Mixtral-8x7B-Instruct-v0.1 \
--do_train \
--dataset alpaca_zh,alpaca_gpt4_zh,sharegpt \
--finetuning_type lora \
--quantization_bit 4 \
--overwrite_cache \
--output_dir output/ \
--per_device_train_batch_size 2 \
--gradient_accumulation_steps 4 \
--lr_scheduler_type cosine \
--logging_steps 100 \
--save_steps 1000 \
--learning_rate 5e-5 \
--num_train_epochs 3.0 \
--plot_loss \
--fp16 \
--template mistral \
--lora_target q_proj,v_proj
--quantization_bit 4
means you will use QLoRA
, If you have a larger GPU memory size you can remove it and use LoRA
.
Evaluation your MoE model
CUDA_VISIBLE_DEVICES=0 python src/evaluate.py \
--model_name_or_path ./Mixtral-8x7B-Instruct-v0.1 \
--checkpoint_dir Aurora/checkpoint-5000 \
--finetuning_type lora \
--quantization_bit 4 \
--template mistral \
--task cmmlu \ # cmmlu, mmlu, ceval
--split test \
--lang en \ # zh, en
--n_shot 5 \
--batch_size 8
Acknowledgments
This work is mainly done by the Faculty of Applied Sciences of the Macao Polytechnic University. The computational resources used in this work were obtained from AWS servers. The fine-tuning framework we used is LLaMA-Factory, which brings a lot of convenience to our work. We also thank the public datasets from the open source community, such as shareAI, stanford_alpaca and GPT-4-LLM. Most importantly we are very grateful to Mistral AI, who are leading a new technology boom that will dramatically change the future of technology development.
Citation
If you find our work helpful, feel free to give us a cite.
@misc{wang2023auroraactivating,
title={Aurora:Activating Chinese chat capability for Mistral-8x7B sparse Mixture-of-Experts through Instruction-Tuning},
author={Rongsheng Wang and Haoming Chen and Ruizhe Zhou and Yaofei Duan and Kunyan Cai and Han Ma and Jiaxi Cui and Jian Li and Patrick Cheong-Iao Pang and Yapeng Wang and Tao Tan},
year={2023},
eprint={2312.14557},
archivePrefix={arXiv},
primaryClass={cs.CL}
}
License
Please follow the Apache 2.0 License.
评论