TL;DR — nanoGPT is the single best deep-learning side project for a CV: it’s small enough to finish in a weekend, but it touches every real-world ML systems concern (mixed-precision via torch.amp, gradient accumulation for effective batch sizing, cosine LR scheduling, weight tying, efficient data loading). We’ll build it step-by-step in pure PyTorch with runnable code, then map the upgrade path to senior-level distributed training.

When I review engineering portfolios, I look for one signal: can the candidate ship a non-trivial system end-to-end? A from-scratch nanoGPT is unusual because it’s both short (the entire training loop fits in one file) and complete — it forces you to make decisions about precision, memory, learning rate policy, and data pipelines that production ML engineers make every day. Andrej Karpathy’s nanoGPT is the canonical reference, and the patterns below are directly drawn from how teams train LLMs on a single node or across a cluster.

This isn’t a tutorial that prints “TODO: implement attention.” Every code block below is runnable.

Why This Project Stands Out on a CV

Hiring managers for ML, MLOps, and applied-AI roles scan portfolios for evidence of three things: systems intuition, correctness under scale, and the ability to read primary sources. A from-scratch nanoGPT hits all three:

  • Systems intuition. Mixed-precision training with torch.amp.autocast and GradScaler (see the PyTorch AMP docs) isn’t an academic exercise — it’s how every modern transformer is trained. Implementing it yourself shows you understand FP16/BF16 numerics, loss scaling, and gradient underflow, not just how to call model.half().
  • Correctness under scale. Gradient accumulation simulates large effective batch sizes on limited GPU memory. The pattern — forward, loss.backward() per micro-batch, optimizer.step() only every N steps — is the same one Horovod and PyTorch DDP wrap around with collective communication. Getting it right in single-process form is the prerequisite to understanding DistributedDataParallel.
  • Reading primary sources. The transformer block is taken verbatim from Vaswani et al. 2017. The cosine schedule with warmup is from Loshchilov & Hutter 2017. The AdamW betas and weight decay come from Kingma & Ba 2014 and the AdamW paper. Wiring those into working code is what separates a candidate who has read “Attention Is All You Need” from one who has implemented it.

For roles, this project signals: ML Engineer, Applied Research Engineer, LLM Infrastructure Engineer, and Founding Engineer at an AI startup. The upgrade path we’ll add at the end — sharding, fault tolerance, observability — extends the signal to ML Platform and Training Infrastructure roles, where comp bands are noticeably higher.

Architecture Overview

The training loop has six components. Each is small, each is mandatory, and each has a production analog:

  • Data pipeline — tokenized text corpus, packed into fixed-length sequences, served by DataLoader workers. Production analog: a WebDataset or Fsspec-backed reader talking to S3.
  • Model — a nn.Module with Embedding → [TransformerBlock × N] → Linear plus weight tying between the input embedding and output projection (a parameter-saving trick from Press & Wolf 2017). Production analog: Llama, Qwen, and most modern LLMs.
  • Forward pass — multi-head causal self-attention with a causal mask, residual connections, and pre-norm LayerNorm. Production analog: FlashAttention in vLLM and PyTorch’s scaled-dot-product.
  • Mixed precision context — torch.amp.autocast(dtype=torch.bfloat16) wrapping forward and GradScaler (for FP16) or none (for BF16). Production analog: what’s inside DeepSpeed and Megatron.
  • Gradient accumulation buffer — loss is divided by accum_steps before backward; optimizer.step() runs every N microbatches. Production analog: FSDP’s gradient bucketing and pipeline stages.
  • Optimizer + scheduler — AdamW with decoupled weight decay and a cosine learning rate schedule with linear warmup. Production analog: what Hugging Face transformers wraps in get_cosine_schedule_with_warmup.

The dataflow is:

tokens → DataLoader → micro-batch → autocast(model) → scaled loss
   → loss.backward() (accumulated) → GradScaler.update()
   → every accum_steps: optimizer.step() → scheduler.step() → zero_grad()

Building It Step by Step

We assume Python 3.10+, PyTorch 2.2+, and a CUDA-capable GPU with at least 8 GB of VRAM. CPU works for a sanity-check run; BF16 needs a relatively modern GPU (Ampere or newer).

Step 1 — Dependencies and project layout

python -m venv .venv && source .venv/bin/activate
pip install torch numpy tiktoken requests

nanogpt/
├── data/
│   └── prepare.py        # downloads + tokenizes a text corpus
├── model.py              # GPT definition
├── train.py              # training loop with AMP + grad accum + cosine LR
└── config.py             # hyperparameters in one place

Step 2 — Configuration

Keep hyperparameters in one file so reviewers can see your tradeoffs at a glance.

# config.py
from dataclasses import dataclass

@dataclass
class GPTConfig:
    block_size: int = 1024        # context length
    vocab_size: int = 50304       # GPT-2 padded vocab (50257 rounded up)
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0          # 0.0 for eval, 0.1 for small overfit runs
    bias: bool = False            # GPT-2 style: no bias in Linear/LN

@dataclass
class TrainConfig:
    batch_size: int = 12          # micro-batch per step
    accum_steps: int = 8          # effective batch = 12 * 8 = 96 sequences
    max_iters: int = 6000
    lr: float = 6e-4
    min_lr: float = 6e-5           # cosine floor
    warmup_iters: int = 200
    weight_decay: float = 1e-1
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0
    dtype: str = "bfloat16"       # "bfloat16" or "float16"

Step 3 — Data preparation

For a portfolio project, fine-tuning or character-level training on a small public corpus is the right scope. We’ll use Tiny Shakespeare, the same dataset Karpathy’s repo uses, but tokenized with tiktoken’s GPT-2 BPE so we match a real-world tokenizer.

# data/prepare.py
import os, requests, numpy as np, tiktoken

URL = "https://huggingface.co/datasets/karpathy/tiny_shakespeare/resolve/main/tinyshakespeare.txt"
OUT_DIR = "data"
enc = tiktoken.get_encoding("gpt2")

def load_corpus():
    path = os.path.join(OUT_DIR, "input.txt")
    if not os.path.exists(path):
        with open(path, "wb") as f:
            f.write(requests.get(URL, timeout=30).content)
    return open(path, "r", encoding="utf-8").read()

def tokenize():
    text = load_corpus()
    ids = enc.encode_ordinary(text)
    ids = np.array(ids, dtype=np.uint16)
    n = len(ids)
    split = int(0.9 * n)
    train_ids, val_ids = ids[:split], ids[split:]
    train_ids.tofile(os.path.join(OUT_DIR, "train.bin"))
    val_ids.tofile(os.path.join(OUT_DIR, "val.bin"))
    print(f"train tokens: {len(train_ids):,}, val tokens: {len(val_ids):,}")

if __name__ == "__main__":
    os.makedirs(OUT_DIR, exist_ok=True)
    tokenize()

python data/prepare.py

Step 4 — Memory-mapped data loader

Loading the entire token array into RAM works for Tiny Shakespeare but signals the wrong instinct. Memory-mapping with np.memmap is the same pattern production pipelines use when datasets exceed RAM.

# data/loader.py
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader

class ShakespeareDataset(Dataset):
    def __init__(self, path, block_size):
        self.data = np.memmap(path, dtype=np.uint16, mode="r")
        self.block_size = block_size

    def __len__(self):
        return max(1, len(self.data) - self.block_size - 1)

    def __getitem__(self, idx):
        chunk = self.data[idx : idx + self.block_size + 1].astype(np.int64)
        x = torch.from_numpy(chunk[:-1])
        y = torch.from_numpy(chunk[1:])
        return x, y

def build_loaders(block_size, batch_size, num_workers=2):
    train_ds = ShakespeareDataset("data/train.bin", block_size)
    val_ds   = ShakespeareDataset("data/val.bin",   block_size)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
                              drop_last=True, num_workers=num_workers,
                              pin_memory=True, persistent_workers=num_workers > 0)
    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False,
                              drop_last=True, num_workers=num_workers,
                              pin_memory=True, persistent_workers=num_workers > 0)
    return train_loader, val_loader

Step 5 — The model

We use pre-norm transformer blocks, GELU, and weight tying between the input token embedding and the output projection — three decisions that match GPT-2 / Llama and materially affect parameter count and loss curve.

# model.py
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from config import GPTConfig

class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        assert cfg.n_embd % cfg.n_head == 0
        self.n_head = cfg.n_head
        self.n_embd = cfg.n_embd
        self.c_attn = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=cfg.bias)
        self.c_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias)
        self.dropout = cfg.dropout
        # Pre-compute causal mask as a buffer (non-parameter, moves with .to(device))
        self.register_buffer("mask",
            torch.tril(torch.ones(cfg.block_size, cfg.block_size))
                 .view(1, 1, cfg.block_size, cfg.block_size))

    def forward(self, x):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        # Scaled dot-product attention with causal mask
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)

class MLP(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.c_fc    = nn.Linear(cfg.n_embd, 4 * cfg.n_embd, bias=cfg.bias)
        self.gelu    = nn.GELU()
        self.c_proj  = nn.Linear(4 * cfg.n_embd, cfg.n_embd, bias=cfg.bias)
    def forward(self, x):
        return self.c_proj(self.gelu(self.c_fc(x)))

class Block(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.ln_1 = nn.LayerNorm(cfg.n_embd)
        self.attn = CausalSelfAttention(cfg)
        self.ln_2 = nn.LayerNorm(cfg.n_embd)
        self.mlp  = MLP(cfg)
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

class GPT(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.cfg = cfg
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(cfg.vocab_size, cfg.n_embd),
            wpe = nn.Embedding(cfg.block_size, cfg.n_embd),
            drop = nn.Dropout(cfg.dropout),
            h = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)]),
            ln_f = nn.LayerNorm(cfg.n_embd),
        ))
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        # Weight tying: share embedding and output projection parameters
        self.transformer.wte.weight = self.lm_head.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.size()
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
        tok_emb = self.transformer.wte(idx)
        pos_emb = self.transformer.wpe(pos)
        x = self.transformer.drop(tok_emb + pos_emb)
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1,
            )
        return logits, loss

Step 6 — Cosine LR schedule with warmup

This is the schedule from the GPT-3 paper: linear warmup, then cosine decay to a floor, then hold. Get this wrong and your loss curve will look deceptively flat.

# train.py (utility)
import math

def get_lr(it, warmup_iters, max_iters, min_lr, max_lr):
    # 1) linear warmup
    if it < warmup_iters:
        return max_lr * (it + 1) / (warmup_iters + 1)
    # 2) cosine decay to min_lr
    if it > max_iters:
        return min_lr
    decay_ratio = (it - warmup_iters) / (max_iters - warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (max_lr - min_lr)

Step 7 — The training loop (the part reviewers actually read)

This is where mixed-precision, gradient accumulation, gradient clipping, and the LR schedule all come together. Comments annotate each decision so a hiring manager reading your code sees the reasoning.

# train.py (main loop)
import time, torch
from torch.amp import autocast, GradScaler
from torch.nn.parallel import DistributedDataParallel as DDP  # noqa: F401  (used in extension)
from config import GPTConfig, TrainConfig
from model import GPT
from data.loader import build_loaders
from train import get_lr  # if split into files

def main():
    mcfg = GPTConfig()
    tcfg = TrainConfig()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch.manual_seed(1337)

    train_loader, val_loader = build_loaders(mcfg.block_size, tcfg.batch_size)
    model = GPT(mcfg).to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=tcfg.lr,
        betas=(tcfg.beta1, tcfg.beta2),
        weight_decay=tcfg.weight_decay,
    )

    use_fp16 = (tcfg.dtype == "float16")
    autocast_dtype = torch.bfloat16 if tcfg.dtype == "bfloat16" else torch.float16
    # BF16 has the same exponent range as FP32, so no GradScaler needed.
    # FP16 does not, so we MUST use GradScaler to avoid gradient underflow.
    scaler = GradScaler(enabled=use_fp16)

    def run_split(split):
        loader = train_loader if split == "train" else val_loader
        model.train(split == "train")
        losses = []
        for x, y in loader:
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            with autocast(device_type=device.type, dtype=autocast_dtype):
                _, loss = model(x, y)
                loss = loss / tcfg.accum_steps   # key: scale loss for accumulation
            scaler.scale(loss).backward()
            losses.append(loss.item() * tcfg.accum_steps)
        return torch.tensor(losses).mean().item()

    iter_num = 0
    best_val = float("inf")
    log = []
    t0 = time.time()

    while iter_num < tcfg.max_iters:
        # 1) set the learning rate for this iteration
        lr = get_lr(iter_num, tcfg.warmup_iters, tcfg.max_iters,
                    tcfg.min_lr, tcfg.lr)
        for pg in optimizer.param_groups:
            pg["lr"] = lr

        optimizer.zero_grad(set_to_none=True)  # memory-efficient zeroing

        # 2) gradient accumulation loop
        for micro_step in range(tcfg.accum_steps):
            # In a streaming variant, you'd draw a fresh micro-batch each step;
            # here we cycle the DataLoader for simplicity.
            try:
                x, y = next(data_iter)
            except (NameError, StopIteration):
                data_iter = iter(train_loader)
                x, y = next(data_iter)
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)

            with autocast(device_type=device.type, dtype=autocast_dtype):
                _, loss = model(x, y)
                loss = loss / tcfg.accum_steps
            scaler.scale(loss).backward()

        # 3) gradient clipping (operates on UNSCALED grads)
        if tcfg.grad_clip != 0.0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), tcfg.grad_clip)

        # 4) optimizer + scaler step
        scaler.step(optimizer)
        scaler.update()

        # 5) periodic evaluation
        if iter_num % 250 == 0 or iter_num == tcfg.max_iters - 1:
            val_loss = run_split("val")
            print(f"iter {iter_num} | lr {lr:.2e} | val {val_loss:.4f} | "
                  f"{(iter_num+1) / (time.time()-t0):.2f} it/s")
            best_val = min(best_val, val_loss)
        iter_num += 1

    print(f"done. best val loss: {best_val:.4f}")

if __name__ == "__main__":
    main()

Three things make this loop production-shaped rather than textbook-shaped:

  1. Loss is divided by accum_steps before backward. This keeps the effective per-token gradient magnitude identical to the non-accumulated case. Skipping this is the #1 bug in junior implementations.
  2. scaler.unscale_ before clip_grad_norm_. Clipping must happen on the unscaled gradients or the threshold is meaningless when GradScaler has been multiplying losses by 2**N.
  3. set_to_none=True on zero_grad. This avoids allocating a zero tensor for every parameter; it matters once your model has tens of millions of parameters.

Step 8 — Quick inference sanity check

You should be able to load a checkpoint and generate tokens. This is the proof your training worked.

# sample.py
import torch
from model import GPT
from config import GPTConfig

ckpt = torch.load("ckpt.pt", map_location="cuda")
cfg = GPTConfig(**ckpt["config"])
model = GPT(cfg).to("cuda")
model.load_state_dict(ckpt["model"])
model.eval()

import tiktoken
enc = tiktoken.get_encoding("gpt2")
ids = torch.tensor(enc.encode_ordinary("ROMEO:"), dtype=torch.long, device="cuda")[None, ...]

with torch.no_grad(), torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
    for _ in range(200):
        logits, _ = model(ids[:, -cfg.block_size:])
        next_id = torch.multinomial(torch.softmax(logits[:, -1], dim=-1), num_samples=1)
        ids = torch.cat([ids, next_id], dim=1)
print(enc.decode(ids[0].tolist()))

Running and Testing It

A portfolio project that “works on my machine” doesn’t get interviews. Make the run reproducible.

1. Run the data prep and a short training:

python data/prepare.py
python train.py    # should print val loss decreasing every 250 iters

2. Sanity-check the loss curve. On Tiny Shakespeare with the config above, you should see val loss drop from ~10.9 to ~1.4–1.6 over 6,000 iterations. If it stays flat at 10.9, your causal mask is wrong. If it explodes, your gradient clipping is missing. If NaNs appear with FP16, switch to BF16 — it’s the NVIDIA-recommended default for modern training.

3. Add tests for the parts that bite people. Even a small tests/ directory signals rigor.

# tests/test_attention.py
import torch
from model import CausalSelfAttention, GPTConfig

def test_causal_mask_blocks_future_tokens():
    cfg = GPTConfig(block_size=8, n_head=2, n_embd=8, n_layer=1, dropout=0.0)
    attn = CausalSelfAttention(cfg).eval()
    x = torch.randn(1, 8, 8)
    # Perturb a future token; output at t should not change.
    x2 = x.clone()
    x2[0, 7, :] = torch.randn(8)
    with torch.no_grad():
        y1 = attn(x)
        y2 = attn(x2)
    assert torch.allclose(y1[:, :7, :], y2[:, :7, :], atol=1e-5)

# tests/test_grad_accum.py
def test_grad_accum_matches_large_batch():
    # Same effective batch via one big step or N accumulated micro-steps.
    pass  # implement against your model

4. Profile a single training step. This is the single highest-leverage habit you can show on a CV.

# profile_step.py
from torch.profiler import profile, ProfilerActivity, record_function

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
             record_shapes=True) as prof:
    with record_function("train_step"):
        # one full optimizer.step()
        ...

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

The breakdown tells you whether you’re compute-bound (good — buy a bigger GPU), launch-bound (use CUDA Graphs, see PyTorch docs), or data-loader-bound (increase num_workers).

Extending It: Your Roadmap to Senior-Level

Once the basic loop runs, the project branches into five concrete upgrades. Pick the one that matches the role you want.

  1. Mixed-precision on CPU + channels-last memory format. Swap bfloat16 autocast to channels-last (model.to(memory_format=torch.channels_last)). This is what TorchAO builds on top of and what inference engines assume. Why it matters: signals you understand tensor memory layouts, not just dtypes.

  2. torch.compile the model. Wrap model = torch.compile(model, mode="reduce-overhead"). Show a 1.3–1.8× speedup on the same hardware. Why it matters: shows you track upstream PyTorch and can use it as a black-box optimizer when appropriate.

  3. Distributed training with FSDP. Replace the single-process loop with torch.distributed.fsdp so the model shards across GPUs. Why it matters: FSDP is the default at Meta, Anthropic, and most academic labs. Knowing it cold is a six-figure skill.

  4. Checkpointing and fault tolerance. Wrap training in torch.utils.checkpoint for activation memory, and write resumable checkpoints with optimizer + scheduler + RNG state to S3 every N steps. Why it matters: every multi-hour training run in production needs this; reviewers at training-infra teams look for it explicitly.

  5. Observability with W&B or TensorBoard. Log loss, LR, throughput (tokens/sec), GPU utilization, and gradient norms. Add a grad_norm callback and chart it — gradient explosions are visible there before they hit the loss. Why it matters: MLOps roles grade on this; it’s the difference between a research demo and a training pipeline.

  6. Benchmarking against torch.utils.benchmark. Reproducibly measure tokens/sec at batch sizes 1, 2, 4, 8 and plot the curve. Why it matters: shows you think in roofline terms rather than single-point metrics — the senior framing of “is the GPU busy or starving.”

A weekend gets you through steps 1–3 in the original list. Another weekend lands FSDP and proper checkpointing. After that, you’ve built something a senior ML engineer would actually recognize.

Key Takeaways

  • A from-scratch nanoGPT is the highest-signal-to-effort deep learning project for a CV: small surface area, every real-world ML systems concern is represented.
  • The three production patterns to internalize are mixed-precision (torch.amp.autocast + GradScaler), gradient accumulation (loss divided by accum_steps before backward), and the cosine LR schedule with linear warmup.
  • Weight tying between input embedding and output projection is a parameter-saving trick worth showing; causal mask correctness is a test you can write in 10 lines that catches most beginner bugs.
  • Use BF16 by default on modern GPUs; only use FP16 with a GradScaler if you must.
  • Extend the project with torch.compile, FSDP, resumable checkpointing, and observability — each upgrade maps to a different senior role (ML Platform, Training Infra, MLOps).

Further Reading