gitmyhub

MiniLLM

Python ★ 14 updated 25d ago

单卡完整LLM训练流程:手写LLaMA2风格Transformer,预训练→SFT→LoRA→DPO,RTX 4050 6GB即可运行

MiniLLM is an educational project that walks through building a small language model entirely from scratch using Python and PyTorch. The goal is to make every step of modern AI training reproducible and understandable on consumer hardware. The entire pipeline runs on a single graphics card with 6 GB of memory, such as an RTX 4050.

The model itself has about 38 million parameters and uses the same architectural style as LLaMA2, a well-known open AI model. This means it uses grouped-query attention (a way of organizing how the model attends to different parts of text), a specific type of feed-forward network called SwiGLU, rotary position encoding, and layer normalization via RMSNorm. All of these components are written by hand in the codebase rather than imported from existing libraries, so readers can follow exactly what each piece does.

Training proceeds through four stages that mirror what large commercial AI labs do at much larger scale. First, pre-training teaches the model basic language patterns by predicting the next token across a large text dataset for 50,000 steps. Second, supervised fine-tuning (SFT) loads the pre-trained model and trains it on instruction-following examples so it learns to respond to prompts. Third, LoRA fine-tuning applies a parameter-efficient technique that updates only a small fraction of the weights (about 2.28%) rather than the whole model. Fourth, DPO (direct preference optimization) trains the model to prefer human-approved responses over rejected ones, which is a technique used to align AI outputs with human preferences.

The repository also includes a baseline track using QLoRA on Qwen2.5-1.5B, a much larger existing model, for comparison. Five ablation experiments cover choices like learning rate, LoRA rank, and the effect of the DPO alignment step. Eleven annotated code analysis documents walk through each major source file in detail. The model can be exported in HuggingFace format and used via a command-line chat interface.