I Built GPT-2 from Scratch in PyTorch

GPT : Generative Pre-trained Transformer is a decoder-only transformer architecture which takes the tokens as input and predicts the next token. In this tutorial we will build GPT-2 in a single file : architecture, training loop, checkpointing, inference.

Full Code : Github Gist

Image Source : https://cameronrwolfe.substack.com/p/decoder-only-transformers-the-workhorse

Dataset : We will be working on english dataset with paragraphs , with each paragraph is separated by \n\n spaces . We split the data 90/10 for training and validation sets. Add a special <|endoftext|> token to separate paragraphs , inspired by GPT-2 architecture.

    # reading file 
    with open('data/english_big.txt', 'r', encoding='utf-8') as f:
        lines = f.read()

    # split on paragraph boundaries
    paragraphs = lines.split('\n\n')

    print(f'\n Total paragraphs : {len(paragraphs)} \n')

    random.seed(42)
    random.shuffle(paragraphs)

    split_idx = int(0.9*len(paragraphs))

    # join paragraphs with EOT token

    train_text = "<|endoftext|>".join(paragraphs[:split_idx])
    val_text = "<|endoftext|>".join(paragraphs[split_idx:])

Tokenization : Here we are using OpenAI tiktoken It’s wrapped in tiktokenizer class to convert tokens into torch.tensors directly

import tiktoken

class tiktokenizer:

    def __init__(self, model):
        self.enc = tiktoken.get_encoding(model)

    def encode_t(self, text, allowed_special=set()):
        tokens = self.enc.encode(text, allowed_special=allowed_special)
        tokens = torch.tensor(tokens, dtype=torch.long)
        return tokens # T

    def decode_t(self, tokens):
        out = self.enc.decode(tokens)
        return out

Model Architecture : The GPT model consists of *Token Embeddings : Positional Embeddings : Nx Decoder-only Transformer Blocks : Layer Normalization : Output Linear Layer. *In GPT-2 Model the weights of starting embedding layer and final output liner layers are shared.

class GPT(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.config = config  

        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        # weight sharing between embedding layer and logits layer - better result + reduce parameters
        self.lm_head.weight = self.transformer.wte.weight

        # init parameters
        self.apply(self._init_weights)

The Decoder Block consists of *LayerNormalization : Causal Self Attention : Residual Connections : LayerNormalization : MLP layer : Residual Connections*

class Block(nn.Module):
    """
    Decoder Only Transformer Block (GPT-Style)
    |
    ├── LayerNorm → Self-Attention → Residual → LayerNorm → MLP → Residual
    |
    ├── Input:  (batch, seq, n_embd)
    └── Output: (batch, seq, n_embd)
    """

    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd) # In GPT-2 layer-normalization blocks are before attention 
        self.attn = CausalSelfAttention(config) # multi-head self-attention
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        # x = residual path + (layer_norm -> attention)
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

MLP : The MLP / Linear Layer used in decoder block consists of first Linear Layer (which projects the embeddings into higher dimesional embeddings : 4*emb_dim . Then GELU activation is used then second Linear Layer (which projects those higher dimensional embeddings into their original embeddings) emb_dim.

class MLP(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4*config.n_embd)
        self.gelu = nn.GELU(approximate='tanh')                 # suggest to use approximate='none'
        self.c_proj = nn.Linear(4*config.n_embd, config.n_embd)

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        return x

Causal Self Attention : This is the most important part of the decoder-only transformer block. It allows each token to attend all previous tokens masking future tokens so each token see only past context and not the future tokens. It takes the token embeddings, using Wq,Wk,Wv matrices creates Q, K, V vectors. Applies multi-head attention, finally projects back to the original dimension.

Attention(Q, K, V) = softmax((Q.Kᵀ / √dk) + M) V . *here M : mask*

class CausalSelfAttention(nn.Module):
    """
    Causal (masked) Multi-Head Self-Attention
    Input:
        x: (B, T, C) 
            B = batch size
            T = sequence length
            C = embedding dimension (n_embd)
    Output:
        y: (B, T, C)

    x (B, T, C)
    |
    Linear → QKV (B, T, 3C)
    |
    split → q, k, v
    |
    reshape → (B, nh, T, hs)
    |
    attention → (B, nh, T, hs)
    |
    transpose → (B, T, nh, hs)
    |
    reshape → (B, T, C)
    |
    c_proj → (B, T, C)
    """

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # key, query, value projections for all heads,but in batch
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd) # [Wq | Wk | Wv] concatenated 
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd) # after multi-head attention we merge heads → (B, T, C)
        # regularization
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        # not really a 'bias' more of a mask, but following the OpenAI/HF naming though
        self.resid_dropout = nn.Dropout(0.2)

    def forward(self, x):
        B, T, C = x.size() # batch_size, sequence_lenght, n_embd
        # calculate query, key, values for all heads in batch and move head forward to be the bat.............
        # nh = 'number of heads', hs = 'head size'= C / n_head, C = (number of channels) == nh * ns
        # in GPT-2 (124M), n_head = 12, hs = 64, so nh*hs = 768 channels in the Transformer
        qkv = self.c_attn(x)    # (B, T, 3C)
        q, k, v = qkv.split(self.n_embd, dim=2)     # 3 * (B, T, C)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        # ============= Implementing Flash Attention ===================
        y = F.scaled_dot_product_attention( q, k, v, is_causal=True, dropout_p=0.2 if self.training else 0.0 )
        # ==============================================================
        # now re-assembling all head outputs side by side
        y = y.transpose(1, 2) # (B, T, nh, hs)
        y = y.contiguous().view(B, T, C) # (B, T, C)
        # output projection - after concatenation, c_proj mixes information from all heads together
        y = self.c_proj(y)
        # ====================================================================
        y = self.resid_dropout(y)
        # ====================================================================
        return y

DataLoader : Here custom sequential DataLoader is used for training , which takes :block_size tokens as input and same number of output tokens shifted by 1 token .

class DataLoader:
    """
    Sequential DataLoader : GPT style 
    """

    def __init__(self, tokens, batch_size, block_size):

        self.tokens = tokens
        self.batch_size = batch_size
        self.block_size = block_size
        self.pos = 0

    def next_batch(self):
        B, T = self.batch_size, self.block_size
        # grab chunk from current position
        chunk = self.tokens[self.pos : self.pos + B*T + 1]
        x = chunk[:-1].view(B,T)
        y = chunk[1:].view(B,T)
        # advance position
        self.pos += B*T
        # reset if we'd go out of bounds next call
        if self.pos + B*T + 1 > len(self.tokens):
            self.pos = 0
        return x, y

Training Loop : Now comes training loop. Here we run training loop for epochs times and in each epoch we load the data sequentially using custom dataloader in batches for steps_per_epoch = len(train_dataloader.tokens) // (B*T) times. during each epoch runs training and evaluates validation loss.

def train(model:GPT, optimizer,  train_dataloader, val_dataloader,  epochs:int):

    B, T = train_dataloader.batch_size, train_dataloader.block_size
    steps_per_epoch = len(train_dataloader.tokens) // (B*T)
    val_steps = len(val_dataloader.tokens) // (B*T)

    if val_steps == 0:
        print("Warning: val set too small for even one batch, skipping val")
        val_steps = 1

    print(f"Total train steps per epoch : {steps_per_epoch}")
    print(f"Total val steps per epoch    : {val_steps}\n")

    for epoch in range(epochs):

        # ────────── Training ────────────────────
        model.train()
        train_dataloader.pos = 0
        epoch_loss = 0.0 

        for step in range(steps_per_epoch):

            x, y = train_dataloader.next_batch()
            x, y  = x.to(device), y.to(device) # (B, T)

            logits = model(x)
            _, _, vocab_size = logits.shape

            loss = F.cross_entropy(logits.view(B*T, vocab_size), y.view(B*T))

            epoch_loss += loss.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # ────────── Validation ────────────────────
        model.eval()
        val_dataloader.pos=0
        val_epoch_loss = 0.0

        with torch.no_grad():

            for step in range(val_steps):

                x_val, y_val = val_dataloader.next_batch()
                x_val, y_val = x_val.to(device), y_val.to(device)

                val_logits = model(x_val)
                _, _, val_vocab_size = val_logits.shape
                val_loss = F.cross_entropy(val_logits.view(B*T, val_vocab_size), y_val.view(B*T))

                val_epoch_loss += val_loss.item()

        epoch_loss = epoch_loss / steps_per_epoch
        val_epoch_loss = val_epoch_loss / val_steps
        print(f'{epoch}/{epochs} | train loss : {epoch_loss:.5f} | val loss : {val_epoch_loss:.5f}')

Checkpointing : After training model is saved and loaded for future inference

def save_checkpoint(model, optimizer, path):
    torch.save({
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }, path)
    print(f"Checkpoint Saved \n")

def load_checkpoint(model, optimizer, path):
    checkpoint = torch.load(path, map_location=device, weights_only=True)
    model.load_state_dict(checkpoint['model_state_dict'])
    if optimizer is not None:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    print(f"Model Loaded \n")
    return model, optimizer

Training Implementation : The architecture details used for GPT :

from dataclasses import dataclass
import torch
import torch.nn as nn
from torch.nn import functional as F
import math
import os
import random

import tiktoken

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

EOT_TOKEN = 50256

# GPT-2 124M parameter model configs
@dataclass
class GPTConfig:
    block_size: int = 1024
    vocab_size: int = 50257
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768

def main():

    block_size = 64
    batch_size  = 32
    epochs = 50
    learning_rate = 3e-4
    num_return_sequences = 15
    max_length = 80
    max_lr = 6e-4
    min_lr = max_lr * 0.1
    warmup_steps = 100

    MODEL_PATH = 'models/gpt-train.pt'

    config = GPTConfig()
    model = GPT(config)
    model.to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, betas=(0.9, 0.95), eps=1e-8)

    if os.path.exists(MODEL_PATH):
        model, optimizer = load_checkpoint(model, optimizer, MODEL_PATH)

    enc = tiktokenizer('gpt2')

    train_tokens = enc.encode_t(train_text, allowed_special={'<|endoftext|>'}) # (T)
    val_tokens = enc.encode_t(val_text, allowed_special={'<|endoftext|>'}) # (T)

    print(f"Train tokens : {len(train_tokens)}")
    print(f"Val tokens : {len(val_tokens)}")

    train_dataloader = DataLoader(train_tokens, batch_size, block_size)
    val_dataloader = DataLoader(val_tokens, batch_size, block_size)

    train(model, optimizer, train_dataloader, val_dataloader, epochs, max_lr, min_lr, warmup_steps)

    # save after training
    save_checkpoint(model, optimizer, MODEL_PATH)

Inference : Loades the trained model and generates text for given prompt :

project/
├── gpt.py          ← model, training, all classes
└── inference.py    ← this file

from dataclasses import dataclass
import torch
import torch.nn as nn
from torch.nn import functional as F
import math
import os
import random
import tiktoken

# gpt.py must be in the same directory as this file
from gpt import GPTConfig, GPT, tiktokenizer
from gpt import load_checkpoint,  generate

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_PATH = "models/gpt-train.pt"
EOT_TOKEN = 50256

# -----------------------------------------------------------------------------

# model config
config = GPTConfig()

# create model
model = GPT(config)
model.to(device)

# load checkpoints
model, _ = load_checkpoint(model, optimizer=None, path=MODEL_PATH)

# inference mode
model.eval()
# tokenizer

enc = tiktokenizer("gpt2")
# -----------------------------------------------------------------------------

prompt = "But why Mr. Darcy came so often "

tokens = enc.encode_t(prompt)

tokens = tokens.unsqueeze(0).repeat(5, 1)

x = tokens.to(device)

generate(model=model, num_return_sequences=5, x=x, max_length=200, enc=enc, temperature=0.9)

If you find any mistake or issue in the blog, please do comment. It’ll help me and other learners also. If you find this blog helpful then do share it with your peers 😊.

Full Code on Github Gist .

Based on Andrej Karpathy’s Let’s reproduce GPT-2. Modifications include a custom DataLoader, tiktoken wrapper, checkpointing, and EOT-based inference.