ZeroGrad-GPT

A clean, from-scratch implementation of a GPT-style Transformer in PyTorch. Designed for educational purposes, this project implements modern LLM techniques manually to provide a clear understanding of the internal mechanics.

Note: The source code is thoroughly commented in Spanish for educational purposes. 🇪🇸

Features

• Pytorch Native: Built using pure PyTorch tensors and modules.
• Modern Architecture:
  • RMSNorm: Root Mean Square Layer Normalization for better stability.
  • RoPE: Rotary Positional Embeddings for improved relative position handling.
  • Flash Attention: Support for scaled_dot_product_attention (SDPA) when available, with a manual fallback.
  • GELU: Gaussian Error Linear Unit activation.
• Configurable: Easy-to-tune parameters via GPTConfig in config.py.
• Explanatory Code: Clean, modular structure to separate concerns (Model, Layers, Attention, Training).

Quick Start

1. Requirements

• Python 3.8+
• PyTorch
• NumPy

pip install torch numpy

2. Training

To start training the model with the default "Nano" configuration (ideal for CPU testing):

python train.py

• by default, it looks for an input.txt file in the root directory.
• If input.txt is missing, it uses a built-in dummy dataset for demonstration.
• The training script will print the loss every 10 iterations and generate a sample text at the end.

Project Structure

• config.py: Configuration class defining model hyperparameters (layers, heads, embedding size, etc.).
• train.py: The main training loop. Handles data loading, model initialization, optimization, and text generation.
• model.py: The core GPT class assembling the Transformer architecture (Embedding -> Blocks -> RMSNorm -> Head).
• attention.py: Implementation of Causal Self-Attention, including Flash Attention support and RoPE integration.
• rope.py: Mathematics for Rotary Positional Embeddings (precomputing frequencies and applying rotations).
• layers.py: Helper modules like RMSNorm and the Feed-Forward Network (MLP).

License

MIT