Minillm 0.2B WithWudao

🚀 MiniLLM Project

This project aims to build a small-parameter LLM, going through the four stages of pretraining -> instruction fine-tuning -> reward model -> reinforcement learning to complete a chat model capable of simple chat tasks at a controllable cost.

🚀 Quick Start

Environment Installation

pip install bert4torch==0.4.9.post2  # If not found, specify -i https://pypi.org/simple

Script Explanation

# To prevent the terminal from closing, you can use nohup, tmux, or screen to start the script.
# eg. nohup torchrun --standalone --nproc_per_node=4 pretrain.py --name baby > nohup.log&

# Pretraining
cd pretrain
torchrun --standalone --nproc_per_node=4 pretrain.py  # Some reports indicate that DDP training may crash halfway. You need to set `export NCCL_IB_DISABLE=1`.

# Pretraining Inference (Command-line Chat)
cd pretrain
python infer.py  # python infer_transformers.py

# Instruction Fine-tuning Training
cd sft
python sft.py

# Instruction Fine-tuning Inference (Command-line Chat)
cd sft
python infer.py  # python infer_transformers.py

# Convert the checkpoint to a format that can be run by transformers
cd docs
python convert.py

✨ Features

• Efficient Training Framework: Uses the bert4torch training framework, with concise and efficient code.
• Inference Compatibility: The trained checkpoints can be directly used for inference with the transformers package.
• Memory Optimization: Optimizes memory usage during training.
• Reproducibility: Provides complete training logs for reproduction and comparison.

Note: The model trained in this experiment currently only has simple chat capabilities (limited by the size of the corpus, model scale, and the size and quality of the SFT corpus) and is not capable of answering complex questions.

📦 Installation

The installation steps are included in the "Quick Start" section. Please follow the instructions there to install the necessary environment and run the scripts.

💻 Usage Examples

Basic Usage - Pretraining Model Inference

# Add the following two lines depending on the network situation
import os
os.environ['HF_ENDPOINT'] = "https://hf-mirror.com"

from transformers import AutoTokenizer, LlamaForCausalLM
import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model_name = 'Tongjilibo/MiniLLM-L12_H1024_A8-WithWudao'  # 'Tongjilibo/MiniLLM-L12_H1024_A8-NoWudao'

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = LlamaForCausalLM.from_pretrained(model_name).to(device)

query = '王鹏是一名'
inputs = tokenizer.encode(query, return_tensors='pt', add_special_tokens=False).to(device)
output_ids = model.generate(inputs)
response = tokenizer.decode(output_ids[0].cpu(), skip_special_tokens=True)
print(response)

Basic Usage - Instruction Fine-tuning Model Inference

# Add the following two lines depending on the network situation
import os
os.environ['HF_ENDPOINT'] = "https://hf-mirror.com"

from transformers import AutoTokenizer, LlamaForCausalLM
import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model_name = 'Tongjilibo/MiniLLM-L12_H1024_A8-WithWudao-SFT_Alpaca'

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = LlamaForCausalLM.from_pretrained(model_name).to(device)

query = '你好'
query = f'<human>{query}<robot>'
inputs = tokenizer.encode(query, return_tensors='pt', add_special_tokens=False).to(device)
output_ids = model.generate(inputs)
response = tokenizer.decode(output_ids[0].cpu(), skip_special_tokens=True)[len(query):]
print(response)

📚 Documentation

Update History

• 20240316: Initial submission, including the pretrained models MiniLLM-MiniLLM-L12_H1024_A8-NoWudao and MiniLLM-MiniLLM-L12_H1024_A8-WithWudao, and the SFT model MiniLLM-L12_H1024_A8-WithWudao-SFT_Alpaca.

Pretraining

Pretraining Corpus (from baby-llama2-chinese)

The project has open-sourced the pretraining corpus processed by the ChatGLM2-6B tokenizer, with a total data volume of 63.4 billion Tokens. The link is as follows: Corpus (Extraction code: 6unr).

Pretraining Weights

Pretraining Process

• Training Parameter Configuration and Training Duration

Pretraining Model Invocation

The code example is provided in the "Usage Examples" section.

Pretraining Continuation Effect

• MiniLLM-L12_H1024_A8-NoWudao

User: Xiaoming has excellent academic performance and good health. He is a

Continuation: Xiaoming has excellent academic performance and good health. He is an outstanding student at Dalian University of Technology. Since his freshman year in college, he has been very fond of sports, with top-notch grades, and has participated in many competitions. In the school's sports activities, he ranked among the top and was awarded the title of "Outstanding Athlete" by Dalian University of Technology.
In 2006, in his sophomore year of high school, Xiaoming's classmate Yang Fan came to Dalian University of Technology. He was determined to change his fate. Every day after school, he would run to the playground to practice basketball. By around 3 pm, Xiaoming's basketball skills began to show. His basketball moves were agile and powerful. Moreover, he often used his free time to practice jump shots.

User: Before my bed, the moonlight is so bright. I wonder if it's frost on the ground. I raise my head to look at the bright moon.

Continuation: Before my bed, the moonlight is so bright. I wonder if it's frost on the ground. I raise my head to look at the bright moon. Lowering my head, I miss my hometown...
When I read these poems, I also feel very kind and familiar.

• MiniLLM-L12_H1024_A8-WithWudao

User: Xiaoming has excellent academic performance and good health. He is a

Continuation: Xiaoming has excellent academic performance and good health. He is an outstanding teacher. He is honest and responsible. In class, he can listen attentively, think actively, and actively raise his hand to answer questions in every class. After class, he actively asks to help the teacher, carefully completes his homework, and finishes the homework assigned by the teacher on time, which is well-received by his classmates.

User: Before my bed, the moonlight is so bright. I wonder if it's frost on the ground. I raise my head to look at the bright moon.

Continuation: "Before my bed, the moonlight is so bright. I wonder if it's frost on the ground. I raise my head to look at the bright moon. Lowering my head, I miss my hometown." This poem not only describes the warmth and tranquility of the Mid-Autumn night but also expresses the author's deep affection and longing for his hometown. The phrase "The moon rises above the willow branches" describes what is seen under the moon. "I want to pour my feelings into the lute" refers to the desire to express farewell feelings; "I raise my head to look at the bright moon" describes the scene of the moon hanging high in the sky, shining brightly with stars on the Mid-Autumn night; "Lowering my head, I miss my hometown" expresses the deep affection and longing for the hometown.
In terms of writing techniques, this poem mainly uses symbolism.

Instruction Fine-tuning

Instruction Fine-tuning Corpus (Selected Available Datasets)

Instruction Fine-tuning Weights

Instruction Fine-tuning Training Process

• Training Parameter Configuration and Training Duration

Instruction Fine-tuning Model Invocation

The code example is provided in the "Usage Examples" section.

Instruction Fine-tuning Chat Effect

• MiniLLM-L12_H1024_A8-WithWudao-SFT_Alpaca

User: Hello

Assistant: Hello! Nice to meet you.

User: What can you do?

Assistant: Hello! I'm an AI assistant. Is there anything I can do for you?

User: How to make scrambled eggs with tomatoes?

Assistant: If you want to make scrambled eggs with tomatoes, the following steps may help you:

1. Prepare ingredients: Prepare eggs, salt, coriander, shredded carrots, shredded cucumbers, etc.
2. Wash the eggs: Break the eggs into a mixing bowl to clean their surfaces.
3. Prepare materials: Prepare a pot, pour an appropriate amount of water, and pour the egg liquid into the pot.
4. Boil the eggs: Fry the eggs in a soup pot or pan until their surfaces turn golden and they are cooked.
5. Stir-fry the vegetables: Add an appropriate amount of oil to the pot, stir-fry the shredded carrots and shredded cucumbers (according to personal preference) for a few times until the carrots are cooked.
6. Add the eggs: Take out the eggs from the pot and pour them into the pot.
7. Season: Pour the stir-fried eggs into the pot and stir-fry them with the vegetables for a few times.

Note: If you want to stir-fry the vegetables to have a richer flavor, you can add an appropriate amount of salt and coriander to enhance the flavor.

📄 License

This project is licensed under the Apache-2.0 license.