Training a Large Language Model (LLM) From Scratch
This interactive guide provides a comprehensive, step-by-step walkthrough for training your own Large Language Model (LLM) from the ground up. We will cover everything from data preparation and tokenization to model architecture, training, and evaluation. Use the navigation on the left to jump to any step in the process.
What is a Large Language Model?
A Large Language Model (LLM) is a type of artificial intelligence (AI) that has been trained on vast amounts of text data to understand and generate human-like text. These models are built using deep learning architectures, most commonly the Transformer architecture. By learning patterns, grammar, and information from the training data, LLMs can perform a wide range of natural language processing (NLP) tasks, including:
- Text generation
- Translation
- Summarization
- Question answering
In this guide, we will build a GPT-like model, which is a type of decoder-only Transformer.
🎯 What You'll Build: A Real Working Transformer
You'll build a genuine transformer model using the same architecture that powers GPT - just optimized for learning and experimentation!
🚀 What You WILL Build
- • Real transformer architecture: Same building blocks as GPT, just smaller
- • Coherent text generation: Stories, conversations, and creative writing
- • Deep understanding: You'll know how LLMs work under the hood
- • Practical foundation: Skills to work on real LLM projects
🎓 Learning-Focused Approach
- • Manageable scale: ~10M parameters vs GPT-4's trillions
- • Quick training: ~1 hour on GPU vs months for large models
- • Quality dataset: TinyStories for coherent, meaningful outputs
- • Educational focus: Every line of code explained
📊 By the Numbers
~10M
Parameters (vs GPT-3's 175B)
~1 hour
Training on GPU (vs months for GPT-3)
~1MB
Training data (vs TB for GPT-3)
💡 Pro tip: With proper training, even small models can generate surprisingly coherent text! Projects like Microsoft's TinyStories show that small transformers can learn meaningful patterns. Focus on clean data and sufficient training iterations.
Prerequisites
Before we begin, you should have a basic understanding of the following concepts and tools. This foundation will help you grasp the core ideas and successfully follow the implementation steps.
🐍
Python Programming
This guide uses Python and the PyTorch library for all code examples.
🧠
Deep Learning Concepts
Familiarity with neural networks, tensors, and training loops is helpful.
🔥
PyTorch
We will use PyTorch to build and train our model from scratch.
Step 0: Environment Setup
Before we start building our LLM, let's set up our environment with all the necessary dependencies.
🚀 Quick Setup
Install all required packages with a single command:
# Install all dependencies
pip install torch>=2.0.0 datasets>=2.14.0 tiktoken>=0.5.0 matplotlib>=3.7.0 numpy>=1.24.0 tqdm>=4.65.0
# Or save this as requirements.txt and run:
# pip install -r requirements.txt
✅ Environment Check
Run this to verify your setup and check hardware availability:
# Check PyTorch and environment setup
import torch
import datasets
import tiktoken
import matplotlib.pyplot as plt
import numpy as np
# Check PyTorch installation and hardware
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"CUDA device: {torch.cuda.get_device_name(0)}")
print(f"CUDA memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
# Enhanced device detection for 2025
if torch.cuda.is_available():
device = 'cuda'
print(f"🚀 CUDA GPU available for training")
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = 'mps'
print(f"🍎 Apple Silicon MPS available for training")
else:
device = 'cpu'
print(f"🖥️ Using CPU for training")
print(f"📋 Selected device: {device}")
print("✅ Environment setup complete!")
Step 1: Data Preparation
We'll use the TinyStories dataset - a collection of simple, coherent stories perfect for training small language models.
🎯 Why TinyStories?
- • High quality: Generated by GPT-3.5/GPT-4
- • Simple vocabulary: Understandable by 3-4 year olds
- • Coherent output: Your model will generate readable stories
- • Full dataset: 2.1 million stories (~2GB), same as Microsoft's 2023 paper
📊 Dataset Size & Training Modes
- • Demo Mode (1k-10k stories): Quick testing, trains in minutes
- • Learning Mode (50k-200k stories): Good results, trains in hours
- • Full Dataset (2.1M stories): Best results, trains in 1-3 days
- • Microsoft's training: 21 days on V100 GPU for their paper
💡 We'll use Learning Mode (50k stories) for this tutorial, but show you how to use the full dataset!
📥 Download Dataset
One line of code downloads and prepares the entire dataset:
# Download and load the TinyStories dataset
from datasets import load_dataset
# Download TinyStories dataset (this may take a few minutes)
print("Downloading TinyStories dataset...")
dataset = load_dataset("roneneldan/TinyStories")
# Check dataset info
print(f"Dataset loaded!")
print(f"Train samples: {len(dataset['train']):,}") # Full dataset: 2,119,719 stories
print(f"Validation samples: {len(dataset['validation']):,}") # Validation: 21,990 stories
# Look at a sample story
sample_story = dataset['train'][0]['text']
print(f"\nSample story ({len(sample_story)} characters):")
print("-" * 50)
print(sample_story[:500] + "..." if len(sample_story) > 500 else sample_story)
🔧 Data Preprocessing
Convert the dataset into a format suitable for training:
# Choose your training mode:
# Demo mode: 1,000 stories (trains in ~5 minutes)
# train_subset_size = 1000
# Learning mode: 50,000 stories (trains in ~1 hour) - RECOMMENDED FOR TUTORIAL
train_subset_size = 50000 # Use 50k stories for balanced speed/quality
val_subset_size = 5000 # Use 5k stories for validation
# Production mode: Full dataset (trains in 1-3 days)
# train_subset_size = len(dataset['train']) # All 2.1M stories
# val_subset_size = len(dataset['validation']) # All 22k validation stories
# Extract text from dataset
train_texts = [item['text'] for item in dataset['train'].select(range(train_subset_size))]
val_texts = [item['text'] for item in dataset['validation'].select(range(val_subset_size))]
# Combine all training text
train_text = '\n\n'.join(train_texts)
val_text = '\n\n'.join(val_texts)
print(f"Training text: {len(train_text):,} characters")
print(f"Validation text: {len(val_text):,} characters")
print(f"Total vocabulary preview: {len(set(train_text))} unique characters")
# Save processed data (optional)
with open('train_data.txt', 'w', encoding='utf-8') as f:
f.write(train_text)
with open('val_data.txt', 'w', encoding='utf-8') as f:
f.write(val_text)
print("✅ Data preprocessing complete!")
Step 2: Tokenization
Modern LLMs use sophisticated tokenization. We'll use tiktoken, the same tokenizer as GPT-4, for professional-quality results.
🚀 Why tiktoken?
- • Industry standard: Same tokenizer used by GPT-4
- • Efficient: Handles subwords better than character-level
- • Large vocabulary: ~50k tokens vs ~100 characters
- • Better performance: Model learns faster with proper tokenization
🔧 Setup Tokenizer
Initialize the GPT-4 tokenizer and explore how it works:
# Setup GPT-4 tokenizer
import tiktoken
# Initialize the GPT-4 tokenizer
tokenizer = tiktoken.get_encoding("cl100k_base")
# Test tokenization on sample text
sample_text = "Once upon a time, there was a brave little mouse."
tokens = tokenizer.encode(sample_text)
decoded_text = tokenizer.decode(tokens)
print(f"Original: {sample_text}")
print(f"Tokens: {tokens}")
print(f"Decoded: {decoded_text}")
print(f"Number of tokens: {len(tokens)}")
# Show individual token breakdown
print("\nToken breakdown:")
for i, token in enumerate(tokens):
print(f" {i}: {token} -> '{tokenizer.decode([token])}'")
vocab_size = tokenizer.n_vocab
print(f"\nTokenizer vocabulary size: {vocab_size:,}")
print("✅ Tokenizer setup complete!")
📊 Tokenize Dataset
Convert our text data into tokens for training:
# Tokenize dataset and prepare training data
import torch
import numpy as np
from tqdm import tqdm
def tokenize_text(text, tokenizer, max_length=1024):
"""Tokenize text and split into chunks of max_length"""
tokens = tokenizer.encode(text)
# Split into chunks
chunks = []
for i in range(0, len(tokens), max_length):
chunk = tokens[i:i + max_length]
if len(chunk) == max_length: # Only keep full chunks
chunks.append(chunk)
return chunks
print("Tokenizing training data...")
train_chunks = tokenize_text(train_text, tokenizer, max_length=1024)
print("Tokenizing validation data...")
val_chunks = tokenize_text(val_text, tokenizer, max_length=1024)
# Convert to tensors
train_data = torch.tensor(np.array(train_chunks), dtype=torch.long)
val_data = torch.tensor(np.array(val_chunks), dtype=torch.long)
print(f"Training data shape: {train_data.shape}")
print(f"Validation data shape: {val_data.shape}")
print(f"Total training tokens: {train_data.numel():,}")
print(f"Total validation tokens: {val_data.numel():,}")
# Save tokenized data
torch.save(train_data, 'train_tokens.pt')
torch.save(val_data, 'val_tokens.pt')
print("✅ Dataset tokenization complete!")
🎯 Quick Test
Verify our tokenization works correctly:
# Test with a sample chunk
test_chunk = train_data[0]
print(f"Sample chunk shape: {test_chunk.shape}")
print(f"First 20 tokens: {test_chunk[:20].tolist()}")
# Decode back to text
decoded_sample = tokenizer.decode(test_chunk.tolist())
print(f"\nDecoded sample (first 200 chars):")
print("-" * 50)
print(decoded_sample[:200] + "...")
# Verify tokenizer parameters for model
print(f"\nTokenizer info for model:")
print(f"- Vocabulary size: {vocab_size:,}")
print(f"- Context length: {train_data.shape[1]:,}")
print(f"- Training sequences: {len(train_data):,}")
print(f"- Validation sequences: {len(val_data):,}")
Step 3: Model Architecture
We will build a decoder-only Transformer model, similar to the architecture of GPT. This architecture is the foundation of most modern LLMs. Click on the components below to learn more about them.
High-Level Architecture
Embedding Layer
↓
Positional Encoding
↓
Transformer Blocks
Click a component above to see its description.
Embedding Layer
This is the first layer of the model. It converts the input tokens (which are just numbers) into dense vectors of a fixed size. These vectors, or embeddings, capture semantic meaning, so similar tokens will have similar vectors.
Positional Encoding
The standard Transformer architecture doesn't inherently know the order of tokens. Positional encodings are vectors that are added to the token embeddings to give the model information about the position of each token in the sequence.
Transformer Blocks
The core of the model is a stack of identical Transformer blocks. Each block contains two main sub-layers: a multi-head self-attention mechanism and a position-wise feed-forward neural network. The self-attention mechanism allows the model to weigh the importance of different tokens in the sequence, and the feed-forward network processes this information further.
A Deeper Look: The Transformer Block
Let's zoom into a single Transformer Block. This is where the core computation happens. The block processes the input embeddings and passes its output to the next block in the stack. It uses two key concepts: self-attention and feed-forward networks, with residual connections and layer normalization to help with training stability.
Positional embeddings
Token embeddings
Tokenized text
↑
Input text for LLM
Transformer Block
+
Dropout
Feed forward NN
Layer Norm 2
+
Dropout
Multi-head Attention
Layer Norm 1
Click a component in the diagram to see its description.
Layer Norm 1
Layer Normalization stabilizes the network by normalizing the features for each token across the embedding dimension. It's applied before the main sub-layer (in this case, Multi-head Attention) to ensure a consistent distribution of inputs.
Multi-head Attention
This is the core of the Transformer. It allows the model to weigh the importance of different tokens in the input when processing a specific token. "Multi-head" means it does this multiple times in parallel (with different learned weights) and combines the results, allowing it to focus on different aspects of the relationships between tokens.
Dropout
A regularization technique to prevent overfitting. During training, it randomly sets a fraction of neuron activations to zero, forcing the network to learn more robust and redundant features.
Layer Norm 2
Another Layer Normalization step, applied before the feed-forward network. This again helps to stabilize the gradients and improve the training process.
Feed Forward NN
A simple fully connected neural network applied independently to each token's representation. It consists of two linear layers with a non-linear activation function (like ReLU or GeLU) in between. Its role is to further process the output from the attention layer, transforming it into a more useful representation.
Dropout
Another Dropout layer is applied after the feed-forward network for further regularization.
Step 4: Training the Model
Now, we will train our model on the prepared data. The goal of training is to adjust the model's parameters (or weights) so that it gets better at predicting the next token in a sequence.
Setting up the Training Loop
The training loop is an iterative process where the model learns from the data. It consists of these key steps:
- Sample a batch of data: Get a random chunk of text from the dataset.
- Forward pass: Pass the batch through the model to get predictions (logits).
- Calculate the loss: Compare the model's predictions with the actual next tokens.
- Backward pass: Calculate gradients to see how each parameter contributed to the error.
- Update parameters: Adjust the model's parameters using an optimizer to reduce the loss.
# Complete Training Setup with Progress Tracking
import torch
import torch.nn as nn
from torch.nn import functional as F
import matplotlib.pyplot as plt
from tqdm import tqdm
import time
# Hyperparameters (optimized for good results)
# Device-optimized hyperparameters
if torch.cuda.is_available():
device = 'cuda'
batch_size = 16 # CUDA optimized
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = 'mps'
batch_size = 20 # Apple Silicon optimized
else:
device = 'cpu'
batch_size = 8 # CPU conservative
block_size = 1024 # Use our tokenized chunk size
max_iters = 3000 # Fewer iterations but better data
eval_interval = 100 # More frequent evaluation
learning_rate = 3e-4 # Better learning rate
eval_iters = 50
n_embd = 384 # Larger model for better results
n_head = 6 # More attention heads
n_layer = 6 # Deeper model
dropout = 0.1
print(f"Training on device: {device}")
print(f"Model size: ~{(n_embd * n_layer * 4 * 1e-6):.1f}M parameters")
# Data loading function (works with our tokenized data)
def get_batch(split):
data = train_data if split == 'train' else val_data
ix = torch.randint(len(data), (batch_size,))
x = data[ix, :-1].to(device) # Input sequence
y = data[ix, 1:].to(device) # Target sequence (shifted by 1)
return x, y
# Loss estimation with progress tracking
@torch.no_grad()
def estimate_loss():
out = {}
model.eval()
for split in ['train', 'val']:
losses = torch.zeros(eval_iters)
for k in range(eval_iters):
X, Y = get_batch(split)
logits, loss = model(X, Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train()
return out
# Create the model (updated for tiktoken vocab size)
model = GPTLanguageModel(vocab_size, n_embd, block_size, n_head, n_layer, dropout)
model = model.to(device)
# Count parameters
total_params = sum(p.numel() for p in model.parameters())
print(f"Model parameters: {total_params/1e6:.2f}M")
# Create optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
# Training tracking
train_losses = []
val_losses = []
learning_rates = []
times = []
start_time = time.time()
print("\n🚀 Starting training...")
print("=" * 60)
# Training loop with progress bar
progress_bar = tqdm(range(max_iters), desc="Training")
for iter in progress_bar:
# Evaluation and logging
if iter % eval_interval == 0 or iter == max_iters - 1:
losses = estimate_loss()
train_loss = losses['train']
val_loss = losses['val']
# Track metrics
train_losses.append(train_loss)
val_losses.append(val_loss)
learning_rates.append(optimizer.param_groups[0]['lr'])
times.append(time.time() - start_time)
# Update progress bar
progress_bar.set_postfix({
'train_loss': f'{train_loss:.4f}',
'val_loss': f'{val_loss:.4f}'
})
# Print detailed progress
if iter % (eval_interval * 5) == 0:
elapsed = time.time() - start_time
print(f"\nStep {iter}: train={train_loss:.4f}, val={val_loss:.4f}, time={elapsed:.1f}s")
# Training step
xb, yb = get_batch('train')
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
print(f"\n✅ Training complete! Total time: {time.time() - start_time:.1f}s")
📊 Visualize Training Progress
Plot the training curves to see how well your model learned:
# Plot training progress
plt.figure(figsize=(15, 5))
# Loss curves
plt.subplot(1, 3, 1)
steps = [i * eval_interval for i in range(len(train_losses))]
plt.plot(steps, train_losses, label='Training Loss', color='#B4654A')
plt.plot(steps, val_losses, label='Validation Loss', color='#5A7D7C')
plt.xlabel('Training Steps')
plt.ylabel('Loss')
plt.title('Training Progress')
plt.legend()
plt.grid(True, alpha=0.3)
# Learning rate
plt.subplot(1, 3, 2)
plt.plot(steps, learning_rates, color='green')
plt.xlabel('Training Steps')
plt.ylabel('Learning Rate')
plt.title('Learning Rate Schedule')
plt.grid(True, alpha=0.3)
# Time per iteration
plt.subplot(1, 3, 3)
if len(times) > 1:
time_diffs = [times[i] - times[i-1] if i > 0 else times[i] for i in range(len(times))]
plt.plot(steps, time_diffs, color='purple')
plt.xlabel('Training Steps')
plt.ylabel('Time (seconds)')
plt.title('Training Speed')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('training_progress.png', dpi=150, bbox_inches='tight')
plt.show()
# Print final metrics
print(f"\n📈 Final Results:")
print(f"- Final training loss: {train_losses[-1]:.4f}")
print(f"- Final validation loss: {val_losses[-1]:.4f}")
print(f"- Best validation loss: {min(val_losses):.4f}")
print(f"- Total training time: {times[-1]:.1f} seconds")
print(f"- Average time per step: {times[-1]/max_iters:.2f}s")
Step 5: Evaluation
After training, we need to evaluate our model's performance and consider how to improve it. Evaluation helps us understand how well the model has learned and where it can be made better.
Measuring Performance with Training Loss
The loss on the training and validation sets gives us an idea of how well the model is learning. A lower loss generally means the model is better at predicting the next token. We expect to see the loss decrease over time, as shown in the example chart below.
Building the Model in Code
Here is a complete implementation of a small GPT-like model in PyTorch. This is a real transformer that you can actually train! We'll build it step by step.
First, the Self-Attention Block
This is the heart of the transformer. It allows the model to look at all previous tokens when predicting the next one.
# Multi-Head Self-Attention
import torch
import torch.nn as nn
from torch.nn import functional as F
class Head(nn.Module):
""" One head of self-attention """
def __init__(self, head_size, n_embd, block_size, dropout=0.1):
super().__init__()
self.key = nn.Linear(n_embd, head_size, bias=False)
self.query = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
# Lower triangular matrix for masking future tokens
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
self.dropout = nn.Dropout(dropout)
def forward(self, x):
B, T, C = x.shape # batch, time-step, channels
k = self.key(x) # (B, T, head_size)
q = self.query(x) # (B, T, head_size)
# Attention scores: "how much should we look at each token?"
wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5 # (B, T, T)
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
wei = F.softmax(wei, dim=-1) # (B, T, T)
wei = self.dropout(wei)
# Weighted aggregation of values
v = self.value(x) # (B, T, head_size)
out = wei @ v # (B, T, head_size)
return out
class MultiHeadAttention(nn.Module):
""" Multiple heads of self-attention in parallel """
def __init__(self, num_heads, head_size, n_embd, block_size, dropout=0.1):
super().__init__()
self.heads = nn.ModuleList([Head(head_size, n_embd, block_size, dropout) for _ in range(num_heads)])
self.proj = nn.Linear(head_size * num_heads, n_embd)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1)
out = self.dropout(self.proj(out))
return out
Next, the Feed-Forward Network
This processes the attention output through a simple neural network.
# Feed-Forward Network
class FeedForward(nn.Module):
""" A simple linear layer followed by ReLU """
def __init__(self, n_embd, dropout=0.1):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd), # Expansion
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd), # Projection back
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
Now, the Complete Transformer Block
This combines attention and feed-forward with residual connections and layer normalization.
# Transformer Block
class Block(nn.Module):
""" Transformer block: communication followed by computation """
def __init__(self, n_embd, n_head, block_size, dropout=0.1):
super().__init__()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head, head_size, n_embd, block_size, dropout)
self.ffwd = FeedForward(n_embd, dropout)
self.ln1 = nn.LayerNorm(n_embd)
self.ln2 = nn.LayerNorm(n_embd)
def forward(self, x):
# Residual connections: "add & norm"
x = x + self.sa(self.ln1(x))
x = x + self.ffwd(self.ln2(x))
return x
Finally, the Complete GPT Model
This puts everything together into a working language model!
# Complete GPT Language Model
class GPTLanguageModel(nn.Module):
def __init__(self, vocab_size, n_embd=384, block_size=1024, n_head=6, n_layer=6, dropout=0.1):
super().__init__()
# Token embedding: converts token IDs to vectors
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
# Position embedding: gives the model a sense of word order
self.position_embedding_table = nn.Embedding(block_size, n_embd)
# Stack of transformer blocks
self.blocks = nn.Sequential(*[Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
# Final layer norm
self.ln_f = nn.LayerNorm(n_embd)
# Output projection to vocabulary
self.lm_head = nn.Linear(n_embd, vocab_size)
self.block_size = block_size
def forward(self, idx, targets=None):
B, T = idx.shape
# Get embeddings
tok_emb = self.token_embedding_table(idx) # (B, T, n_embd)
pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device)) # (T, n_embd)
x = tok_emb + pos_emb # (B, T, n_embd)
# Forward through transformer blocks
x = self.blocks(x) # (B, T, n_embd)
x = self.ln_f(x) # (B, T, n_embd)
logits = self.lm_head(x) # (B, T, vocab_size)
# Calculate loss if targets provided
if targets is None:
loss = None
else:
B, T, C = logits.shape
logits = logits.view(B*T, C)
targets = targets.view(B*T)
loss = F.cross_entropy(logits, targets)
return logits, loss
def generate(self, idx, max_new_tokens):
# idx is (B, T) array of indices in the current context
for _ in range(max_new_tokens):
# Crop context to block_size
idx_cond = idx[:, -self.block_size:]
# Get predictions
logits, loss = self(idx_cond)
# Focus on last time step
logits = logits[:, -1, :] # (B, C)
# Apply softmax to get probabilities
probs = F.softmax(logits, dim=-1) # (B, C)
# Sample from the distribution
idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
# Append to the sequence
idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
return idx
🎭 Generate Stories
Now for the exciting part - let's see what stories your model can create!
# Generate stories with your trained model
model.eval()
def generate_story(prompt="Once upon a time", max_new_tokens=300, temperature=0.8):
"""Generate a story starting with the given prompt"""
# Encode the prompt
context = torch.tensor(tokenizer.encode(prompt), dtype=torch.long, device=device).unsqueeze(0)
# Generate new tokens
with torch.no_grad():
generated = model.generate(context, max_new_tokens)
# Decode and return the story
story = tokenizer.decode(generated[0].tolist())
return story
# Generate multiple stories
print("🎭 Story Generation Results")
print("=" * 60)
prompts = [
"Once upon a time",
"The little girl",
"In a magical forest",
"The brave mouse",
"One sunny day"
]
for i, prompt in enumerate(prompts, 1):
print(f"\n📖 Story #{i}: Starting with '{prompt}'")
print("-" * 40)
story = generate_story(prompt, max_new_tokens=200, temperature=0.8)
print(story)
print()
# Save your best stories
with open('generated_stories.txt', 'w', encoding='utf-8') as f:
for prompt in prompts:
story = generate_story(prompt, max_new_tokens=300)
f.write(f"Prompt: {prompt}\n")
f.write(f"Story: {story}\n")
f.write("-" * 60 + "\n\n")
print("✅ Stories saved to 'generated_stories.txt'!")
🎉 What to Expect
If your model trained well on TinyStories, you should see:
- • Coherent sentences: Proper grammar and sentence structure
- • Story elements: Characters, settings, simple plots
- • Child-like vocabulary: Simple words appropriate for young children
- • Narrative flow: Stories that make sense and have a beginning, middle, end
💡 Pro tip: Adjust the temperature parameter - lower values (0.3-0.5) make more focused text, higher values (0.8-1.2) make more creative but potentially less coherent text.
🍎 Apple Silicon Optimization (2025 Research)
The final code section incorporates the latest 2025 research on Apple Silicon optimization:
🚀 Performance Gains
- • 2-3x faster training on M1/M2/M3/M4 chips
- • Unified memory eliminates CPU-GPU transfers
- • TF32 acceleration for matrix operations
- • Optimized batch sizes based on 2025 benchmarks
🔧 Key Optimizations
- • Automatic MPS detection with fallback support
- • Float32 precision for stable MPS training
- • Memory management with periodic cache clearing
- • Device-specific hyperparameters for optimal performance
Based on ICML 2025 research: Apple demonstrated fine-tuning 7B parameter LLMs on iPhone, showcasing MLX and MPS capabilities for on-device AI training.
🚀 Complete Working Example
Here's everything put together in one script that you can actually run! This trains a tiny GPT on the TinyStories dataset with full Apple Silicon optimization.
📓 Jupyter Notebook Version Available!
For a more interactive experience with step-by-step execution, progress visualization, and detailed explanations, download our complete Jupyter notebook:
📥 Download Jupyter NotebookIncludes: Environment setup, progress tracking, interactive story generation, and detailed explanations
💡 Try this: Save this code as train_gpt.py and run it! This complete example uses TinyStories dataset for training.
"""
Complete GPT training script with TinyStories dataset.
Usage: python train_gpt.py
This will automatically download the TinyStories dataset and train a small GPT model.
"""
# Setup GPT-4 tokenizer
import tiktoken
import torch
import torch.nn as nn
from torch.nn import functional as F
from datasets import load_dataset
# Enhanced device detection with Apple Silicon support (2025 research-based)
import os
import platform
def get_optimal_device():
"""Detect and configure the best available device based on 2025 research"""
# Set MPS fallback environment variable for unsupported operations
os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK', '1')
if torch.cuda.is_available():
device = 'cuda'
device_name = torch.cuda.get_device_name(0)
memory_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
print(f"🚀 CUDA detected: {device_name}")
print(f"💾 GPU Memory: {memory_gb:.1f} GB")
# Enable CUDA optimizations
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = 'mps'
print(f"🍎 MPS detected: Apple Silicon with unified memory")
print(f"💻 System: {platform.processor()}")
print(f"📊 Unified Memory: Optimized for M1/M2/M3/M4 chips")
# Enable MPS optimizations (based on 2025 research)
torch.backends.mps.allow_tf32 = True
print(f"⚡ TF32 acceleration enabled")
else:
device = 'cpu'
print(f"🖥️ Using CPU: {platform.processor()}")
torch.set_num_threads(4)
print(f"📋 Selected device: {device}")
return device
# Detect optimal device
device = get_optimal_device()
# Device-specific hyperparameters (optimized based on 2025 research)
if device == 'cuda':
# CUDA optimizations
batch_size = 16
block_size = 1024
max_iters = 3000
dtype = torch.float16 # Mixed precision for CUDA
use_amp = True
print("🚀 CUDA configuration: Mixed precision enabled")
elif device == 'mps':
# Apple Silicon optimizations (based on ICML 2025 research)
batch_size = 20 # MPS-optimized batch size
block_size = 1024 # Good context length for unified memory
max_iters = 3500 # Slightly more iterations for MPS
dtype = torch.float32 # Better stability on MPS (2025 research)
use_amp = False # MPS doesn't support all AMP operations
print("🍎 Apple Silicon configuration: Unified memory optimized")
else:
# CPU fallback
batch_size = 8 # Smaller for CPU
block_size = 512 # Reduced context for memory efficiency
max_iters = 1000 # Fewer iterations for time constraints
dtype = torch.float32
use_amp = False
print("🖥️ CPU configuration: Memory efficient")
# Common hyperparameters
eval_interval = 100
learning_rate = 3e-4
eval_iters = 200
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2
print(f"✅ Optimized for {device}: batch_size={batch_size}, dtype={dtype}")
# Memory optimization for Apple Silicon
if device == 'mps':
# Clear MPS cache for optimal performance
torch.mps.empty_cache()
print(f"🧹 MPS memory cache cleared")
elif device == 'cuda':
# Clear CUDA cache
torch.cuda.empty_cache()
print(f"🧹 CUDA memory cache cleared")
# Download and load the TinyStories dataset
print("Downloading TinyStories dataset...")
dataset = load_dataset("roneneldan/TinyStories")
train_data = dataset['train']
# Setup tokenizer
tokenizer = tiktoken.get_encoding("cl100k_base")
vocab_size = tokenizer.n_vocab
# Prepare training data
print("Tokenizing TinyStories dataset...")
# Choose how many stories to use:
# Demo: 1000 stories (quick test, ~5 min training)
# Learning: 50000 stories (good results, ~1 hour training)
# Full: len(train_data) stories (best results, 1-3 days training)
num_stories = 1000 # Using demo mode for this example
tokenized_texts = []
for i, example in enumerate(train_data):
if i >= num_stories:
break
tokens = tokenizer.encode(example['text'])
tokenized_texts.extend(tokens)
print(f"Using {num_stories:,} stories ({num_stories/len(train_data)*100:.1f}% of full dataset)")
# Convert to tensor
data = torch.tensor(tokenized_texts, dtype=torch.long)
n = int(0.9 * len(data))
train_data_tensor = data[:n]
val_data_tensor = data[n:]
print(f"Dataset size: {len(data)} tokens")
print(f"Training set: {len(train_data_tensor)} tokens")
print(f"Validation set: {len(val_data_tensor)} tokens")
# Data loading
def get_batch(split):
data = train_data_tensor if split == 'train' else val_data_tensor
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([data[i:i+block_size] for i in ix])
y = torch.stack([data[i+1:i+block_size+1] for i in ix])
x, y = x.to(device), y.to(device)
return x, y
# Model definition (same as above, but all in one place)
class Head(nn.Module):
""" One head of self-attention """
def __init__(self, head_size, n_embd, block_size, dropout=0.1):
super().__init__()
self.key = nn.Linear(n_embd, head_size, bias=False)
self.query = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
self.dropout = nn.Dropout(dropout)
def forward(self, x):
B, T, C = x.shape
k = self.key(x) # (B, T, head_size)
q = self.query(x) # (B, T, head_size)
# Compute attention scores
wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5 # (B, T, T)
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
wei = F.softmax(wei, dim=-1) # (B, T, T)
wei = self.dropout(wei)
# Apply attention to values
v = self.value(x) # (B, T, head_size)
out = wei @ v # (B, T, head_size)
return out
class MultiHeadAttention(nn.Module):
""" Multiple heads of self-attention in parallel """
def __init__(self, num_heads, head_size, n_embd, block_size, dropout=0.1):
super().__init__()
self.heads = nn.ModuleList([Head(head_size, n_embd, block_size, dropout) for _ in range(num_heads)])
self.proj = nn.Linear(head_size * num_heads, n_embd)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1)
out = self.dropout(self.proj(out))
return out
class FeedForward(nn.Module):
""" A simple linear layer followed by a non-linearity """
def __init__(self, n_embd, dropout=0.1):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
class Block(nn.Module):
""" Transformer block: communication followed by computation """
def __init__(self, n_embd, n_head, block_size, dropout=0.1):
super().__init__()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head, head_size, n_embd, block_size, dropout)
self.ffwd = FeedForward(n_embd, dropout)
self.ln1 = nn.LayerNorm(n_embd)
self.ln2 = nn.LayerNorm(n_embd)
def forward(self, x):
# Residual connections: "add & norm"
x = x + self.sa(self.ln1(x))
x = x + self.ffwd(self.ln2(x))
return x
class GPTLanguageModel(nn.Module):
def __init__(self, vocab_size, n_embd=384, block_size=1024, n_head=6, n_layer=6, dropout=0.1):
super().__init__()
# Token embedding: converts token IDs to vectors
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
# Position embedding: gives the model a sense of word order
self.position_embedding_table = nn.Embedding(block_size, n_embd)
# Stack of transformer blocks
self.blocks = nn.Sequential(*[Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
# Final layer norm
self.ln_f = nn.LayerNorm(n_embd)
# Output projection to vocabulary
self.lm_head = nn.Linear(n_embd, vocab_size)
self.block_size = block_size
def forward(self, idx, targets=None):
B, T = idx.shape
# Get embeddings
tok_emb = self.token_embedding_table(idx) # (B, T, n_embd)
pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device)) # (T, n_embd)
x = tok_emb + pos_emb # (B, T, n_embd)
# Forward through transformer blocks
x = self.blocks(x) # (B, T, n_embd)
x = self.ln_f(x) # (B, T, n_embd)
logits = self.lm_head(x) # (B, T, vocab_size)
if targets is None:
loss = None
else:
B, T, C = logits.shape
logits = logits.view(B*T, C)
targets = targets.view(B*T)
loss = F.cross_entropy(logits, targets)
return logits, loss
def generate(self, idx, max_new_tokens):
for _ in range(max_new_tokens):
idx_cond = idx[:, -block_size:]
logits, loss = self(idx_cond)
logits = logits[:, -1, :]
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
return idx
# Create model
model = GPTLanguageModel(vocab_size, n_embd, block_size, n_head, n_layer, dropout)
m = model.to(device)
print(f"{sum(p.numel() for p in m.parameters())/1e6:.2f}M parameters")
# Create optimizer with device-specific optimizations
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate,
betas=(0.9, 0.95), weight_decay=0.1)
# Mixed precision setup (CUDA only, based on 2025 research)
scaler = None
if use_amp and device == 'cuda':
scaler = torch.cuda.amp.GradScaler()
print("⚡ Mixed precision training enabled for CUDA")
# Memory-efficient loss estimation
@torch.no_grad()
def estimate_loss():
out = {}
model.eval()
for split in ['train', 'val']:
losses = torch.zeros(eval_iters)
for k in range(eval_iters):
X, Y = get_batch(split)
# Use appropriate context for mixed precision
if use_amp and device == 'cuda':
with torch.cuda.amp.autocast():
logits, loss = model(X, Y)
else:
logits, loss = model(X, Y)
losses[k] = loss.item()
# Clear cache periodically for memory efficiency
if k % 20 == 0 and device in ['mps', 'cuda']:
if device == 'mps':
torch.mps.empty_cache()
elif device == 'cuda':
torch.cuda.empty_cache()
out[split] = losses.mean()
model.train()
return out
# Enhanced training loop with Apple Silicon optimization
print(f"🚀 Starting optimized training on {device}...")
print(f"📊 Configuration: batch_size={batch_size}, max_iters={max_iters}")
for iter in range(max_iters):
# Evaluate model performance
if iter % eval_interval == 0 or iter == max_iters - 1:
losses = estimate_loss()
print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
# Memory management for Apple Silicon
if device == 'mps' and iter % (eval_interval * 2) == 0:
torch.mps.empty_cache()
# Forward pass with device-specific optimization
xb, yb = get_batch('train')
if use_amp and device == 'cuda':
# Mixed precision training for CUDA
with torch.cuda.amp.autocast():
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
# Standard precision for MPS and CPU (2025 research recommendation)
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
print(f"✅ Training completed on {device}!")
# Performance summary based on device
print(f"\n📈 Training Performance Summary:")
if device == 'mps':
print(f"🍎 Apple Silicon Benefits Utilized:")
print(f" ✅ Unified Memory Architecture - No CPU-GPU data transfers")
print(f" ✅ TF32 Acceleration - Optimized matrix operations")
print(f" ✅ MPS-Optimized Batch Size - {batch_size} (vs 16 default)")
print(f" ✅ Float32 Precision - Stable training (2025 research)")
print(f" ⚡ Expected speedup: 2-3x faster than CPU")
elif device == 'cuda':
print(f"🚀 CUDA Benefits Utilized:")
print(f" ✅ Mixed Precision Training - Float16 acceleration")
print(f" ✅ TensorCore Optimization - Enhanced matrix ops")
print(f" ✅ Memory Management - Automatic cache clearing")
print(f" ⚡ Expected speedup: 5-10x faster than CPU")
else:
print(f"🖥️ CPU Training - Consider upgrading to Apple Silicon or CUDA for better performance")
# Generate some text
print("\n" + "="*50)
print("Generating sample text...")
print("="*50 + "\n")
context = torch.zeros((1, 1), dtype=torch.long, device=device)
generated = model.generate(context, max_new_tokens=500)[0].tolist()
print(tokenizer.decode(generated))
# Save the model
torch.save(model.state_dict(), 'gpt_model.pth')
print("\nModel saved to gpt_model.pth")
📝 What to Expect
- • Training time: ~5-30 minutes on CPU, ~1-5 minutes on GPU
- • Loss values: Should start around 4.5 and drop to ~1.5
- • Generated text: With enough training, expect coherent sentences and patterns
- • Memory usage: ~500MB-1GB RAM
⚡ Hardware Optimization: CUDA vs MPS vs CPU
Learn how to maximize training performance across different hardware setups and choose the best configuration for your system.
🔍 Device Detection & Setup
First, let's detect what hardware acceleration is available on your system and configure PyTorch optimally:
# Hardware detection and optimal device selection
import torch
import platform
def get_optimal_device():
"""Detect and configure the best available device"""
# Check for CUDA (NVIDIA GPUs)
if torch.cuda.is_available():
device = 'cuda'
device_name = torch.cuda.get_device_name(0)
memory_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
print(f"🚀 CUDA detected: {device_name}")
print(f"💾 GPU Memory: {memory_gb:.1f} GB")
# Enable optimizations for CUDA
torch.backends.cudnn.benchmark = True # Optimize for consistent input sizes
torch.backends.cuda.matmul.allow_tf32 = True # Faster matmul on Ampere GPUs
# Check for MPS (Apple Silicon M1/M2/M3)
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = 'mps'
print(f"🍎 MPS detected: Apple Silicon")
print(f"💻 System: {platform.processor()}")
# MPS optimizations
torch.backends.mps.allow_tf32 = True # Enable TF32 if available
# Fallback to CPU
else:
device = 'cpu'
print(f"🖥️ Using CPU: {platform.processor()}")
# CPU optimizations
torch.set_num_threads(4) # Limit threads to prevent oversubscription
print(f"📋 Selected device: {device}")
return device
# Detect optimal device
device = get_optimal_device()
# Additional memory optimization
if device == 'cuda':
# Clear cache and set memory fraction
torch.cuda.empty_cache()
print(f"🧹 GPU memory cleared")
elif device == 'mps':
# MPS memory is managed automatically
print(f"🔄 MPS memory management: automatic")
print("✅ Hardware optimization complete!")
📊 Performance Comparison
🚀 NVIDIA CUDA
Training Speed:
⭐⭐⭐⭐⭐
Memory Efficiency:
⭐⭐⭐⭐⭐
Model Size Limit:
8-80+ GB
Best for: Large models, long training sessions
🍎 Apple MPS
Training Speed:
⭐⭐⭐⭐
Memory Efficiency:
⭐⭐⭐⭐
Model Size Limit:
16-64 GB
Best for: MacBooks, energy efficiency
🖥️ CPU Only
Training Speed:
⭐⭐
Memory Efficiency:
⭐⭐⭐
Model Size Limit:
System RAM
Best for: Small models, learning
⚙️ Device-Specific Training Configurations
Here are optimized hyperparameters for different hardware setups:
# Device-optimized training configurations
def get_training_config(device, total_params):
"""Get optimized training configuration based on device and model size"""
base_config = {
'learning_rate': 3e-4,
'max_iters': 3000,
'eval_interval': 100,
'batch_size': 16,
'block_size': 1024,
}
if device == 'cuda':
# CUDA optimizations
gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9
if gpu_memory >= 12: # High-end GPU (RTX 4090, A100, etc.)
config = {
**base_config,
'batch_size': 32, # Larger batches
'block_size': 2048, # Longer context
'max_iters': 5000, # More training
'grad_accumulation': 1, # No accumulation needed
}
print("🚀 High-end GPU config: Large batches, long context")
elif gpu_memory >= 8: # Mid-range GPU (RTX 3080, 4070, etc.)
config = {
**base_config,
'batch_size': 24, # Medium batches
'block_size': 1024, # Standard context
'max_iters': 4000,
'grad_accumulation': 2, # Some accumulation
}
print("⚡ Mid-range GPU config: Balanced performance")
else: # Budget GPU (RTX 3060, etc.)
config = {
**base_config,
'batch_size': 16, # Smaller batches
'block_size': 512, # Shorter context
'max_iters': 3000,
'grad_accumulation': 4, # More accumulation
}
print("💾 Budget GPU config: Memory-efficient")
# Enable mixed precision for all CUDA configs
config['use_amp'] = True
config['dtype'] = 'float16'
elif device == 'mps':
# MPS (Apple Silicon) optimizations
config = {
**base_config,
'batch_size': 20, # MPS-optimized batch size
'block_size': 1024, # Good context length
'max_iters': 3500,
'grad_accumulation': 2,
'use_amp': False, # MPS doesn't support all AMP ops
'dtype': 'float32', # Better stability on MPS
}
print("🍎 Apple Silicon config: MPS-optimized")
else: # CPU
# CPU optimizations - focus on memory efficiency
config = {
**base_config,
'batch_size': 8, # Small batches for CPU
'block_size': 256, # Short context
'max_iters': 1000, # Fewer iterations
'grad_accumulation': 8, # High accumulation
'use_amp': False, # No AMP on CPU
'dtype': 'float32',
}
print("🖥️ CPU config: Memory and time efficient")
# Adjust for model size
param_mb = total_params * 4 / 1e6 # Rough memory estimate
print(f"📊 Model size: ~{param_mb:.0f}MB ({total_params/1e6:.1f}M parameters)")
return config
# Get device-optimized configuration
training_config = get_training_config(device, sum(p.numel() for p in model.parameters()))
# Apply configuration
batch_size = training_config['batch_size']
block_size = training_config['block_size']
max_iters = training_config['max_iters']
learning_rate = training_config['learning_rate']
eval_interval = training_config['eval_interval']
print(f"✅ Configuration applied:")
print(f" - Batch size: {batch_size}")
print(f" - Context length: {block_size}")
print(f" - Training iterations: {max_iters}")
print(f" - Learning rate: {learning_rate}")
🧠 Memory Optimization Techniques
🚀 GPU Memory Optimization
# GPU memory optimization
if device == 'cuda':
# Clear cache before training
torch.cuda.empty_cache()
# Enable memory efficient attention
torch.backends.cuda.enable_flash_sdp(True)
# Gradient checkpointing for large models
# model.gradient_checkpointing = True
# Monitor memory usage
def print_gpu_memory():
if torch.cuda.is_available():
allocated = torch.cuda.memory_allocated() / 1e9
cached = torch.cuda.memory_reserved() / 1e9
print(f"GPU Memory - Allocated: {allocated:.1f}GB, Cached: {cached:.1f}GB")
print_gpu_memory()
💾 General Memory Tips
# General memory optimization
import gc
# Reduce model precision where possible
if training_config.get('use_amp', False):
scaler = torch.cuda.amp.GradScaler()
print("⚡ Mixed precision enabled")
# Clear unnecessary variables
def cleanup_memory():
gc.collect()
if device == 'cuda':
torch.cuda.empty_cache()
elif device == 'mps':
torch.mps.empty_cache()
# Use gradient accumulation for effective larger batches
effective_batch_size = batch_size * training_config.get('grad_accumulation', 1)
print(f"📊 Effective batch size: {effective_batch_size}")
# Pin memory for faster data loading (CUDA only)
pin_memory = (device == 'cuda')
print(f"📌 Pin memory: {pin_memory}")
📈 Performance Monitoring
Monitor your training performance and system resources:
# Performance monitoring setup
import time
import psutil
class PerformanceMonitor:
def __init__(self, device):
self.device = device
self.start_time = None
self.iteration_times = []
def start_training(self):
self.start_time = time.time()
print("🚀 Training performance monitoring started")
def log_iteration(self, iteration, loss):
current_time = time.time()
if len(self.iteration_times) > 0:
iter_time = current_time - self.iteration_times[-1]
else:
iter_time = 0
self.iteration_times.append(current_time)
if iteration % 100 == 0:
# System metrics
cpu_percent = psutil.cpu_percent()
memory_percent = psutil.virtual_memory().percent
# Device-specific metrics
if self.device == 'cuda':
gpu_memory = torch.cuda.memory_allocated() / 1e9
print(f"Iter {iteration}: Loss={loss:.4f}, Time={iter_time:.2f}s, GPU={gpu_memory:.1f}GB, CPU={cpu_percent:.1f}%, RAM={memory_percent:.1f}%")
else:
print(f"Iter {iteration}: Loss={loss:.4f}, Time={iter_time:.2f}s, CPU={cpu_percent:.1f}%, RAM={memory_percent:.1f}%")
def get_summary(self):
if self.start_time:
total_time = time.time() - self.start_time
avg_iter_time = total_time / len(self.iteration_times) if self.iteration_times else 0
print(f"\n📊 Training Summary:")
print(f" - Total time: {total_time:.1f} seconds")
print(f" - Average per iteration: {avg_iter_time:.3f} seconds")
print(f" - Iterations per second: {1/avg_iter_time:.2f}")
# Initialize performance monitor
perf_monitor = PerformanceMonitor(device)
🔧 Common Issues & Solutions
❌ CUDA Out of Memory
- Reduce batch_size (try 8, 4, or even 2)
- Reduce block_size (try 512 or 256)
- Increase gradient_accumulation_steps
- Enable gradient checkpointing
🍎 MPS Compatibility Issues
- Use dtype='float32' instead of 'float16'
- Disable mixed precision (use_amp=False)
- Update to latest PyTorch version
- Some operations may fall back to CPU automatically
🖥️ CPU Training Too Slow
- Reduce model size (n_embd=256, n_layer=4)
- Use smaller batch_size=4 and block_size=128
- Reduce max_iters to 500-1000
- Consider using smaller dataset subset