This repository provides a complete implementation of a Transformer-based Large Language Model (LLM) using pure PyTorch, covering the entire pipeline from raw text pretraining to post-training alignment. The architecture is built from fundamental components, including Multi-Head Attention with causal masking, MLP blocks with ReLU activation, and pre-norm residual connections. The author demonstrates the process using three model scales: a 13 million parameter model, a 77 million parameter base, and a 406 million parameter post-training configuration. Pretraining is performed via next-token prediction using cross-entropy loss on datasets like The Pile. The repository follows a "wrap, do not rewrite" design philosophy, allowing the same backbone to be used for various post-training stages by adding specialized heads. Key alignment algorithms implemented from scratch include Supervised Fine-Tuning (SFT) using a loss mask on assistant tokens, Reward Model training via Bradley-Terry loss, and Direct Preference Optimization (DPO). Additionally, the framework supports Reinforcement Learning from Human Feedback (RLHF) through Proximal Policy Optimization (PPO) with Generalized Advantage Estimation (GAE) and Group Relative Policy Optimization (GRPO). The implementation includes utility scripts for tokenization using OpenAI's r50k_base tokenizer and HDF5-based data streaming to optimize memory efficiency during training.

Read original