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 a Large Language Model (LLM) built from scratch using PyTorch. It removes the dependency on high-level libraries like transformers, trl, or peft, allowing users to understand and implement the entire pipeline from raw text processing to an aligned, reasoning-style model.

How it works

The project follows a sequential pipeline to transform raw text into a functional assistant:

  1. Data Preparation: Raw text is tokenized using OpenAI's tiktoken and stored in HDF5 files. This includes pretraining data, instruction data for SFT, preference pairs for reward modeling, and RL prompts.
  2. Model Architecture: A Transformer model is built from basic PyTorch modules, including a Multi-Layer Perceptron (MLP), Single Head Attention with causal masking, Multi-Head Attention, and Transformer Blocks with pre-norm residual connections.
  3. Pretraining: The base model is trained on a next-token prediction loss using a dataset like The Pile.
  4. Post-Training: The base model is converted into an assistant through several stages:
    • SFT (Supervised Fine-Tuning): Training on instruction data using a loss mask to focus only on assistant responses.
    • Reward Modeling: Training a Bradley-Terry reward model on preference pairs.
    • Alignment: Applying techniques such as PPO (Proximal Policy Optimization), DPO (Direct Preference Optimization), ORPO, KTO, and GRPO (Group Relative Policy Optimization) to align the model with human preferences and reasoning capabilities.

Who it’s for

  • Students: Those wanting a step-by-step guide with plain explanations and expected outputs for every code block.
  • Developers: Those looking for runnable scripts and clear file paths to implement their own LLM.
  • Researchers: Those interested in the from-scratch implementation of post-training algorithms like PPO, DPO, and GRPO on a small Transformer.

Highlights

  • Zero-Dependency Core: Implemented entirely in plain PyTorch without using transformers or peft libraries.
  • Full Pipeline: Covers the entire journey from raw text $\rightarrow$ tokens $\rightarrow$ base model $\rightarrow$ SFT $\rightarrow$ Reward Model $\rightarrow$ RLHF/Alignment $\rightarrow$ Evaluation.
  • Scalable: Supports training models ranging from 13 million to several billion parameters depending on the available GPU memory.
  • Integrated Tooling: Includes a Streamlit control panel for management and a dedicated documentation site for theory and diagrams.

Related

  • Project
    lucidrains/PaLM-rlhf-pytorch

    A PyTorch implementation of RLHF on top of the PaLM architecture, enabling the alignment of large language models with human preferences.

  • Project
    jingyaogong/minimind

    A lightweight framework for training ultra-small LLMs from scratch using native PyTorch, designed as a practical tutorial for understanding the full LLM lifecycle.

  • 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/pytorch-lightning

    A deep learning framework that abstracts away the engineering boilerplate of PyTorch, enabling researchers to pretrain and finetune AI models at scale across various hardware.

  • Project
    HorizonRobotics/alf

    ALF (Agent Learning Framework) is a PyTorch‑based reinforcement‑learning library offering a large collection of algorithms (A2C, PPO, SAC, MuZero, etc.), modular components, Python‑style configuration, and support for single‑ and multi‑GPU/cluster training. It includes many ready‑to‑run examples (CartPole, Atari, MuJoCo, MetaDrive) and provides Docker, Nix, and uv‑based installation options.