Train a language model on a laptop CPU in ten minutes to generate short stories
You can train a small language model in under ten minutes on a laptop without a GPU to generate short stories. The project implements everything in PyTorch, from the basic principle of converting text to numbers and predicting the next word to the structures and reinforcement learning techniques used in modern open-source models. By focusing on core algorithms rather than complex libraries, it provides a clear understanding of how the model internals work.
A training pipeline implemented without high-level libraries
This repository writes all algorithms in pure PyTorch without using high-level libraries like transformers or trl. The code for each step, from data preprocessing to model structure, optimization, and inference, is separated, making it easy to modify or extend specific parts. It supports both the original 2017 Transformer structure and modern model structures from 2023 onwards (such as RoPE, RMSNorm, and SwiGLU), allowing you to directly compare the differences between training models and actual deployment models. The training process follows a flow that converts raw text into tokens, trains a base model, and then proceeds through post-training stages like instruction tuning (SFT), reward models, PPO, DPO, and GRPO to complete the final assistant model.
Training strategies for environments from laptops to data centers
In environments without a GPU, a 'tiny' preset allows you to train an ultra-small model with 369K parameters in about five minutes using only the CPU. This is achieved by directly training a BPE tokenizer and using a small vocabulary to reduce the size of the embedding table. The trainable size of larger models varies depending on GPU memory capacity, and the README includes a table summarizing the memory capacity and maximum trainable model size (e.g., 1B~8B) for various GPUs from A100 to RTX 4060. If memory is insufficient, you can significantly reduce memory usage by enabling options like --amp or --grad-checkpointing. Additionally, running scripts/benchmark.py lets you measure the training speed of each preset on your current hardware, allowing you to check performance before starting a long training run.
Data preparation and loss function design for each training stage
The data the model learns from is divided into four flows depending on the purpose. Pre-training text is stored in HDF5 files, and instruction tuning data is masked so that only the assistant's response part is learned. Preference data is used for reward model and DPO training, while reinforcement learning prompts are stored in a format that includes the correct answer. In particular, in dialogue formats, markers distinguishing the roles of the user and the assistant are explicitly stated in the text, and the loss function is designed to apply only to the parts spoken by the assistant. For reinforcement learning in math problem solving, the model is guided to place the reasoning process inside think tags and the final answer inside answer tags, applying a reward function that evaluates both the accuracy of the answer and the appropriateness of the format.
Prerequisites and difficulty level for training
To use this repository, you need basic knowledge of object-oriented programming, fundamental neural network concepts, and how to use PyTorch. The README suggests which parts to read depending on the target audience, such as students, developers, and researchers, and all code blocks include expected outputs so you can verify the learning effect immediately. While it covers content at a level applicable to actual research or development, going beyond a simple tutorial by including modern model structures and post-training techniques, all code is implemented directly, which can require significant time and effort if debugging or modifications are needed. The trainable size for specific GPU models is for reference only, and in actual environments, you should appropriately adjust memory management options.