train-llm-from-scratch

## Train LLM From Scratch This repository provides a straightforward, end-to-end method for training your own large language model (LLM), from downloading raw data to generating text. It was created by Fareed Khan and is based on the Transformer architecture from the "Attention is All You Need" paper. The project is designed to be accessible to students, developers, and researchers, with every algorithm implemented from scratch in plain PyTorch (without using libraries like `trl`, `peft`, or `transformers`). ### Core Pipeline The repository guides you through a complete LLM training journey: ``` raw text -> tokens -> a Transformer -> next-token loss -> a base model base model -> SFT -> Reward Model -> {PPO, DPO} -> GRPO -> evaluation and chat ``` ### Key Features - **From-Scratch Implementation**: All models and algorithms are hand-written in PyTorch, providing a deep understanding of the underlying mechanisms. - **Step-by-Step Guide**: The README and code are structured to follow a logical path, with clear explanations and code blocks. - **Multiple Training Stages**: Covers pretraining, supervised fine-tuning (SFT), reward model training, and advanced reinforcement learning methods like DPO, PPO, and GRPO. - **Practical and Flexible**: Includes optional memory-saving features (AMP, gradient checkpointing, gradient accumulation) to train larger models on limited hardware. - **Comprehensive Tooling**: Includes scripts for data preparation, training, evaluation, and a Streamlit control panel for monitoring. ### Code Structure The repository is organized into clear modules: - `src/models/`: Contains the Transformer model built from small, reusable components (MLP, attention, blocks). - `src/post_training/`: Implements SFT, reward models, PPO, DPO, GRPO, evaluation, and inference. - `scripts/`: Contains all runnable scripts for each stage of the pipeline. - `config/` and `configs/`: Configuration files (Python and JSON) for model hyperparameters and training settings. - `data_loader/`: Batch iterators for different data types. - `ui/`: A Streamlit control panel for interacting with and monitoring the model. - `docs/`: A documentation site with theory and diagrams. ### Getting Started 1. **Clone and Install**: Clone the repository and install it in editable mode using `pip install -e .`. Optional extras are available for specific features (training, UI, docs). 2. **Prepare Data**: Use the provided scripts to tokenize datasets like The Pile (for pretraining), Alpaca, Dolly, and GSM8K (for instruction tuning), and Anthropic HH-RLHF (for preference learning). 3. **Build and Train the Model**: Start with a small 13M parameter model using `scripts/train_transformer.py` or use the more advanced `scripts/pretrain_base.py` for larger models with features like distributed training and gradient accumulation. 4. **Post-Train and Use the Model**: Fine-tune the base model with SFT, train a reward model, and apply RLHF techniques like DPO or PPO to align the model with human preferences. Finally, use the evaluation and chat scripts to assess and interact with your model. ### Example Output Here is an example of text generated by a trained 13M parameter model: ``` 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. ``` ### Documentation The repository includes a documentation site with more detailed theory, diagrams, and explanations, available at the provided link in the README.