82b1759c5a
- Changed from Apache 2.0 only to dual license model - Apache 2.0 for research, education, and non-commercial use - Commercial license required for profit-making use - Added citation requirement as condition of use for academic purposes - Created CITATION.cff file for automatic citation suggestions - Updated LICENSE, LICENSE.txt, README.md, and documentation - Citation formats provided (BibTeX and text) - Contact information for commercial licensing inquiries
861 lines
31 KiB
Markdown
861 lines
31 KiB
Markdown
# SheepOp LLM 🐑➡️🤖
|
||
|
||
A modern language model implementation from scratch, incorporating insights from recent research papers.
|
||
|
||
## License
|
||
|
||
This project is available under a **dual license**:
|
||
|
||
### Apache 2.0 License (Research & Non-Commercial Use)
|
||
|
||
**Free for:** Academic research, educational purposes, personal projects, open source contributions, and non-commercial use.
|
||
|
||
### Commercial License (Commercial Use)
|
||
|
||
**Requires a commercial license for:** Commercial products, SaaS applications, revenue-generating applications, and any profit-making use.
|
||
|
||
**Contact:** carlos.gutierrez@carg.dev for commercial licensing inquiries.
|
||
|
||
### Citation
|
||
|
||
Please cite this work in academic publications.
|
||
|
||
See [LICENSE](../LICENSE) or [LICENSE.txt](../LICENSE.txt) for the full license text.
|
||
|
||
## Table of Contents
|
||
|
||
- [What is This?](#what-is-this)
|
||
- [Features](#features)
|
||
- [Quick Start](#quick-start)
|
||
- [Mathematical Foundations](#mathematical-foundations)
|
||
- [Architecture Explained](#architecture-explained)
|
||
- [Project Structure](#project-structure)
|
||
- [Installation](#installation)
|
||
- [Usage](#usage)
|
||
- [Configuration](#configuration)
|
||
- [Diagrams](#diagrams)
|
||
- [References](#references)
|
||
|
||
## What is This?
|
||
|
||
A Transformer-based language model implementing autoregressive next-token prediction using multi-head self-attention, positional encoding, and modern training optimizations (mixed precision, gradient accumulation, KV caching).
|
||
|
||
The model learns to write by reading large amounts of text, discovering patterns like "after 'the cat' usually comes 'sat' or 'ran'", enabling it to generate coherent sentences. It processes text sequentially, predicting each next word based on the context provided by previous words.
|
||
|
||
## Features
|
||
|
||
- **Transformer Architecture**: Multi-head self-attention mechanism from "Attention Is All You Need"
|
||
- **Long Context Support**: Efficient handling of long sequences with Rotary Positional Encoding (RoPE)
|
||
- **Training Optimizations**: Mixed precision training, gradient accumulation, gradient clipping
|
||
- **Modern Best Practices**: Pre-norm architecture, GELU activation, weight tying
|
||
- **Comprehensive Evaluation**: Perplexity, accuracy metrics, and generation utilities
|
||
|
||
## Quick Start
|
||
|
||
**Option 1: Automated Setup (Recommended)**
|
||
|
||
```bash
|
||
# Run the setup script - it handles everything automatically!
|
||
./setup.sh
|
||
|
||
# Then activate the virtual environment
|
||
source venv/bin/activate
|
||
|
||
# Download data and train
|
||
python3 download_large_data.py wiki --version 103
|
||
python3 train.py --data data/wikitext_103.txt --config config.json --device cuda
|
||
```
|
||
|
||
**Option 2: Manual Setup**
|
||
|
||
```bash
|
||
# 1. Create and activate virtual environment (REQUIRED on modern Linux systems)
|
||
python3 -m venv venv
|
||
source venv/bin/activate
|
||
|
||
# 2. Upgrade pip
|
||
pip install --upgrade pip
|
||
|
||
# 3. Install dependencies
|
||
pip install -r requirements.txt
|
||
|
||
# 4. Download data
|
||
python3 download_large_data.py wiki --version 103
|
||
|
||
# 5. Train
|
||
python3 train.py --data data/wikitext_103.txt --config config.json --device cuda
|
||
|
||
# 6. Generate
|
||
python3 inference.py --checkpoint checkpoints/checkpoint_epoch_10.pt \
|
||
--prompt "The future of artificial intelligence" --device cuda
|
||
```
|
||
|
||
**Note:** On modern Debian/Ubuntu systems (especially Python 3.12+), pip prevents system-wide installations to protect system packages. Always use a virtual environment (`python3 -m venv venv`) before installing dependencies.
|
||
|
||
See [GETTING_STARTED.md](GETTING_STARTED.md) for detailed instructions.
|
||
|
||
## Mathematical Foundations
|
||
|
||
### 1. Token Embedding
|
||
|
||
Words are converted into numerical representations that the model can process. Each word (token) is assigned a unique ID, which is then mapped to a dense vector representation - think of it as converting words into a format the computer can understand and manipulate mathematically.
|
||
|
||
**Mathematical Formulation**:
|
||
|
||
Given a vocabulary of size $V$ and token ID $t \in \{0, 1, \ldots, V-1\}$:
|
||
|
||
$$
|
||
\mathbf{E}_t = \text{EmbeddingTable}[t] \in \mathbb{R}^{d_{\text{model}}}
|
||
$$
|
||
|
||
where $\mathbf{E} \in \mathbb{R}^{V \times d_{\text{model}}}$ is the learnable embedding matrix.
|
||
|
||
### 2. Positional Encoding
|
||
|
||
Word order is crucial in language - "Cat bites dog" is very different from "Dog bites cat". Since transformers process all tokens simultaneously, we need to inject positional information so the model knows which word comes first, second, third, etc.
|
||
|
||
**Mathematical Formulation**:
|
||
|
||
**Sinusoidal Positional Encoding** (for position $pos$ and dimension $i$):
|
||
|
||
$$
|
||
PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)
|
||
$$
|
||
|
||
$$
|
||
PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)
|
||
$$
|
||
|
||
The final embedding combines token and position:
|
||
|
||
$$
|
||
\mathbf{h}_i = \mathbf{E}_{t_i} + PE(i)
|
||
$$
|
||
|
||
where $t_i$ is the token at position $i$.
|
||
|
||
### 3. Multi-Head Self-Attention
|
||
|
||
Attention mechanisms allow the model to understand relationships between words in a sentence. When reading "The cat sat on the mat", the model learns to connect "cat" with "sat" because they're related. Multi-head attention enables the model to focus on different types of relationships simultaneously - syntax, semantics, and context.
|
||
|
||
**Mathematical Formulation**:
|
||
|
||
Given input $\mathbf{X} \in \mathbb{R}^{n \times d_{\text{model}}}$ with $n$ tokens:
|
||
|
||
1. **Project to Query, Key, Value**:
|
||
|
||
```math
|
||
\mathbf{Q} = \mathbf{X}\mathbf{W}_Q, \quad \mathbf{K} = \mathbf{X}\mathbf{W}_K, \quad \mathbf{V} = \mathbf{X}\mathbf{W}_V
|
||
|
||
where $\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$
|
||
```
|
||
|
||
2. **Scaled Dot-Product Attention**:
|
||
|
||
```math
|
||
\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}\right)\mathbf{V}
|
||
```
|
||
|
||
3. **Multi-Head Attention** splits into $h$ heads:
|
||
|
||
```math
|
||
\text{head}_i = \text{Attention}(\mathbf{Q}_i, \mathbf{K}_i, \mathbf{V}_i)
|
||
|
||
\text{MultiHead}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\mathbf{W}_O
|
||
|
||
where $d_k = d_{\text{model}} / h$ (each head has dimension $d_k$).
|
||
```
|
||
|
||
4. **Causal Masking** (for autoregressive generation):
|
||
|
||
```math
|
||
M_{ij} = \begin{cases}
|
||
0 & \text{if } i \geq j \\
|
||
-\infty & \text{if } i < j
|
||
\end{cases}
|
||
\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}, M) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}} + M\right)\mathbf{V}
|
||
```
|
||
|
||
### 4. Feed-Forward Network
|
||
|
||
After attention identifies which words relate to each other, the feed-forward network performs non-linear transformations on the information. This step allows the model to synthesize and combine the attended information, similar to mixing ingredients to create a final output.
|
||
|
||
**Mathematical Formulation**:
|
||
|
||
```math
|
||
\text{FFN}(\mathbf{x}) = \text{GELU}(\mathbf{x}\mathbf{W}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2
|
||
|
||
where $\mathbf{W}_1 \in \mathbb{R}^{d_{\text{model}} \times d_{ff}}$, $\mathbf{W}_2 \in \mathbb{R}^{d_{ff} \times d_{\text{model}}}$, and typically $d_{ff} = 4 \times d_{\text{model}}$.
|
||
```
|
||
|
||
**GELU Activation**:
|
||
|
||
```math
|
||
\text{GELU}(x) = x \cdot \Phi(x) = x \cdot \frac{1}{2}\left(1 + \text{erf}\left(\frac{x}{\sqrt{2}}\right)\right)
|
||
|
||
where $\Phi(x)$ is the cumulative distribution function of the standard normal distribution.
|
||
```
|
||
|
||
### 5. Transformer Block (Pre-Norm Architecture)
|
||
|
||
**For Kids**: Imagine building with LEGO blocks. Each block does something special (attention or thinking), and you can stack many blocks on top of each other to build something amazing!
|
||
|
||
**Mathematical Formulation**:
|
||
|
||
For a transformer block with input $\mathbf{x}$:
|
||
|
||
1. **Self-Attention Sublayer**:
|
||
|
||
```math
|
||
\mathbf{x}' = \mathbf{x} + \text{Dropout}\left(\text{MultiHead}(\text{LayerNorm}(\mathbf{x}))\right)
|
||
```
|
||
|
||
2. **Feed-Forward Sublayer**:
|
||
|
||
```math
|
||
\mathbf{x}'' = \mathbf{x}' + \text{Dropout}\left(\text{FFN}(\text{LayerNorm}(\mathbf{x}'))\right)
|
||
```
|
||
|
||
**Layer Normalization**:
|
||
|
||
```math
|
||
\text{LayerNorm}(\mathbf{x}) = \gamma \odot \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta
|
||
|
||
where $\mu = \frac{1}{d}\sum_{i=1}^d x_i$, $\sigma^2 = \frac{1}{d}\sum_{i=1}^d (x_i - \mu)^2$, and $\gamma, \beta$ are learnable parameters.
|
||
```
|
||
|
||
### 6. Complete Forward Pass
|
||
|
||
**For Kids**: Think of it like an assembly line! The words go through many stations (layers), each doing something special, until finally we get a prediction of what word comes next.
|
||
|
||
**Mathematical Formulation**:
|
||
|
||
Given input token sequence $\mathbf{t} = [t_1, t_2, \ldots, t_n]$:
|
||
|
||
$$
|
||
\mathbf{h}_0 = \mathbf{E}[\mathbf{t}] + PE(\mathbf{t}) \quad \text{(token embeddings + positional encoding)}
|
||
$$
|
||
|
||
$$
|
||
\mathbf{h}_1 = \text{TransformerBlock}_1(\mathbf{h}_0)
|
||
$$
|
||
|
||
$$
|
||
\mathbf{h}_2 = \text{TransformerBlock}_2(\mathbf{h}_1)
|
||
$$
|
||
|
||
$$
|
||
\vdots
|
||
$$
|
||
|
||
$$
|
||
\mathbf{h}_L = \text{TransformerBlock}_L(\mathbf{h}_{L-1})
|
||
$$
|
||
|
||
$$
|
||
\mathbf{o} = \text{LayerNorm}(\mathbf{h}_L) \quad \text{(final normalization)}
|
||
$$
|
||
|
||
$$
|
||
\mathbf{y} = \mathbf{o}\mathbf{W}_{\text{out}} \quad \text{(output projection)}
|
||
$$
|
||
|
||
$$
|
||
p(t_{n+1} | t_1, \ldots, t_n) = \text{softmax}(\mathbf{y}_n) \quad \text{(next token probability)}
|
||
$$
|
||
|
||
### 7. Training Objective (Cross-Entropy Loss)
|
||
|
||
**For Kids**: When we train, we show the computer a sentence and ask "what word comes next?" If it guesses wrong, we say "try again!" and it learns from its mistake.
|
||
|
||
**Mathematical Formulation**:
|
||
|
||
For a sequence of tokens $\mathbf{t} = [t_1, t_2, \ldots, t_n]$:
|
||
|
||
$$
|
||
\mathcal{L} = -\frac{1}{n-1}\sum_{i=1}^{n-1} \log p(t_{i+1} | t_1, \ldots, t_i)
|
||
$$
|
||
|
||
**Perplexity** (measure of model confidence):
|
||
|
||
$$
|
||
\text{Perplexity} = \exp(\mathcal{L}) = \exp\left(-\frac{1}{n-1}\sum_{i=1}^{n-1} \log p(t_{i+1} | t_1, \ldots, t_i)\right)
|
||
$$
|
||
|
||
Lower perplexity = better model!
|
||
|
||
### 8. Text Generation (Autoregressive Sampling)
|
||
|
||
**For Kids**: The computer writes one word at a time. After each word, it thinks "what word makes sense next?" and picks one.
|
||
|
||
**Mathematical Formulation**:
|
||
|
||
Given prompt $\mathbf{p} = [p_1, \ldots, p_k]$:
|
||
|
||
1. Initialize: $\mathbf{s} = \mathbf{p}$
|
||
2. For $i = k+1, \ldots, k+n$:
|
||
|
||
$$
|
||
\mathbf{h}_i = \text{Transformer}(\mathbf{s})
|
||
$$
|
||
|
||
$$
|
||
\mathbf{p}_i = \text{softmax}(\mathbf{h}_i / \tau) \quad \text{(with temperature $\tau \geq 1$)}
|
||
$$
|
||
|
||
$$
|
||
t_i \sim \text{Sample}(\mathbf{p}_i) \quad \text{(sample next token)}
|
||
$$
|
||
|
||
$$
|
||
\mathbf{s} = \mathbf{s} \cup \{t_i\} \quad \text{(append to sequence)}
|
||
$$
|
||
|
||
**Top-k Sampling**:
|
||
|
||
$$
|
||
\mathbf{p}_i' = \begin{cases}
|
||
\mathbf{p}_i[j] & \text{if } j \in \text{top}_k(\mathbf{p}_i) \\
|
||
0 & \text{otherwise}
|
||
\end{cases}
|
||
$$
|
||
|
||
**Top-p (Nucleus) Sampling**:
|
||
Find smallest set $S$ such that $\sum_{j \in S} \mathbf{p}_i[j] \geq p$, then:
|
||
|
||
$$
|
||
\mathbf{p}_i'[j] = \begin{cases}
|
||
\mathbf{p}_i[j] & \text{if } j \in S \\
|
||
0 & \text{otherwise}
|
||
\end{cases}
|
||
$$
|
||
|
||
### 9. Optimization (AdamW)
|
||
|
||
**For Kids**: Learning is like climbing a mountain. You take small steps in the right direction. AdamW is like having a smart guide that knows which direction is best and adjusts your step size automatically!
|
||
|
||
**Mathematical Formulation**:
|
||
|
||
For parameter $\theta$ with gradient $\mathbf{g}_t$ at step $t$:
|
||
|
||
**Momentum**:
|
||
|
||
$$
|
||
\mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1 - \beta_1) \mathbf{g}_t
|
||
$$
|
||
|
||
**RMSprop**:
|
||
|
||
$$
|
||
\mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1 - \beta_2) \mathbf{g}_t^2
|
||
$$
|
||
|
||
**Bias Correction**:
|
||
|
||
$$
|
||
\hat{\mathbf{m}}_t = \frac{\mathbf{m}_t}{1 - \beta_1^t}, \quad \hat{\mathbf{v}}_t = \frac{\mathbf{v}_t}{1 - \beta_2^t}
|
||
$$
|
||
|
||
**Parameter Update**:
|
||
|
||
$$
|
||
\theta_t = \theta_{t-1} - \eta \left(\frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t} + \epsilon} + \lambda \theta_{t-1}\right)
|
||
$$
|
||
|
||
where $\eta$ is learning rate, $\lambda$ is weight decay, and typically $\beta_1 = 0.9$, $\beta_2 = 0.999$.
|
||
|
||
### 10. Gradient Clipping
|
||
|
||
**For Kids**: Sometimes the computer gets too excited and tries to learn too fast (like running too fast down a hill). We slow it down so it doesn't fall!
|
||
|
||
**Mathematical Formulation**:
|
||
|
||
$$
|
||
\mathbf{g}_{\text{clipped}} = \begin{cases}
|
||
\mathbf{g} & \text{if } \|\mathbf{g}\| \leq \theta_{\max} \\
|
||
\mathbf{g} \cdot \frac{\theta_{\max}}{\|\mathbf{g}\|} & \text{if } \|\mathbf{g}\| > \theta_{\max}
|
||
\end{cases}
|
||
$$
|
||
|
||
where $\theta_{\max}$ is the maximum gradient norm (typically 1.0).
|
||
|
||
## Architecture Explained
|
||
|
||
### High-Level Overview
|
||
|
||
**For Kids**:
|
||
|
||
```
|
||
📚 Text Input → 🔤 Turn into Numbers → 🧠 Thinking Layers → ✨ Predict Next Word → 📝 Generate Text
|
||
```
|
||
|
||
**For Scientists**:
|
||
|
||
```
|
||
Token IDs → Embeddings → Positional Encoding → N× Transformer Blocks → Output Projection → Logits → Sampling → Text
|
||
```
|
||
|
||
### Detailed Architecture
|
||
|
||
```
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ INPUT TEXT │
|
||
│ "The cat sat on the mat" │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ TOKENIZATION │
|
||
│ [1, 45, 123, 67, 45, 234] (each word → number) │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ TOKEN EMBEDDING │
|
||
│ E ∈ ℝ^(V×d_model) | Each token → vector of size d_model │
|
||
│ [1, 45, ...] → [[0.1, 0.3, ...], [0.2, -0.1, ...], ...] │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ POSITIONAL ENCODING │
|
||
│ PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) │
|
||
│ h_i = E[t_i] + PE(i) (add position info) │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ TRANSFORMER BLOCK 1 │
|
||
│ ┌──────────────────────────────────────────────┐ │
|
||
│ │ Pre-Norm Multi-Head Attention │ │
|
||
│ │ Attention(Q, K, V) = softmax(QK^T/√d_k)V │ │
|
||
│ │ + Residual Connection │ │
|
||
│ └──────────────────────────────────────────────┘ │
|
||
│ ┌──────────────────────────────────────────────┐ │
|
||
│ │ Pre-Norm Feed-Forward Network │ │
|
||
│ │ FFN(x) = GELU(xW₁ + b₁)W₂ + b₂ │ │
|
||
│ │ + Residual Connection │ │
|
||
│ └──────────────────────────────────────────────┘ │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ TRANSFORMER BLOCK 2 │
|
||
│ (Same structure as Block 1) │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
... (N-2 more blocks) ...
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ TRANSFORMER BLOCK N │
|
||
│ (Same structure) │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ FINAL LAYER NORM │
|
||
│ LayerNorm(h_L) │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ OUTPUT PROJECTION │
|
||
│ y = h_L · W_out (W_out ∈ ℝ^(d_model × V)) │
|
||
│ Output: logits for each token in vocabulary │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ SOFTMAX │
|
||
│ p(t | context) = softmax(y) │
|
||
│ Probability distribution over vocabulary │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ SAMPLING │
|
||
│ t_{next} ~ Sample(p) (with temperature, top-k, top-p) │
|
||
└──────────────────────────┬──────────────────────────────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ OUTPUT TEXT │
|
||
│ "The cat sat on the mat and..." │
|
||
└─────────────────────────────────────────────────────────────────┘
|
||
```
|
||
|
||
## Project Structure
|
||
|
||
```
|
||
sheepOp/
|
||
├── models/ # Model architectures
|
||
│ ├── __init__.py # Module exports
|
||
│ ├── transformer.py # Main transformer model
|
||
│ ├── attention.py # Attention mechanisms
|
||
│ ├── blocks.py # Building blocks (FFN, TransformerBlock)
|
||
│ ├── optimized_attention.py # KV caching, optimized inference
|
||
│ └── prefetching.py # Data prefetching utilities
|
||
├── data/ # Data loading utilities
|
||
│ └── __init__.py # Dataset and tokenizer
|
||
├── training/ # Training utilities
|
||
│ ├── __init__.py # Trainer class
|
||
│ └── metrics.py # Training metrics and plotting
|
||
├── config.py # Configuration management
|
||
├── config.json # Example configuration file
|
||
├── train.py # Training script
|
||
├── inference.py # Inference script
|
||
├── example.py # Usage examples
|
||
├── utils.py # Evaluation utilities
|
||
└── requirements.txt # Dependencies
|
||
```
|
||
|
||
## Installation
|
||
|
||
**Important:** On modern Debian/Ubuntu systems (Python 3.12+), you must use a virtual environment. The system prevents system-wide pip installations to protect system packages.
|
||
|
||
```bash
|
||
# 1. Create virtual environment
|
||
python3 -m venv venv
|
||
|
||
# 2. Activate virtual environment
|
||
source venv/bin/activate
|
||
|
||
# 3. Upgrade pip
|
||
pip install --upgrade pip
|
||
|
||
# 4. Install dependencies
|
||
pip install -r requirements.txt
|
||
|
||
# 5. For large dataset downloads (optional)
|
||
pip install datasets
|
||
```
|
||
|
||
**Alternative:** Use the automated setup script which handles everything:
|
||
```bash
|
||
./setup.sh
|
||
source venv/bin/activate
|
||
```
|
||
|
||
## Usage
|
||
|
||
### Training
|
||
|
||
```bash
|
||
python3 train.py \
|
||
--data data/amazon_reviews.txt \
|
||
--config config.json \
|
||
--device cuda
|
||
```
|
||
|
||
### Inference
|
||
|
||
```bash
|
||
python3 inference.py \
|
||
--checkpoint checkpoints/checkpoint_epoch_10.pt \
|
||
--prompt "The future of artificial intelligence" \
|
||
--optimized \
|
||
--device cuda
|
||
```
|
||
|
||
## Configuration
|
||
|
||
See `config.json` for all available settings. Key parameters:
|
||
|
||
- **Model**: `vocab_size`, `d_model`, `num_layers`, `num_heads`
|
||
- **Training**: `batch_size`, `learning_rate`, `max_epochs`
|
||
- **Data**: `max_length`, `data_dir`
|
||
|
||
## Diagrams
|
||
|
||
### Complete Model Architecture Flow
|
||
|
||
```mermaid
|
||
graph TB
|
||
subgraph "Input Processing"
|
||
A[Text Input] --> B[Tokenization]
|
||
B --> C[Token Embedding<br/>E ∈ ℝ^V×d]
|
||
C --> D[Positional Encoding<br/>PE pos,2i = sin pos/10000^2i/d]
|
||
D --> E[Embedding + Position<br/>h = E + PE]
|
||
end
|
||
|
||
subgraph "Transformer Stack"
|
||
E --> F[Transformer Block 1]
|
||
F --> G[Transformer Block 2]
|
||
G --> H[...]
|
||
H --> I[Transformer Block N]
|
||
end
|
||
|
||
subgraph "Transformer Block Detail"
|
||
F --> F1[Layer Norm 1]
|
||
F1 --> F2[Multi-Head Attention<br/>QK^T/√d_k → softmax → V]
|
||
F2 --> F3[Residual + Dropout]
|
||
F --> F3
|
||
F3 --> F4[Layer Norm 2]
|
||
F4 --> F5[Feed-Forward<br/>GELU xW₁ + b₁W₂ + b₂]
|
||
F5 --> F6[Residual + Dropout]
|
||
F3 --> F6
|
||
end
|
||
|
||
subgraph "Output Processing"
|
||
I --> J[Final Layer Norm]
|
||
J --> K[Output Projection<br/>y = hW_out]
|
||
K --> L[Softmax<br/>p = softmax y]
|
||
L --> M[Sampling<br/>t ~ p with temp, top-k, top-p]
|
||
M --> N[Generated Text]
|
||
end
|
||
|
||
style A fill:#e1f5ff
|
||
style N fill:#fff4e1
|
||
style F fill:#e8f5e9
|
||
style F2 fill:#ffe1f5
|
||
style F5 fill:#ffe1f5
|
||
```
|
||
|
||
### Attention Mechanism Visualization
|
||
|
||
```mermaid
|
||
graph LR
|
||
subgraph "Input"
|
||
X[Input Tokens<br/>x₁, x₂, ..., xₙ]
|
||
end
|
||
|
||
subgraph "Q, K, V Projections"
|
||
X --> Q[Query Q<br/>Q = XW_Q]
|
||
X --> K[Key K<br/>K = XW_K]
|
||
X --> V[Value V<br/>V = XW_V]
|
||
end
|
||
|
||
subgraph "Attention Computation"
|
||
Q --> M[Matrix Multiply<br/>QK^T]
|
||
K --> M
|
||
M --> S[Scale<br/>÷√d_k]
|
||
S --> Mask{Causal Mask?}
|
||
Mask -->|Yes| CM[Mask<br/>M_ij = -∞ if i < j]
|
||
Mask -->|No| SM[Skip Mask]
|
||
CM --> Soft[Softmax<br/>exp scores / Σexp]
|
||
SM --> Soft
|
||
Soft --> WV[Weighted Sum<br/>Attention = softmax scores × V]
|
||
V --> WV
|
||
end
|
||
|
||
subgraph "Multi-Head"
|
||
WV --> Split[Split into h heads]
|
||
Split --> H1[Head 1]
|
||
Split --> H2[Head 2]
|
||
Split --> H3[...]
|
||
Split --> Hh[Head h]
|
||
H1 --> Concat[Concatenate]
|
||
H2 --> Concat
|
||
H3 --> Concat
|
||
Hh --> Concat
|
||
Concat --> Out[Output Projection<br/>WO]
|
||
end
|
||
|
||
style X fill:#e1f5ff
|
||
style Out fill:#fff4e1
|
||
style Soft fill:#ffe1f5
|
||
style WV fill:#e8f5e9
|
||
```
|
||
|
||
### Training Loop Flow
|
||
|
||
```mermaid
|
||
graph TD
|
||
A[Start Training] --> B[Load Data]
|
||
B --> C[Create Tokenizer]
|
||
C --> D[Build Vocabulary]
|
||
D --> E[Create DataLoader]
|
||
E --> F[Initialize Model]
|
||
F --> G[Setup Optimizer<br/>AdamW]
|
||
G --> H[Setup LR Scheduler<br/>Cosine Annealing]
|
||
|
||
H --> I[For Each Epoch]
|
||
I --> J[For Each Batch]
|
||
|
||
J --> K[Forward Pass<br/>y = Model x]
|
||
K --> L[Compute Loss<br/>L = -log p t_next]
|
||
L --> M[Backward Pass<br/>∇L]
|
||
M --> N[Gradient Accumulation]
|
||
N --> O{Accumulation<br/>Complete?}
|
||
O -->|No| J
|
||
O -->|Yes| P[Gradient Clipping<br/>Clip gradient norm]
|
||
P --> Q[Optimizer Step<br/>θ = θ - η∇L]
|
||
Q --> R[LR Scheduler Step]
|
||
R --> S{End of<br/>Epoch?}
|
||
S -->|No| J
|
||
S -->|Yes| T[Evaluate on<br/>Validation Set]
|
||
T --> U[Compute Perplexity<br/>exp L]
|
||
U --> V{Best<br/>Model?}
|
||
V -->|Yes| W[Save Checkpoint]
|
||
V -->|No| X[Save Regular Checkpoint]
|
||
W --> Y{More<br/>Epochs?}
|
||
X --> Y
|
||
Y -->|Yes| I
|
||
Y -->|No| Z[Training Complete]
|
||
|
||
style A fill:#e1f5ff
|
||
style Z fill:#fff4e1
|
||
style K fill:#e8f5e9
|
||
style L fill:#ffe1f5
|
||
style P fill:#fff4e1
|
||
```
|
||
|
||
### Inference Flow with Sampling
|
||
|
||
```mermaid
|
||
graph TD
|
||
A[Load Checkpoint] --> B[Initialize Model]
|
||
B --> C[Set Eval Mode]
|
||
C --> D[Input Prompt]
|
||
D --> E[Tokenize<br/>t = t1, t2, ..., tk]
|
||
E --> F[Encode<br/>h = E t + PE]
|
||
|
||
F --> G[Forward Pass<br/>y = Model h]
|
||
G --> H[Get Logits<br/>y ∈ ℝ^V]
|
||
H --> I[Apply Temperature<br/>p = softmax y/τ]
|
||
|
||
I --> J{Top-k<br/>Filter?}
|
||
J -->|Yes| K[Keep Top-k<br/>p' = filter p, k]
|
||
J -->|No| L[p' = p]
|
||
K --> M{Top-p<br/>Filter?}
|
||
L --> M
|
||
M -->|Yes| N[Nucleus Sampling<br/>p' = filter p', p]
|
||
M -->|No| O[p'' = p']
|
||
N --> O
|
||
|
||
O --> P[Sample Token<br/>t_i ~ p'']
|
||
P --> Q[Append to Sequence<br/>s = s ∪ t_i]
|
||
Q --> R{Max Length<br/>Reached?}
|
||
R -->|No| G
|
||
R -->|Yes| S[Decode Tokens<br/>text = decode s]
|
||
S --> T[Output Text]
|
||
|
||
style A fill:#e1f5ff
|
||
style T fill:#fff4e1
|
||
style G fill:#e8f5e9
|
||
style P fill:#ffe1f5
|
||
```
|
||
|
||
### Component Interaction
|
||
|
||
```mermaid
|
||
graph TB
|
||
subgraph "Configuration"
|
||
CFG[config.py<br/>Config Classes]
|
||
CFGJSON[config.json<br/>JSON Settings]
|
||
end
|
||
|
||
subgraph "Models"
|
||
TRANS[transformer.py<br/>TransformerModel]
|
||
ATT[attention.py<br/>MultiHeadAttention]
|
||
BLOCKS[blocks.py<br/>TransformerBlock, FFN]
|
||
OPT[optimized_attention.py<br/>KV Cache, Optimized Inference]
|
||
end
|
||
|
||
subgraph "Data"
|
||
DATA[data/__init__.py<br/>TextDataset, SimpleTokenizer]
|
||
TOKEN[SimpleTokenizer<br/>encode/decode]
|
||
DATASET[TextDataset<br/>PyTorch Dataset]
|
||
end
|
||
|
||
subgraph "Training"
|
||
TRAIN[training/__init__.py<br/>Trainer]
|
||
METRICS[training/metrics.py<br/>TrainingMetrics]
|
||
end
|
||
|
||
subgraph "Scripts"
|
||
TRAIN_SCRIPT[train.py<br/>Training Entry Point]
|
||
INFER[inference.py<br/>Generation Script]
|
||
EX[example.py<br/>Usage Examples]
|
||
end
|
||
|
||
subgraph "Utils"
|
||
UTILS[utils.py<br/>Evaluation Functions]
|
||
end
|
||
|
||
CFG --> TRAIN_SCRIPT
|
||
CFGJSON --> TRAIN_SCRIPT
|
||
TRANS --> TRAIN_SCRIPT
|
||
TRANS --> INFER
|
||
ATT --> TRANS
|
||
BLOCKS --> TRANS
|
||
OPT --> TRANS
|
||
DATA --> TRAIN_SCRIPT
|
||
DATA --> INFER
|
||
TOKEN --> DATASET
|
||
DATASET --> TRAINER
|
||
TRAIN --> TRAIN_SCRIPT
|
||
METRICS --> TRAIN
|
||
UTILS --> TRAIN_SCRIPT
|
||
|
||
style TRANS fill:#e1f5ff
|
||
style TRAINER fill:#fff4e1
|
||
style DATA fill:#e8f5e9
|
||
style ATT fill:#ffe1f5
|
||
```
|
||
|
||
## Mathematical Summary Table
|
||
|
||
<table>
|
||
<tr>
|
||
<th>Concept</th>
|
||
<th>Formula</th>
|
||
<th>Description</th>
|
||
</tr>
|
||
<tr>
|
||
<td><strong>Token Embedding</strong></td>
|
||
<td>$$\mathbf{E}_t = \text{EmbeddingTable}[t]$$</td>
|
||
<td>Maps token ID to vector</td>
|
||
</tr>
|
||
<tr>
|
||
<td><strong>Positional Encoding</strong></td>
|
||
<td>$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d}}\right)$$</td>
|
||
<td>Adds position information</td>
|
||
</tr>
|
||
<tr>
|
||
<td><strong>Attention</strong></td>
|
||
<td>$$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$</td>
|
||
<td>Computes attention weights</td>
|
||
</tr>
|
||
<tr>
|
||
<td><strong>Feed-Forward</strong></td>
|
||
<td>$$\text{FFN}(x) = \text{GELU}(xW_1 + b_1)W_2 + b_2$$</td>
|
||
<td>Non-linear transformation</td>
|
||
</tr>
|
||
<tr>
|
||
<td><strong>Layer Norm</strong></td>
|
||
<td>$$\text{LayerNorm}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$$</td>
|
||
<td>Normalizes activations</td>
|
||
</tr>
|
||
<tr>
|
||
<td><strong>Loss</strong></td>
|
||
<td>$$\mathcal{L} = -\frac{1}{n}\sum_{i=1}^{n} \log p(t_{i+1} \mid t_1, \ldots, t_i)$$</td>
|
||
<td>Cross-entropy loss</td>
|
||
</tr>
|
||
<tr>
|
||
<td><strong>Perplexity</strong></td>
|
||
<td>$$\text{Perplexity} = \exp(\mathcal{L})$$</td>
|
||
<td>Measure of model confidence</td>
|
||
</tr>
|
||
<tr>
|
||
<td><strong>AdamW Update</strong></td>
|
||
<td>$$\theta_t = \theta_{t-1} - \eta\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda\theta_{t-1}\right)$$</td>
|
||
<td>Optimizer step</td>
|
||
</tr>
|
||
</table>
|
||
|
||
## References
|
||
|
||
- **Attention Is All You Need** (Vaswani et al., 2017) - Original Transformer paper
|
||
- **Optimizing LLM Inference and Retrieval** - Production RAG systems optimizations
|
||
- **RoPE: Rotary Position Embedding** - Efficient positional encoding for long sequences
|
||
- Various papers on LLM training, hallucinations, and long context handling
|
||
|
||
---
|
||
|
||
**For Kids**: You've learned how computers can read and write! Just like you practice writing stories, computers practice by reading millions of books. The more they practice, the better they get! 🎉
|
||
|
||
**For Scientists**: This implementation follows modern best practices including pre-norm architecture, weight tying, mixed precision training, and efficient inference optimizations. The codebase is modular, well-documented, and optimized for both training and production deployment.
|