TL;DR — Training a GPT‑style model from scratch reveals the inner workings of attention, optimization, and data pipelines. You’ll end up with a runnable loop that trains a tiny transformer on text, and the code can be extended to real workloads. It’s a concrete, interview‑ready project that proves you can ship AI infrastructure.

Building a minimal GPT‑style training loop is one of the most effective ways to move from “I used a library” to “I understand how the pieces fit together.” In this post we’ll walk through a complete, runnable PyTorch implementation that you can clone, train on a few sentences, and then iterate on. Along the way you’ll see how data, model, optimizer, and loss interlock, and you’ll walk away with a portfolio piece that hiring managers can inspect and run.

Why This Project Stands Out on a CV

  • Deep‑learning fundamentals – You implement forward/backward passes, loss computation, and optimizer steps yourself, proving you understand the math behind modern LLMs.
  • Systems thinking – Choosing a character‑level tokenizer, managing memory with torch.no_grad(), and writing checkpoint I/O shows you can design pipelines that scale.
  • Production‑ready habits – Adding logging, versioned saves, and a simple requirements.txt mirrors the workflow of AI engineers at companies like OpenAI, DeepMind, and GCP‑based ML teams.
  • Signal for hiring managers – Roles such as ML Engineer, AI Infrastructure Engineer, and Research Engineer value candidates who can move from idea to runnable code without relying on high‑level abstractions.
  • Extensibility demo – The modular structure (tokenizer → model → loop) makes it trivial to swap in GPT‑2‑style BPE, mixed‑precision, or distributed training later.

Architecture Overview

The loop consists of four core components that map directly to a real LLM pipeline:

[Tokenizer] → [Dataset] → [Model (N transformer blocks)] → [Loss + Optimizer] → [Checkpoint / Log]
  • Tokenizer – Character‑level mapping (stoi/itos) so no external dependency is required.
  • Dataset – Simple torch.utils.data.Dataset that yields fixed‑length sequences of token IDs.
  • Model – Stack of n_layer transformer blocks, each containing a multi‑head self‑attention sub‑layer and a position‑wise feed‑forward network, followed by RMSNorm and a final lm_head linear layer.
  • Training loop – forward(x) → cross‑entropy loss → optimizer.step() → zero_grad().
  • Checkpointing – torch.save(model.state_dict(), "ckpt.pt") after each epoch, enabling resumability.

A text diagram of the data flow:

   raw text
        |
   char‑level tokenizer
        +--> token IDs
        |
   Dataset (batchify) --> Tensor [B, T]
        |
   Transformer model (N blocks)
        |
   logits -> cross‑entropy loss
        |
   optimizer -> weight update
        |
   saved checkpoint

Building It Step By Step

Below are seven concrete steps, each with a runnable Python snippet (all code is Python‑tagged). You can copy‑paste these into a file train_minigpt.py and execute python train_minigpt.py.

Step 1 – Environment & Imports

# step_1_env.py
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import math, random, time
print("PyTorch version:", torch.__version__)

Step 2 – Character‑level Tokenizer

# step_2_tokenizer.py
class CharTokenizer:
    """Very small tokenizer that maps characters → ids and back."""
    def __init__(self, text):
        chars = sorted(set(text))
        self.stoi = {ch: i+1 for i, ch in enumerate(chars)}  # 0 reserved for padding
        self.itos = {i: ch for ch, i in self.stoi.items()}
        self.vocab_size = len(chars) + 1

    def encode(self, s):
        return [self.stoi[c] for c in s]

    def decode(self, ids):
        return "".join(self.itos[i] for i in ids if i in self.itos)

tok = CharTokenizer("hello world this is a minimal gpt experiment")
print("vocab size:", tok.vocab_size)

Step 3 – Toy Dataset

# step_3_dataset.py
class GPTDataset(Dataset):
    def __init__(self, text, tokenizer, block_size=8):
        self.ids = torch.tensor(tokenizer.encode(text), dtype=torch.long)
        self.block_size = block_size

    def __len__(self):
        return max(0, len(self.ids) - self.block_size)

    def __getitem__( idx):
        # input = ids[idx : idx+block_size], target = ids[idx+1 : idx+block_size+1]
        x = self.ids[idx: idx+self.block_size]
        y = self.ids[idx+1: idx+self.block_size+1]
        return x, y

text = "hello world! this is a tiny gpt training loop."
dataset = GPTDataset(text, tok, block_size=8)
dloader = DataLoader(dataset, batch_size=2, shuffle=True)
for batch in dloader:
    print("batch x:", batch[0])
    print("batch y:", batch[1])
    break

Step 4 – Transformer Block

# step_4_block.py
class TransformerBlock(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout=0.1):
        super().__init__()
        # multi‑head self‑attention
        self.attn = nn.MultiheadAttention(n_embd, n_head, dropout=dropout, batch_first=True)
        # feed‑forward
        self.ff = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
        )
        # normalization
        self.ln_1 = nn.RMSNorm(n_embd)
        self.ln_2 = nn.RMSNorm(n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # residual path
        attn_out, _ = self.attn(self.ln_1(x), self.ln_1(x), self.ln_1(x))
        x = x + self.dropout(attn_out)
        ff_out = self.ff(self.ln_2(x))
        x = x + self.dropout(ff_out)
        return x

Step 5 – Tiny GPT Model

# step_5_model.py
class MiniGPT(nn.Module):
    def __init__(self, vocab_size, n_embd=32, n_head=4, n_layer=2, block_size=8):
        super().__init__()
        self.block_size = block_size
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        self.position_embedding = nn.Embedding(block_size, n_embd)
        self.blocks = nn.ModuleList([
            TransformerBlock(n_embd, n_head, block_size) for _ in range(n_layer)
        ])
        self.ln_f = nn.RMSNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)

    def forward(self, idx):
        # idx: [B, T] token ids
        B, T = idx.shape
        # truncate if longer than block_size
        idx = idx[:, -self.block_size:]
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))
        x = tok_emb + pos_emb
        for blk in self.blocks:
            x = blk(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)  # [B, T, vocab_size]
        return logits

Step 6 – Training Step & Loop

# step_6_train.py
@torch.no_grad()
def estimate_loss(model, dataset, tokenizer, eval_iters=20):
    out = {}
    model.eval()
    for split in ["train"]:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            x, y = dataset[random.randint(0, len(dataset)-1)]
            x, y = x.unsqueeze(0), y.unsqueeze(0)  # add batch dim
            logits = model(x)
            loss = nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

model = MiniGPT(vocab_size=tok.vocab_size, n_embd=64, n_head=4, n_layer=2, block_size=8)
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
epochs = 5
for epoch in range(epochs):
    t0 = time.time()
    for x, y in dloader:
        optimizer.zero_grad()
        logits = model(x)
        loss = nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
        loss.backward()
        optimizer.step()
    elapsed = time.time() - t0
    print(f"epoch {epoch+1}/{epochs} | loss {loss.item():.4f} | {elapsed:.1f}s")

Step 7 – Checkpoint & Generation

# step_7_checkpoint.py
# after training, save model state
torch.save(model.state_dict(), "minigpt.pt")
# later, load and generate text
model.load_state_dict(torch.load("minigpt.pt"))
model.eval()
context = torch.tensor([tok.stoi[c] for c in "hello"], dtype=torch.long).unsqueeze(0)
generated = context
for _ in range(30):
    logits = model(generated)
    probs = torch.softmax(logits[:, -1, :], dim=-1)
    idx = torch.multinomial(probs, num_samples=1).item()
    generated = torch.cat([generated, torch.tensor([[idx]], dtype=torch.long)], dim=1)
print("generated text:", tok.decode(generated.squeeze().tolist()))

Running and Testing It

  1. Install dependencies
    pip install torch==2.3.0  # or the latest stable
  2. Save the snippets above into train_minigpt.py (or separate files as shown) and run:
    python train_minigpt.py
  3. Expected output – You’ll see a decreasing cross‑entropy loss each epoch, e.g.
    epoch 1/5 | loss 3.4217 | 12.3s
epoch 2/5 | loss 2.9845 | 11.7s
epoch 3/5 | loss 2.6172 | 11.2s
…
  4. Verify generation – At the end of the run the script prints a short continuation of the prompt (e.g. “hello world! this is a…”). If the loss is falling and the generated text makes semantic sense (even if trivial), the loop is working.
  5. Debug tips –
    • If you hit CUDA out of memory, reduce n_embd or batch_size.
    • If loss stalls, increase the learning rate or add a learning‑rate scheduler (torch.optim.lr_scheduler.StepLR).
    • Use torch.autograd.set_detect_anomaly(True) for quick gradient‑check diagnostics.

Extending It: Your Roadmap to Senior‑Level

  1. Mixed‑precision training (torch.cuda.amp) – Cuts memory usage and speeds up convergence on modern GPUs, a common requirement in production training pipelines.
  2. Checkpoint resumption & early stopping – Persist optimizer state and epoch number; enables long‑run experiments without losing progress.
  3. Distributed data‑parallel training (torch.distributed.launch or torchrun) – Scales the loop across multiple GPUs, mirroring how teams train GPT‑2‑size models on clusters.
  4. TensorBoard / MLflow logging – Track loss, learning‑rate, and gradient norms over time; essential for reproducibility and performance reviews.
  5. **BPE tokenizer (e.g., using tiktoken or HuggingFace Tokenizer) – Moves from character‑level to word‑piece, dramatically reducing sequence length and improving model efficiency.
  6. Evaluation metric: perplexity on a held‑out validation set – Provides a standardized way to compare different model sizes or training regimes, a skill hiring managers look for in senior ML engineers.

Key Takeaways

  • Implementing a training loop from scratch demystifies attention, loss back‑prop, and optimizer mechanics.
  • A modular design (tokenizer → dataset → model → loop) makes it trivial to swap components—BPE, mixed precision, or distributed training.
  • Checkpointing and logging turn a toy script into a reproducible experiment you can discuss in interviews.
  • The project signals deep‑learning fundamentals, systems thinking, and production‑ready habits—exactly the traits hiring managers seek in AI/ML engineers.
  • Extending the loop with any of the six upgrades above transforms it from a “hello‑world” demo into a building block for real‑world LLM workloads.

Further Reading