FareedKhan-dev/train-llm-from-scratch

Train LLM From Scratch I am Looking for a PhD position in AI . GitHub I implemented a transformer model from scratch using PyTorch, based on the paper Attention is All You Need . You can use my scripts to train your own billion or million parameter LLM using a single GPU. This started as a pretraining tutorial. It now goes all the way from raw text to an aligned, reasoning style model, with every algorithm hand written in plain PyTorch (no trl , no peft , no transformers ). The whole journey is one idea repeated: turn text into numbers, predict the next token, then keep changing the data and the loss until the model does what we want. Here is the path we will walk, end to end: raw text -> tokens -> a Transformer -> next-token loss -> a base model base model -> SFT -> Reward Model -> {PPO, DPO} -> GRPO -> evaluation and chat Below is the output of a trained 13 million parameter LLM, just so you can see where the small end of this starts: In ***1978, The park was returned to the factory-plate that the public share to the lower of the electronic fence that follow from the Station's cities. The Canal of ancient Western nations were confined to the city spot. The villages were directly linked to cities in China that revolt that the US budget and in Odambinais is uncertain and fortune established in rural areas. Table of Contents Who this is for Prerequisites and Training Time Setup Code Structure Step 1: Preparing the Data Step 2: The Model, Built From Small Pieces Multi Layer Perceptron (MLP) Single Head Attention Multi Head Attention The Transformer Block The Full Transformer Step 3: Pretraining the Base Model Step 4: Generating Text Step 5: Post-Training, Turning a Base Model Into an Assistant SFT (Supervised Fine-Tuning) The Reward Model DPO, ORPO and KTO PPO GRPO / RLVR Step 6: Evaluation Step 7: Talking to the Model The Streamlit Control Panel The Documentation Site Run the Whole Thing What's Next Who this is for I tried to write this so one page works for very different readers: If you are a student , read top to bottom. Every block of code comes after a plain explanation of what it does and why, and most blocks are followed by the output you should expect. If you are a developer , the commands and file paths are all here. You can copy, run, and read the referenced source files directly. If you are a researcher , the post-training half is the interesting part: SFT, a Bradley-Terry reward model, PPO with GAE, DPO/ORPO/KTO, and GRPO, all from scratch on the same small Transformer, trained on real public datasets. Every diagram in this README is colored the same way, so the colors mean something: green is raw data teal is stored, tokenized data on disk blue is a plain processing step yellow is the model or a training step orange is the reinforcement learning and reward parts red is a loss grey is a saved checkpoint purple is the final output or evaluation Prerequisites and Training Time You need a basic understanding of object oriented programming, neural networks, and PyTorch. Below are some resources to help you get started: Topic Video Link OOP OOP Video Neural Network Neural Network Video Pytorch Pytorch Video You will need a GPU to train. A free Colab or Kaggle T4 is enough for the 13 million parameter model, but it will not fit a billion parameter model. Here is a rough guide: GPU Name Memory 2B LLM Training 13M LLM Training Max Practical LLM Size (Training) NVIDIA A100 40 GB ✔ ✔ ~6B to 8B NVIDIA V100 16 GB ✘ ✔ ~2B NVIDIA RTX 4090 24 GB ✔ ✔ ~4B NVIDIA RTX 5090 32 GB ✔ ✔ 13M verified, larger configs TBD NVIDIA RTX 3090 24 GB ✔ ✔ ~3.5B to 4B NVIDIA RTX 4080 16 GB ✘ ✔ ~2B NVIDIA RTX 4060 8 GB ✘ ✔ ~1B Tesla T4 16 GB ✘ ✔ ~1.5B to 2B If a large config runs out of memory, the pretraining script has opt-in flags ( --amp , --grad-checkpointing , --grad-accum ) that bring the memory down a lot. More on those later. Setup Clone the repository and install it in editable mode. The editable install puts config , src , data_loader , and ui on your import path, so you do not need to set PYTHONPATH by hand anymore: git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git cd train-llm-from-scratch pip install -e . There are optional extras for the parts you want: pip install -e " .[train] " # datasets + wandb, for downloading data and logging pip install -e " .[ui] " # streamlit + pandas + altair, for the control panel pip install -e " .[docs] " # mkdocs, for the documentation site pip install -e " .[all] " # everything There are two config systems, and it helps to know which is which from the start: config/config.py is the original, simple config for the legacy pretraining script scripts/train_transformer.py . It is plain Python constants. config/post_training_config.py plus the JSON files in configs/ drive everything else (pretraining the bigger base, SFT, reward, DPO, PPO, GRPO). You edit a small JSON file per stage, and any field can also be overridden on the command line, for example --lr 2e-5 --batch_size 16 . For fast checks there is a tiny configs/smoke/ variant of every stage that shrinks the model so a full run finishes in seconds on a CPU or a single GPU. Code Structure train-llm-from-scratch/ ├── src/ │ ├── models/ # the Transformer, built from small pieces │ │ ├── mlp.py # the feed-forward block │ │ ├── attention.py # single head and multi head attention │ │ ├── transformer_block.py # one block: attention + MLP + residuals │ │ └── transformer.py # the full model: embeddings + blocks + lm_head │ └── post_training/ # SFT, reward model, PPO, DPO, GRPO, eval, inference ├── config/ │ ├── config.py # legacy pretraining config (plain constants) │ ├── post_training_config.py # dataclasses for every post-training stage │ └── loader.py # merges defaults < base.json < stage.json < CLI ├── configs/ # editable JSON, one file per stage (+ smoke/) ├── data_loader/ # batch iterators for each kind of data ├── scripts/ # every runnable step lives here ├── ui/ # the Streamlit control panel ├── docs/ # the MkDocs site (theory + diagrams) ├── images/ # the diagrams in this README (+ the generator) └── pyproject.toml # pip install -e . Step 1: Preparing the Data A model only ever sees integers. So the first job is always the same: take text, turn it into token ids, and store those ids on disk in a format that is fast to read during training. We do this four times, once for each kind of training we will do later. The four streams are: Pretraining text from The Pile , stored as a flat array of token ids in an HDF5 file. Instruction data (Alpaca, Dolly, GSM8K) for SFT, packed into fixed length rows with a mask that says which tokens are the assistant's answer. Preference pairs (Anthropic HH-RLHF, UltraFeedback) for the reward model and DPO, stored as {prompt, chosen, rejected} . RL prompts (GSM8K and a small arithmetic warm-up) for PPO and GRPO, stored as {prompt, gold} . Tokenization We use the r50k_base tokenizer from OpenAI's tiktoken , the same one GPT-3 used. Text becomes a list of integers, and we append a special <|endoftext|> token (id 50256) at the end of every document so the model learns where one piece of text stops and the next begins. For the legacy path, download a slice of The Pile and tokenize it into HDF5: python scripts/data_download.py # downloads the validation file + 1 training shard python scripts/data_preprocess.py # tokenizes to data/train/pile_train.h5 and data/val/pile_dev.h5 The newer, faster path streams and batch-encodes the same data straight into a flat token array: python scripts/prepare_pretrain_data.py --split val --out data/pile_dev.h5 python scripts/prepare_pretrain_data.py --split train --num_shards 1 --out data/pile_train.h5 Once tokenized, the data is just a long line of integers. Here is a real peek at the validation file I prepared for this README (8.