DEV Community

Ankit Arora
Ankit Arora

Posted on

Building a ~10M-Parameter GPT-Style Language Model on Simple Home Computer

A Complete Guide with Code Examples

Generative AI has transformed the way we build intelligent applications. Large Language Models (LLMs) such as GPT, Llama, and Gemma can write code, answer questions, summarize documents, and even reason through complex problems. While using these existing models is straightforward, understanding how they actually work is a different challenge.

Many engineers begin by fine-tuning existing models or using frameworks like Hugging Face. While these tools are incredibly valuable, they often hide the mechanics that make language models possible. Building a small GPT-style model from scratch is one of the best ways to understand the underlying concepts.

This article walks through the complete process of building a miniature Generative AI model using only PyTorch. By the end, you'll understand every major component of a Transformer-based language model and have a solid foundation for exploring larger architectures.

Although the model we'll build contains around 10 million parameters—tiny compared to today's billion-parameter models—it contains the same fundamental building blocks used in modern LLMs.


What You'll Build

Our goal is to create a character-level GPT model capable of learning patterns from text and generating new text one token at a time.

The journey consists of five major stages:

  1. Build a tokenizer
  2. Create a dataset for next-token prediction
  3. Implement the Transformer architecture
  4. Train the model
  5. Generate text

Each stage builds on the previous one, gradually transforming raw text into an intelligent language model.


Setting Up Your Environment

First, prepare your Mac for machine learning:

# Install Xcode Command Line Tools
xcode-select --install

# Create project directory
mkdir my-gpt-from-scratch && cd my-gpt-from-scratch

# Create virtual environment
python3 -m venv .venv
source .venv/bin/activate

# Install dependencies
pip install torch numpy tqdm
Enter fullscreen mode Exit fullscreen mode

Verify PyTorch can use your Mac's GPU:

import torch
print(f"MPS available: {torch.backends.mps.is_available()}")
print(f"MPS built: {torch.backends.mps.is_built()}")
# Should output: True, True
Enter fullscreen mode Exit fullscreen mode

Project Structure

my-gpt-from-scratch/
│
├── tokenizer.py      # Character-level tokenizer
├── data.py           # Dataset and batching
├── model.py          # GPT architecture
├── train.py          # Training loop
├── generate.py       # Text generation
├── config.py         # Hyperparameters
├── utils.py          # Helper functions
└── data/
    └── input.txt     # Training corpus
Enter fullscreen mode Exit fullscreen mode

Stage 1: Teaching the Computer to Read

Before a neural network can understand language, it must convert text into numbers. This process is called tokenization. For simplicity, we'll use a character-level tokenizer, where every unique character receives a unique integer.

# tokenizer.py

class CharTokenizer:
    def __init__(self, text):
        chars = sorted(list(set(text)))

        self.stoi = {ch:i for i,ch in enumerate(chars)}
        self.itos = {i:ch for ch,i in self.stoi.items()}

        self.vocab_size = len(chars)

    def encode(self, text):
        """Convert text to token IDs"""
        return [self.stoi[c] for c in text]

    def decode(self, ids):
        """Convert token IDs back to text"""
        return ''.join(self.itos[i] for i in ids)
Enter fullscreen mode Exit fullscreen mode

Example Usage:

text = "hello world"
tok = CharTokenizer(text)

ids = tok.encode("hello")
print(ids)  # [1, 0, 2, 2, 3]

print(tok.decode(ids))  # hello
Enter fullscreen mode Exit fullscreen mode

Stage 2: Creating Training Examples

Language models learn a surprisingly simple task: predict the next token.

Consider the sentence: "Machine learning is amazing"

During training, the model repeatedly sees examples like:

Input Target
Machine learning
Machine learning is
Machine learning is amazing

Instead of feeding the entire dataset every iteration, we'll create small chunks called contexts (or sequences).

# data.py

import torch

class TextDataset:
    def __init__(self, data, block_size):
        """
        data: list of token IDs
        block_size: context length (e.g., 256)
        """
        self.data = torch.tensor(data, dtype=torch.long)
        self.block_size = block_size

    def get_batch(self, batch_size):
        """Get a batch of input-target pairs"""
        # Random starting positions
        ix = torch.randint(
            len(self.data) - self.block_size - 1,
            (batch_size,)
        )

        # Input sequences
        x = torch.stack([
            self.data[i:i+self.block_size]
            for i in ix
        ])

        # Target sequences (shifted by 1)
        y = torch.stack([
            self.data[i+1:i+self.block_size+1]
            for i in ix
        ])

        return x, y
Enter fullscreen mode Exit fullscreen mode

Stage 3: Building the Transformer

Modern language models are built using the Transformer architecture, introduced in the landmark paper "Attention Is All You Need."

Step 3.1: Token Embedding

Neural networks cannot operate directly on integer IDs. Instead, every token is mapped to a dense vector called an embedding.

import torch.nn as nn

# In your model class:
token_embedding = nn.Embedding(
    vocab_size,     # Number of unique tokens
    embedding_dim   # Size of each embedding vector (e.g., 384)
)
Enter fullscreen mode Exit fullscreen mode

Step 3.2: Positional Embedding

Unlike humans, Transformers have no built-in understanding of sequence order. Without positional information, "dog bites man" and "man bites dog" appear identical.

# Positional embedding
position_embedding = nn.Embedding(
    block_size,     # Maximum sequence length
    embedding_dim   # Same as token embedding
)

# In forward pass:
B, T = x.shape  # Batch size, Sequence length

tok_emb = token_embedding(x)  # Shape: (B, T, n_embed)

pos = position_embedding(
    torch.arange(T, device=x.device)
)  # Shape: (T, n_embed)

x = tok_emb + pos  # Add position information
Enter fullscreen mode Exit fullscreen mode

Step 3.3: Single-Head Self-Attention

Self-attention allows the model to weigh the importance of different tokens. Each token asks: "Which previous tokens should I pay attention to?"

class Head(nn.Module):
    def __init__(self, n_embed, head_size, block_size):
        super().__init__()

        self.key = nn.Linear(n_embed, head_size, bias=False)
        self.query = nn.Linear(n_embed, head_size, bias=False)
        self.value = nn.Linear(n_embed, head_size, bias=False)

        # Causal mask (prevents looking at future tokens)
        self.register_buffer(
            "tril",
            torch.tril(torch.ones(block_size, block_size))
        )

    def forward(self, x):
        B, T, C = x.shape

        # Compute key, query, value
        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)  # (B, T, T)

        # Scale (stabilizes gradients)
        wei = wei / (k.shape[-1] ** 0.5)

        # Apply causal mask
        wei = wei.masked_fill(
            self.tril[:T, :T] == 0,
            float("-inf")
        )

        # Convert to probabilities
        wei = wei.softmax(dim=-1)  # (B, T, T)

        # Apply attention to values
        v = self.value(x)  # (B, T, head_size)
        out = wei @ v      # (B, T, head_size)

        return out
Enter fullscreen mode Exit fullscreen mode

Step 3.4: Multi-Head Attention

Because language contains many types of relationships, we use multiple attention heads. Different heads can specialize in different patterns.

class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, n_embed, block_size):
        super().__init__()

        head_size = n_embed // num_heads

        self.heads = nn.ModuleList([
            Head(n_embed, head_size, block_size)
            for _ in range(num_heads)
        ])

        self.proj = nn.Linear(n_embed, n_embed)

    def forward(self, x):
        # Run all heads and concatenate
        out = torch.cat(
            [h(x) for h in self.heads],
            dim=-1
        )

        return self.proj(out)
Enter fullscreen mode Exit fullscreen mode

Step 3.5: Feed-Forward Network

After attention gathers information, each token passes through a small neural network that transforms the representation.

class FeedForward(nn.Module):
    def __init__(self, n_embed):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(n_embed, 4 * n_embed),  # Expand
            nn.GELU(),                         # Activation
            nn.Linear(4 * n_embed, n_embed)    # Contract
        )

    def forward(self, x):
        return self.net(x)
Enter fullscreen mode Exit fullscreen mode

Step 3.6: Transformer Block

Layer normalization and residual connections make training stable.

class Block(nn.Module):
    def __init__(self, n_embed, n_head, block_size):
        super().__init__()

        self.sa = MultiHeadAttention(n_head, n_embed, block_size)
        self.ffwd = FeedForward(n_embed)

        self.ln1 = nn.LayerNorm(n_embed)
        self.ln2 = nn.LayerNorm(n_embed)

    def forward(self, x):
        # Attention with residual connection
        x = x + self.sa(self.ln1(x))

        # Feed-forward with residual connection
        x = x + self.ffwd(self.ln2(x))

        return x
Enter fullscreen mode Exit fullscreen mode

Step 3.7: Complete GPT Model

class GPT(nn.Module):
    def __init__(self, vocab_size, n_embed, n_head, n_layer, block_size):
        super().__init__()

        self.token_embedding = nn.Embedding(vocab_size, n_embed)
        self.position_embedding = nn.Embedding(block_size, n_embed)

        self.blocks = nn.Sequential(*[
            Block(n_embed, n_head, block_size)
            for _ in range(n_layer)
        ])

        self.ln = nn.LayerNorm(n_embed)
        self.lm_head = nn.Linear(n_embed, vocab_size)

    def forward(self, idx):
        B, T = idx.shape

        # Get embeddings
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))

        x = tok_emb + pos_emb

        # Pass through transformer blocks
        x = self.blocks(x)

        # Final layer norm
        x = self.ln(x)

        # Project to vocabulary
        logits = self.lm_head(x)

        return logits
Enter fullscreen mode Exit fullscreen mode

Stage 4: Training the Model

Configuration

First, define your hyperparameters:

# config.py

# Model architecture
vocab_size = 100     # Character-level vocabulary
block_size = 256     # Context length
n_embed = 384        # Embedding dimension
n_head = 6           # Number of attention heads
n_layer = 8          # Number of transformer layers

# Training
batch_size = 64
learning_rate = 3e-4
max_iters = 10000
eval_interval = 100

# Generation
max_new_tokens = 300
temperature = 0.8
Enter fullscreen mode Exit fullscreen mode

Training Loop

# train.py

import torch
import torch.nn.functional as F
from config import *
from tokenizer import CharTokenizer
from data import TextDataset
from model import GPT

# Load data
with open('data/input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# Create tokenizer
tokenizer = CharTokenizer(text)
vocab_size = tokenizer.vocab_size

# Create dataset
data = tokenizer.encode(text)
dataset = TextDataset(data, block_size)

# Initialize model
model = GPT(vocab_size, n_embed, n_head, n_layer, block_size)

# Move to GPU if available
device = 'mps' if torch.backends.mps.is_available() else 'cpu'
model = model.to(device)

# Optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# Training loop
for step in range(max_iters):
    # Get batch
    xb, yb = dataset.get_batch(batch_size)
    xb, yb = xb.to(device), yb.to(device)

    # Forward pass
    logits = model(xb)

    # Compute loss
    loss = F.cross_entropy(
        logits.view(-1, vocab_size),
        yb.view(-1)
    )

    # Backward pass
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Log progress
    if step % eval_interval == 0:
        print(f"Step {step}: loss = {loss.item():.4f}")

# Save model
torch.save(model.state_dict(), 'model.pt')
Enter fullscreen mode Exit fullscreen mode

Stage 5: Generating Text

# generate.py

import torch
from config import *
from tokenizer import CharTokenizer
from model import GPT

def generate(model, tokenizer, start_text, max_new_tokens, temperature=0.8):
    """Generate text from a starting prompt"""
    model.eval()

    # Encode the start text
    idx = torch.tensor([tokenizer.encode(start_text)], dtype=torch.long)
    idx = idx.to(next(model.parameters()).device)

    for _ in range(max_new_tokens):
        # Get predictions
        with torch.no_grad():
            logits = model(idx)
            logits = logits[:, -1, :]  # Get last token's logits

            # Apply temperature
            logits = logits / temperature

            # Convert to probabilities
            probs = logits.softmax(dim=-1)

            # Sample from the distribution
            next_token = torch.multinomial(probs, num_samples=1)

            # Append to sequence
            idx = torch.cat([idx, next_token], dim=1)

    # Decode and return
    generated = tokenizer.decode(idx[0].tolist())
    return generated

# Load model
model = GPT(vocab_size, n_embed, n_head, n_layer, block_size)
model.load_state_dict(torch.load('model.pt'))
model.to(device)

# Generate text
prompt = "Once upon a time"
generated_text = generate(
    model,
    tokenizer,
    prompt,
    max_new_tokens=300,
    temperature=0.8
)
print(generated_text)
Enter fullscreen mode Exit fullscreen mode

Advanced Sampling: Top-K and Top-P

def generate_with_sampling(model, tokenizer, start_text, max_new_tokens, 
                          temperature=0.8, top_k=None, top_p=None):
    """Generate with Top-K or Nucleus (Top-P) sampling"""
    model.eval()

    idx = torch.tensor([tokenizer.encode(start_text)], dtype=torch.long)
    idx = idx.to(next(model.parameters()).device)

    for _ in range(max_new_tokens):
        with torch.no_grad():
            logits = model(idx)
            logits = logits[:, -1, :] / temperature

            # Top-K filtering
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')

            # Top-P (nucleus) filtering
            if top_p is not None:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

                # Remove tokens with cumulative probability above the threshold
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = 0

                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                logits[indices_to_remove] = float('-inf')

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_token], dim=1)

    return tokenizer.decode(idx[0].tolist())
Enter fullscreen mode Exit fullscreen mode

Configuration for a ~10M Parameter Model

A compact GPT configuration that lands near 10 million parameters:

Hyperparameter Value
Vocabulary Character-level (~80-200 chars)
Context Length (block_size) 256
Embedding Dimension (n_embed) 384
Transformer Layers (n_layer) 8
Attention Heads (n_head) 6
Feed-Forward Hidden Size 1536 (4× embedding)
Batch Size 64
Learning Rate 3e-4
Optimizer AdamW
Parameters ~10 Million

Performance Benchmarks

With the MPS backend enabled, training a 10M parameter model on a MacBook is surprisingly fast:

  • M1 Max: ~7.7x faster than CPU training
  • M2 Pro: A 10M model shows promising results in ~1 hour
  • M3 Pro: Can train the same model in under an hour
# Move model to MPS
device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
model = model.to(device)

# Everything else remains the same!
Enter fullscreen mode Exit fullscreen mode

Where to Go Next

Once you've built a working GPT from scratch, you'll have the foundation to explore more advanced ideas:

Tokenization Improvements

  • Byte Pair Encoding (BPE): More efficient than character-level
  • SentencePiece: Handles unicode and multiple languages
  • Byte-level tokenization: Preserves raw bytes

Architecture Enhancements

  • Rotary Positional Embeddings (RoPE): Better relative positioning
  • FlashAttention: Faster attention computation
  • KV Caching: Faster inference
  • RMSNorm: Faster alternative to LayerNorm
  • Weight Tying: Share embeddings between input and output

Training Optimizations

  • Mixed Precision Training (FP16/BF16)
  • Gradient Checkpointing
  • Learning Rate Schedules: Warmup + cosine decay
  • Gradient Accumulation

Advanced Techniques

  • LoRA and QLoRA: Efficient fine-tuning
  • RLHF: Reinforcement Learning from Human Feedback
  • RAG: Retrieval-Augmented Generation
  • Mixture of Experts (MoE)

Suggested Learning Roadmap

  1. Build the character tokenizer ✓
  2. Implement batch creation and next-token prediction ✓
  3. Add token and positional embeddings ✓
  4. Implement masked self-attention ✓
  5. Extend to multi-head attention ✓
  6. Add feed-forward network, layer norm, residual connections ✓
  7. Stack multiple transformer blocks ✓
  8. Train using cross-entropy loss and AdamW ✓
  9. Implement autoregressive generation with sampling ✓
  10. Next Steps:
    • Add dropout for regularization
    • Implement learning rate scheduling
    • Add validation metrics and early stopping
    • Experiment with different model sizes
    • Try BPE tokenization

Example Training Data Sources

For your first model, start with a single text corpus:

  • Public domain books: Project Gutenberg
  • Shakespeare plays: Complete works
  • Wikipedia extracts: First few MB of text
  • TinyStories: Small story dataset
  • Code datasets: For building a code model

A few hundred kilobytes of clean text is enough to produce recognizable language patterns with a 10M parameter model.


Final Thoughts

The gap between using Generative AI and truly understanding it is smaller than it appears. Every modern language model—from compact open-source models to frontier-scale systems—relies on the same core ideas: tokenization, embeddings, self-attention, feed-forward networks, normalization, residual connections, and next-token prediction.

Building a small GPT model from scratch won't produce ChatGPT-level intelligence, but it will give you something even more valuable: intuition. You'll understand why Transformers work, how they learn, and where the opportunities lie for optimization, fine-tuning, and innovation.

Top comments (0)