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

> Learn to build a Transformer model from scratch in Python. Implement token embeddings, self-attention, and more with this complete guide and code from rohitg00/ai-engineering-from-scratch. Start coding today.

- Repository: [Rohit Ghumare/ai-engineering-from-scratch](https://github.com/rohitg00/ai-engineering-from-scratch)
- Tags: how-to-guide
- Published: 2026-06-15

---

**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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/19-capstone-projects/33-multihead-self-attention/code/main.py), the token embedding is defined at lines 101-108:

```python

# 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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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

```python
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

```python
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

```python
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

```python
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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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.