A PyTorch implementation of GPT (Generative Pre-trained Transformer) from scratch, optimized for training efficiency and loss reduction. This project implements a decoder-only transformer architecture similar to GPT-2, with comprehensive training optimizations.
Features
✅ Full GPT Architecture: Complete implementation of GPT with causal self-attention, MLP blocks, and layer normalization
✅ Optimized Training: Gradient clipping, learning rate scheduling with warmup, and cosine annealing
✅ Automatic Checkpointing: Saves best model automatically when loss improves
✅ Comprehensive Comments: Well-documented code with step-by-step explanations
✅ Flexible Configuration: Easy-to-modify hyperparameters for different model sizes
✅ Device Support: Automatic device selection (CPU, CUDA, MPS for Apple Silicon)
Requirements
Python 3.7+
PyTorch 1.9+
tiktoken (for tokenization)
transformers (optional, for loading pretrained weights)
session12/
├── train_get2-8-init.py # Main training script
├── input.txt # Training data (text corpus)
├── checkpoints/ # Saved model checkpoints
│ ├── best_model.pt # Best model (lowest loss)
│ └── final_model.pt # Final model after training
└── README.md # This file
Usage
Training
Simply run the training script:
python train_get2-8-init.py
The script will:
Load and tokenize the text from input.txt
Initialize a GPT model with default configuration (GPT-2 small: 124M parameters)
Train the model with optimized hyperparameters
Automatically save the best model checkpoint when loss improves
Save the final model at the end of training
Training Configuration
You can modify these hyperparameters in the script:
python
1BATCH_SIZE =16# Batch size (number of sequences per batch)2SEQ_LENGTH =128# Sequence length (context window)3LEARNING_RATE =6e-4# Initial learning rate4WEIGHT_DECAY =0.1# Weight decay for regularization5MAX_STEPS =5000# Maximum training steps6WARMUP_STEPS =100# Learning rate warmup steps7GRAD_CLIP =1.0# Gradient clipping threshold
Loading Saved Models
To load a saved checkpoint:
python
1from train_get2-8-init import load_checkpoint, GPT
23# Load the best model4model, checkpoint = load_checkpoint('checkpoints/best_model.pt', device='cpu')56# Access checkpoint information7print(f"Step: {checkpoint['step']}")8print(f"Loss: {checkpoint['loss']:.4f}")9print(f"Best Loss: {checkpoint['best_loss']:.4f}")10print(f"Learning Rate: {checkpoint['learning_rate']}")1112# Use the model for inference13model.eval()14# ... your inference code here
Model Architecture
The model implements a GPT-style decoder-only transformer:
Components
CausalSelfAttention: Multi-head self-attention with causal masking
Prevents tokens from attending to future tokens
Implements scaled dot-product attention
Supports multiple attention heads
MLP: Feed-forward network
2-layer MLP with GELU activation
4x expansion ratio (n_embd → 4*n_embd → n_embd)
Block: Transformer block
Pre-norm architecture (LayerNorm before attention/MLP)
Residual connections around attention and MLP
GPT: Main model
Token embeddings (vocab_size × n_embd)
Position embeddings (block_size × n_embd)
Stack of transformer blocks
Language modeling head
Default Configuration (GPT-2 Small)
python
1block_size:1024# Maximum sequence length2vocab_size:50257# Vocabulary size3n_layer:12# Number of transformer blocks4n_head:12# Number of attention heads5n_embd:768# Embedding dimension
Total parameters: ~124M
Training Optimizations
The training script includes several optimizations to reduce loss:
Optimized Batch Size: Increased to 16 for better gradient estimates
Longer Sequences: Sequence length of 128 for better context understanding
Scaled Weight Initialization: Proper initialization for residual connections
Training Output
During training, you'll see output like:
Starting training with batch_size=16, seq_length=128
Total steps: 5000, Warmup steps: 100
Checkpoints will be saved to: checkpoints/
------------------------------------------------------------
✓ Saved best model checkpoint at step 0 with loss 10.5234
step 0 | lr: 6.00e-06 | loss: 10.5234 | best: 10.5234 (step 0)
step 100 | lr: 6.00e-04 | loss: 8.2341 | best: 7.8912 (step 87)
step 200 | lr: 5.98e-04 | loss: 6.5432 | best: 6.1234 (step 198)
...
------------------------------------------------------------
Training completed!
Final loss: 0.3453
Best loss: 0.1716 (achieved at step 4517)
Average loss (last 100 steps): 0.3894
Best model saved to: data/checkpoints/best_model.pt
Final model saved to: data/checkpoints/final_model.pt
Loading Pretrained Weights
The code includes a method to load pretrained GPT-2 weights from HuggingFace:
python
1# Uncomment this line in the script to load pretrained weights2model = GPT.from_pretrained('gpt2')# Options: 'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'
Text Generation
The script includes text generation code (currently commented out due to early exit). To use it:
Remove or comment out the sys.exit(0) line
The generation code uses top-k sampling with k=50
Generates multiple sequences in parallel
Notes
The model uses weight sharing between token embeddings and the output projection layer
Weight initialization follows the nanoGPT scheme with proper scaling for residual connections
The data loader implements sequential batching with automatic wrapping
Troubleshooting
Out of Memory Error:
Reduce BATCH_SIZE or SEQ_LENGTH
Use a smaller model configuration
Slow Training:
Ensure you're using GPU (CUDA or MPS) if available
