FareedKhan-dev/train-llm-from-scratch

A straightforward method for training your LLM, from downloading data to generating text.

What it solves

This project provides a complete, end-to-end implementation of Large Language Models (LLMs) built from scratch using PyTorch. It removes the dependency on high-level libraries like transformers, trl, or peft, allowing users to understand and implement every stage of the LLM lifecycle—from raw text processing to an aligned, reasoning-capable model—using plain PyTorch code.

How it works

The project follows a linear pipeline that transforms raw data into a functional assistant:

  1. Data Preparation: Raw text is tokenized (using tiktoken or a custom BPE tokenizer) and stored in HDF5 files. It handles pretraining text, instruction data for SFT, preference pairs for reward models, and RL prompts.
  2. Model Architecture: It implements both a classic Transformer (based on "Attention is All You Need") and a modern decoder-only architecture featuring RoPE, RMSNorm, SwiGLU, Grouped-Query Attention (GQA), Mixture of Experts (MoE), and a KV cache.
  3. Pretraining: The model is trained on a next-token prediction loss to create a base model.
  4. Post-Training: The base model is refined through a sequence of Supervised Fine-Tuning (SFT), Reward Modeling, and Reinforcement Learning techniques including PPO, DPO, ORPO, KTO, and GRPO.
  5. Inference: Implements sampling techniques (top-k, top-p, min-p), speculative decoding, and int8 weights for efficient generation.

Who it’s for

  • Students: Those wanting to learn LLM internals through plain-language explanations and step-by-step code.
  • Developers: Those looking for runnable scripts to train models ranging from 369K to billions of parameters on hardware from laptop CPUs to A100 GPUs.
  • Researchers: Those interested in hand-written implementations of post-training algorithms like GRPO and PPO on a single Transformer architecture.

Highlights

  • Zero-Dependency Core: Every algorithm is handwritten in PyTorch without relying on transformers or peft.
  • Hardware Flexibility: Includes a "laptop track" for training tiny models on CPUs in minutes.
  • Modern Features: Implements advanced architecture components like RoPE, MoE, and the Muon optimizer.
  • Comprehensive RL Suite: Includes a wide array of alignment algorithms (DPO, PPO, GRPO, etc.) from scratch.
  • Built-in Tooling: Provides a Streamlit control panel, a documentation site, and a benchmarking script.

Related

  • Project
    lucidrains/PaLM-rlhf-pytorch

    A PyTorch implementation of RLHF on the PaLM architecture, providing tools to train a base model, a reward model, and align them using a reinforcement learning trainer.

  • Project
    mlwithme/BertWithPretrained

    A PyTorch implementation of BERT that allows users to build the model from scratch and apply it to various NLP tasks like classification, QA, and NER.

  • Project
    datawhalechina/llms-from-scratch-cn

    A Chinese‑language, hands‑on tutorial that walks users from basic PyTorch concepts to a full GPT‑style implementation, then extends the lessons to modern architectures (GLM‑4, Llama 3, RWKV, MiniCPM). Includes notebooks, training scripts, and exercises, all under Apache 2.0 / CC‑BY‑NC‑SA licensing.

  • Dispatch
    Illustrating Reinforcement Learning from Human Feedback (RLHF)

    Hugging Face explains Reinforcement Learning from Human Feedback (RLHF), a three-step process used to align large language models with complex human values by optimizing them using a reward model based on human preferences.

  • Project
    Lightning-AI/litgpt

    A high-performance framework for pretraining, finetuning, and deploying 20+ LLMs with from-scratch implementations and no abstraction layers.