How to Build a Transformer Model from Scratch in Python: A Complete Implementation Guide

You can build a fully functional Transformer model from scratch in Python by implementing modular PyTorch components including token embeddings, sinusoidal positional encodings, multi-head self-attention, and feed-forward networks, as demonstrated in the rohitg00/ai-engineering-from-scratch repository.

The rohitg00/ai-engineering-from-scratch repository provides a pedagogical implementation of the Transformer architecture using only PyTorch and standard Python libraries. This educational codebase exposes the mathematical foundations and engineering decisions behind modern large language models by implementing each component explicitly without high-level abstractions. By following the modular design patterns in this repository, you can construct an end-to-end Transformer language model capable of training on custom datasets.

Token and Positional Embeddings

The input layer begins with two essential components that convert discrete token indices into dense vector representations.

Token Embeddings

The TokenEmbedding class wraps torch.nn.Embedding to convert integer token IDs into dense vectors of dimension d_model. This implementation includes custom weight initialization and serves as the initial lookup table for the vocabulary.

According to the source code in phases/19-capstone-projects/33-multihead-self-attention/code/main.py, the token embedding is defined at lines 101-108:

# Conceptual implementation based on source

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
    
    def forward(self, x):
        return self.embedding(x)

Sinusoidal Positional Encoding

The SinusoidalPositionalEmbedding implements the classic sinusoidal position encoding without learned parameters. It pre-computes a full position table using sine and cosine functions of different frequencies, then slices it according to the sequence length.

This implementation appears at lines 14-33 in phases/19-capstone-projects/33-multihead-self-attention/code/main.py. The combined input representation is computed as: emb = token_emb(ids) + pos_emb(seq_len).

Multi-Head Self-Attention Mechanism

The core of the Transformer is the MultiHeadSelfAttention module, which allows the model to focus on different parts of the input sequence simultaneously.

Query, Key, Value Projections

The attention mechanism projects the input tensor into three separate matrices: Queries, Keys, and Values. The implementation in phases/19-capstone-projects/33-multihead-self-attention/code/main.py (lines 20-98) handles this by:

1. Projecting the input to a concatenated QKV matrix with 3*d_model columns
2. Splitting the resulting tensor into separate Q, K, and V tensors
3. Reshaping to separate n_heads independent attention heads

Causal Masking Implementation

To prevent the model from attending to future tokens during training, the implementation applies a causal mask using torch.tril. This creates a lower-triangular matrix ensuring each token can only attend to itself and previous positions.

The scaled dot-product attention computes: attention_scores = Q @ K.transpose() / sqrt(d_head), applies the causal mask, then softmax and dropout. Finally, the heads are merged back together and projected to the output dimension.

Feed-Forward Networks and Normalization

Each Transformer block contains a position-wise feed-forward network and normalization layers to stabilize training.

Position-wise MLP

The FeedForward module expands the hidden dimension by an mlp_expansion factor (typically 4x), applies the GELU activation function, then projects back to d_model. This implementation is found at lines 15-30 in phases/19-capstone-projects/34-transformer-block/code/main.py.

Layer Normalization and Residual Connections

The LayerNorm class provides two variants: pre-LN (normalization before each sub-layer) and post-LN (normalization after the residual addition). The choice between these architectures affects gradient flow during training.

The TransformerBlock class wires these components together, respecting the pre_ln flag to determine the order of operations. This implementation appears at lines 52-60 in phases/19-capstone-projects/34-transformer-block/code/main.py.

Stacking Transformer Blocks

To create a deep Transformer, individual blocks are stacked using the BlockStack class. This module manages an embedding layer, a list of TransformerBlock instances, and a final layer normalization.

The implementation in phases/19-capstone-projects/34-transformer-block/code/main.py (lines 62-77) demonstrates building a six-layer stack and includes a forward/backward pass example to illustrate gradient flow differences between pre-LN and post-LN configurations.

Training the Complete Model

The TinyAttentionLM class combines all components into a complete language model: token embeddings, positional encoding, a Transformer block, and a language modeling head (linear projection to vocabulary size).

The training loop, implemented in phases/19-capstone-projects/33-multihead-self-attention/code/main.py (lines 45-84), demonstrates:

• Generating a synthetic "repeat" task dataset
• Computing cross-entropy loss
• Optimizing with the Adam optimizer
• Visualizing per-head attention heatmaps

The main() function at lines 44-99 validates the causal mask behavior (ensuring future positions receive zero weight) and confirms loss reduction during training.

Practical Code Examples

Instantiating a Tiny Transformer Language Model

from phases.19_capstone_projects.33_multihead_self_attention.code.main import TinyAttentionLM, DemoConfig

cfg = DemoConfig()  # default hyper-parameters

model = TinyAttentionLM(
    vocab_size=cfg.vocab_size,
    d_model=cfg.d_model,
    n_heads=cfg.n_heads,
    max_context_length=cfg.seq_len,
)

print(model)  # displays the full module tree

Running Forward Pass and Inspecting Attention Weights

import torch
from phases.19_capstone_projects.33_multihead_self_attention.code.main import TinyAttentionLM, DemoConfig

cfg = DemoConfig()
model = TinyAttentionLM(
    vocab_size=cfg.vocab_size,
    d_model=cfg.d_model,
    n_heads=cfg.n_heads,
    max_context_length=cfg.seq_len,
)

# Dummy input: batch of 2 sequences, each 12 tokens

ids = torch.randint(0, cfg.vocab_size, (2, cfg.seq_len))
logits, attn_weights = model(ids, return_weights=True)

print(f"Logits shape: {logits.shape}")           # (2, 12, vocab_size)

print(f"Attention shape: {attn_weights.shape}")  # (2, n_heads, seq_len, seq_len)

Training on the Repeat Task

import torch
from phases.19_capstone_projects.33_multihead_self_attention.code.main import TinyAttentionLM, DemoConfig, _train

cfg = DemoConfig()
model = TinyAttentionLM(
    vocab_size=cfg.vocab_size,
    d_model=cfg.d_model,
    n_heads=cfg.n_heads,
    max_context_length=cfg.seq_len,
)

loss_curve = _train(model, cfg)
print(f"Final loss: {loss_curve[-1]:.4f}")

Building a Full Pre-LN Transformer Stack

from phases.19_capstone_projects.34_transformer_block.code.main import BlockConfig, BlockStack
import torch

cfg = BlockConfig(
    d_model=192,
    num_heads=6,
    context_length=64,
    mlp_expansion=4,
    attn_dropout=0.0,
    residual_dropout=0.0,
    pre_ln=True,
)

stack = BlockStack(cfg, depth=6)
tokens = torch.randint(0, 128, (2, 32))  # (batch, seq_len)

output = stack(tokens)
print(f"Output shape: {output.shape}")  # (2, 32, 192)

Summary

• Modular architecture: The rohitg00/ai-engineering-from-scratch repository implements Transformers as discrete PyTorch modules including TokenEmbedding, SinusoidalPositionalEmbedding, MultiHeadSelfAttention, and FeedForward.

• Attention mechanism: Multi-head self-attention splits queries, keys, and values across parallel heads and applies a causal torch.tril mask to prevent looking ahead during training.

• Normalization strategies: The code supports both pre-LN (normalization before sub-layers) and post-LN (normalization after residuals) configurations, with BlockStack handling deep network composition.

• Training pipeline: The TinyAttentionLM class provides a complete training example using Adam optimization and cross-entropy loss on synthetic sequential data.

Frequently Asked Questions

How does the causal mask work in the multi-head attention implementation?

The causal mask uses torch.tril to create a lower-triangular matrix that is applied to the attention scores. This ensures that when predicting the next token, the model can only attend to the current token and previous tokens, with future positions receiving zero weight. The implementation in phases/19-capstone-projects/33-multihead-self-attention/code/main.py validates this behavior by verifying that attention weights for future positions are exactly zero.

What is the difference between pre-LN and post-LN Transformer blocks?

Pre-LN applies layer normalization before the attention and feed-forward sub-layers, while post-LN applies normalization after the residual connections. According to the source code in phases/19-capstone-projects/34-transformer-block/code/main.py, pre-LN typically provides more stable gradients during training, which is why modern architectures often prefer this configuration over the original post-LN design.

Can I scale this implementation to train on large datasets?

Yes, the modular design allows scaling by adjusting the BlockConfig parameters such as d_model, num_heads, and depth. The BlockStack class in phases/19-capstone-projects/34-transformer-block/code/main.py supports arbitrary depth configurations, though for production-scale training you would need to add data loading, distributed training loops, and gradient accumulation beyond the educational examples provided.

Why use sinusoidal positional encoding instead of learned embeddings?

The SinusoidalPositionalEmbedding uses fixed sine and cosine functions of different frequencies to encode position information without adding trainable parameters. This approach, implemented in phases/19-capstone-projects/33-multihead-self-attention/code/main.py, generalizes to sequence lengths not seen during training and provides unique encoding for each position, though modern architectures often use learned rotary or absolute positional embeddings instead.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}