mlx-gpt2-trainer

A lightweight GPT-2 implementation for training on Apple Silicon using PyTorch. This project provides a minimal, educational implementation of the GPT-2 architecture suitable for training on custom text datasets.

Features

• Minimal Implementation: ~200 lines of clean Python code implementing GPT-2 from scratch
• Character-Level Tokenization: Simple character-level tokenizer for quick experimentation
• Configurable Architecture: Easily adjust model size, number of layers, attention heads, and training parameters
• Training & Generation: Complete training loop with validation and text generation capabilities
• Apple Silicon Optimized: Designed to run efficiently on MacBook Pro/Air with M1/M2/M3 chips

Architecture

The implementation includes all core GPT-2 components:

• Token and positional embeddings
• Multi-head self-attention with causal masking
• Feed-forward MLP blocks
• Layer normalization and residual connections
• Text generation with sampling

Requirements

pip install torch numpy

Quick Start

1. Prepare your training data: Place your text data in input.txt

2. Configure hyperparameters (in train.py):

   ctx_len = 128      # Context length
   n_emb = 128        # Embedding dimension
   dropout = 0.1      # Dropout rate
   head_size = 128    # Attention head size
   n_heads = 4        # Number of attention heads
   n_layers = 3       # Number of transformer blocks
   num_epochs = 20    # Training epochs
   batch_size = 64    # Batch size
   lr = 1e-3          # Learning rate

3. Train the model:

   python train.py

4. Generated text will be saved to completions.txt after training completes

Model Details

• Default Configuration: ~2-3M parameters
• Training Time: ~10 minutes on M2 MacBook Pro (for 1M character dataset)
• Memory: Requires ~2-4GB RAM depending on batch size

Example Output

Training on Shakespeare text for 20 epochs produces coherent (though nonsensical) Shakespeare-like dialogue:

GLOUCESTER:
But accomes mo move it.

KING EDWARD:
Where our that proclaim that I curse, or I sprithe.

CORIOLANUS:
Not want:
His bops to thy father...

Customization

Using Your Own Dataset

Simply replace input.txt with your own text file. The model will automatically:

• Build a character-level vocabulary
• Split data into 90% training / 10% validation
• Tokenize and prepare batches

Adjusting Model Size

For larger/smaller models, modify these parameters:

• n_emb: Embedding dimension (impacts model capacity)
• n_layers: Number of transformer blocks (impacts depth)
• n_heads: Number of attention heads (more heads = more diverse attention patterns)
• ctx_len: Maximum sequence length (longer = more context, more memory)

Technical Notes

• Uses PyTorch for tensor operations and automatic differentiation
• Implements scaled dot-product attention
• Includes proper weight initialization following GPT-2 paper
• Causal masking ensures autoregressive generation

Educational Purpose

This implementation is designed for learning and experimentation. For production use cases, consider:

• Using byte-pair encoding (BPE) tokenization instead of character-level
• Implementing gradient accumulation for larger effective batch sizes
• Adding model checkpointing and resumption
• Using mixed precision training (FP16/BF16)

License

MIT License - see LICENSE file for details

Acknowledgments

• Inspired by Andrej Karpathy's "GPT from scratch" tutorial
• Based on the GPT-2 architecture from OpenAI's Language Models are Unsupervised Multitask Learners

Contributing

Contributions are welcome! Feel free to:

• Report bugs
• Suggest improvements
• Submit pull requests
• Share interesting training results

Author

stukenov - GitHub