Training a Small Model from Scratch—1: Pretraining
Training a small model from scratch is an important part of my study of large language models (LLMs). By reproducing a small-scale model with a minimal implementation, I can gain an in-depth understanding of the basic training process while systematically learning its core principles and actual operating mechanisms.
Read in Chinese→Training a Small Model from Scratch: 1. Pretraining
Introduction
Training a small model from scratch is an important part of my personal study of large language models (LLMs).
By reproducing a small-scale model with a minimal implementation, I can not only gain an in-depth understanding of the basic model training process, but also systematically learn its core principles and actual operating mechanisms. This hands-on approach allows me to experience each stage of model training firsthand while mastering the relevant technical details and implementation methods.
In my current work, there are indeed some tasks that require small models. I also want to understand their operating logic during the learning process and strive to implement minimal models for certain subtasks while consuming as few resources as possible.
ALL in AI
The code has been uploaded: https://github.com/godzeo/miniGPT
The first part of this series is pretraining.
Pretraining
As I understand it, the first step in training a model is learning knowledge.
The first thing an LLM needs to learn is not how to communicate directly with people. Just as students must first acquire foundational knowledge, AI needs to “read” large amounts of material in order to learn. It reads Wikipedia, news articles, and books, learning knowledge and patterns from them. This learning process is autonomous and does not require a teacher—a human—to provide guidance alongside it.
The AI’s main task is simple: learning to “continue the text.” For example, when it sees “The Heavenly King subdues the earthly tiger,” it can automatically follow it with “The pagoda suppresses the river demon.” This is similar to playing a word-chain game: by reading vast amounts of text, the AI learns how to continue such sequences.
Small Models
For small language models, properly optimizing the tokenizer is important, but it is not simply a matter of “compression”:
- Balancing parameter count and vocabulary size: Small models have a limited number of parameters. An excessively large vocabulary will consume too many parameters in the token embedding layer, and the tokenizer affects the entire training process as well as the upper limit of the model’s capabilities.
The tokenizer requires further in-depth research. I plan to use the same dataset and model architecture while changing only the tokenizer configuration parameters to compare the results. For now, this part is mainly based on the experience of others; I will need to study it more systematically later.
I initially trained with an extremely small vocabulary of 6,400 tokens, but found that the results were mediocre.
I rented two RTX 4090 GPUs for training.```
torchrun –nproc_per_node=2 pretrain.py
–output_dir ../outputs/pretrain_full_test
–data_file /root/autodl-fs/data/pretrain_hq.jsonl
–tokenizer_path ../tokenizer
–batch_size 64
–num_epochs 1
–fp16
So, after adjusting the budget and switching to mistralai_tokenizer:
Pretraining on four RTX 4090 GPUs took less than two hours.
# Main Code
[https://github.com/godzeo/miniGPT/blob/main/scripts/pretrain.py](https://github.com/godzeo/miniGPT/blob/main/scripts/pretrain.py)
1. **Main Features**
- Load a pretrained tokenizer
- Create or load a model
- Prepare the pretraining dataset
- Configure training parameters
- Create a Trainer and begin training
- Support resuming from checkpoints and exception handling
### Explanation of Important Training Parameters
1. **Basic Training Parameters**
- `batch_size`: The training batch size per device. This is like the amount of ingredients a chef puts into the wok each time. The larger the number, the more data is processed at once. It mainly depends on available VRAM—exceeding it will immediately cause the process to crash.
- `gradient_accumulation_steps`: The number of gradient accumulation steps. This is like weighing ingredients in several batches and then adjusting the recipe all at once. When VRAM is insufficient, this method can simulate the effect of cooking with a larger wok.
- `num_epochs`: The number of training epochs, similar to how many times a student repeatedly works through practice exam questions. Too many epochs can lead to rote memorization—overfitting—while too few may prevent the model from learning.
- `learning_rate`: The learning rate—the parameter that controls the model’s step size while learning.
- A high learning rate can easily cause the model to miss details.
- A low learning rate produces more thorough learning but takes more time.
- `weight_decay`: Weight decay, which uses regularization to prevent the model from overfitting.
- `warmup_ratio`: The warmup ratio, which allows the learning rate to increase gradually from zero to the configured value.
The above are training parameters for this scenario. At present, I only understand what they mean, but there are no definitive values for how they should be configured to achieve the desired results. In most cases, the appropriate settings can only be determined through experimentation. For now, I use some default values and directly adopt parameter settings that others have found to work well in most cases.
- `data_limit`: Used to limit the amount of training data. When testing, a smaller dataset is used to verify that the entire process runs successfully.
### Training Considerations
1. **Resuming Training from a Checkpoint**
- Supports resuming training from a checkpoint
- Automatically checks checkpoint integrity
- Automatically saves the current state if training is interrupted
Training results:
Let’s load the model and examine its performance. At this point, the model can only play a basic Chinese idiom-chain game.
Results after one epoch:
{'train_runtime': 6075.9813, 'train_samples_per_second': 232.572, 'train_steps_per_second': 0.227, 'train_loss': 3.2094810264697973, 'epoch': 1.0}
