TL;DR — A mini GPT trained from scratch with bf16 mixed precision, a custom autograd engine, gradient checkpointing, and AdamW is a small-but-real systems build: ~500 lines of PyTorch that exercises the same machinery production LLM stacks use. This post walks through every piece with runnable code and ends with a roadmap to senior-level extensions.

Most “built a transformer” portfolio projects stop at nn.TransformerEncoderLayer and a copy of Karpathy’s nanoGPT. That’s fine — it proves you can read code. What hiring managers at ML systems teams actually scan for is whether you understand the infrastructure underneath: how autograd actually backpropagates, where memory goes, why bf16 differs from fp16, and how checkpointing trades compute for RAM. This project delivers all of that in a single repo you can demo in under ten minutes.

Why This Project Stands Out on a CV

Hiring managers don’t read every line of your GitHub. They read the README and skim the code for signals. This project emits several strong ones:

  • You understand mixed precision at the operator level, not just model.half(). The 2023-vintage AMP story in PyTorch is well documented; what hiring managers want to see is someone who can explain when bf16 beats fp16 (no loss scaling, larger dynamic range on H100/A100), why gradients must stay in fp32, and what the autocast context manager actually does. Implementing it manually proves you can.
  • You’ve built an autograd engine. Every ML framework ships one. Writing a 100-line version — even a simplified one — separates you from the 90% of applicants who treat autograd as magic. It also makes PyTorch internals legible: if you can explain torch.autograd.Function.backward, you can debug any custom op.
  • You know gradient checkpointing is a memory/compute trade, not a free lunch. Memory drops by ~√N, recompute cost is roughly 30%. Implementing it shows you reason about resource budgets the way an inference engineer does.
  • AdamW is the default optimizer for serious pretraining (used in GPT-2, BERT, RoBERTa, LLaMA, and most open-source fine-tunes). Knowing why weight decay is decoupled from the gradient update — and showing it in code — signals you’ve read Loshchilov & Hutter 2019, not just copy-pasted torch.optim.AdamW.
  • Roles this signals for: ML infrastructure engineer, training systems engineer, GPU performance engineer, research engineer, applied scientist at model labs, and increasingly the “AI platform engineer” roles at FAANG-adjacent shops.

One caveat: the project must actually run. A broken repo teaches the reviewer to doubt your other work. The testing section below is non-optional.

Architecture Overview

The build is split into six components, each small enough to reason about in isolation but wired together so the training loop is real:

  • autograd.py — Custom autograd engine. A Tensor class that wraps a NumPy or PyTorch array and records operations in a dynamic graph. Each op stores a backward closure. backward() performs reverse-mode autodiff via topological sort. Think tinygrad or micrograd, but biased toward clarity over speed.
  • ops.py — Primitive ops. matmul, add, gelu, softmax, cross_entropy, layer_norm. Each registers itself on the autograd graph. We deliberately skip PyTorch’s autograd in this file so the engine is the source of truth for forward/backward.
  • amp.py — Mixed-precision context. A manual equivalent of torch.autocast. On __enter__ it stashes the original dtypes and a flag; every Tensor constructor checks the flag and downcasts to bf16. Gradients always accumulate in fp32 (we’ll show why).
  • checkpoint.py — Gradient checkpointing. A checkpoint(fn, *args) helper that runs fn without autograd recording, then re-runs it inside a recording scope during backward. This is the exact pattern PyTorch’s torch.utils.checkpoint uses, per the PyTorch docs.
  • model.py — Mini GPT. Token + positional embeddings, N× decoder blocks (multi-head self-attention with causal mask, residual + LayerNorm, MLP), and an LM head. All weights are nn.Parameter with bf16 storage and fp32 master copies (the standard pattern described in Micikevicius et al., 2018).
  • train.py — Loop + AdamW. Pulls everything together: batch sampling from a text file, forward under autocast, backward with checkpointing, AdamW step on fp32 master weights, copy back to bf16.

Data flow per training step:

tokens → embed → [checkpoint(decoder_block) for _ in range(N)]
     → layer_norm → lm_head → cross_entropy → loss
     → backward (recomputes checkpointed blocks) → fp32 grads
     → AdamW on master weights → copy to bf16 params

The whole repo is ~500 lines. Drop it in a directory with a README.md and a small text corpus (TinyShakespeare is the canonical choice) and you’re done.

Building It Step by Step

I’ll walk through the five most instructive pieces. Full code lives in the companion repo; here we focus on the parts that teach.

Step 1 — The Tensor class

The core abstraction is a tensor that records its computational graph:

import numpy as np
from typing import List, Optional

class Tensor:
    def __init__(self, data, requires_grad=False, dtype=None, _children=(), _op=""):
        self.data = np.asarray(data, dtype=dtype) if dtype else np.asarray(data)
        self.requires_grad = requires_grad
        self.grad: Optional[np.ndarray] = None
        self._backward = lambda: None
        self._prev = set(_children)
        self._op = _op

    def __repr__(self):
        return f"Tensor(shape={self.data.shape}, dtype={self.data.dtype})"

    def backward(self):
        topo: List[Tensor] = []
        visited = set()
        def build(t):
            if id(t) in visited or not t.requires_grad:
                return
            visited.add(id(t))
            for child in t._prev:
                build(child)
            topo.append(t)
        build(self)
        self.grad = np.ones_like(self.data)
        for t in reversed(topo):
            t._backward()

That’s the engine. Every op will register a backward closure on its output’s _backward. Reverse-mode autodiff falls out for free from the topological order.

Step 2 — A primitive op, matmul

A worked example shows how an op plugs in:

def matmul(a: Tensor, b: Tensor) -> Tensor:
    out = Tensor(a.data @ b.data, requires_grad=a.requires_grad or b.requires_grad,
                 _children=(a, b), _op="matmul")

    def _backward():
        if a.requires_grad:
            a.grad = (out.grad @ b.data.swapaxes(-1, -2)) if a.grad is None \
                     else a.grad + (out.grad @ b.data.swapaxes(-1, -2))
        if b.requires_grad:
            b.grad = (a.data.swapaxes(-1, -2) @ out.grad) if b.grad is None \
                     else b.grad + (a.data.swapaxes(-1, -2) @ out.grad)
    out._backward = _backward
    return out

The pattern repeats for every op: compute forward with NumPy, write a _backward that uses the chain rule. gelu, softmax, layer_norm, cross_entropy all follow this template.

Step 3 — Manual mixed-precision context

PyTorch’s torch.autocast is convenient; implementing a stripped-down version clarifies what’s happening:

class autocast:
    def __init__(self, enabled=True, dtype="bfloat16"):
        self.enabled = enabled
        self.dtype = np.float32 if dtype == "float32" else np.bfloat16
        self._prev = None

    def __enter__(self):
        self._prev = Tensor._amp_enabled
        Tensor._amp_enabled = self.enabled and self.dtype is np.bfloat16
        return self

    def __exit__(self, *exc):
        Tensor._amp_enabled = self._prev

Tensor._amp_enabled = False

Inside Tensor.__init__, when Tensor._amp_enabled is true, activations downcast to bf16. Crucially, parameter master weights stay fp32; we keep bf16 copies for the forward pass and update the masters inside AdamW. This is the pattern NVIDIA documents in their mixed-precision training guide.

Step 4 — Gradient checkpointing

The checkpoint helper is small but conceptually dense:

def checkpoint(fn, *args):
    # Forward: run fn without recording, stash inputs for recompute
    with torch.no_grad():
        out = fn(*[a.detach() if hasattr(a, "detach") else a for a in args])

    # During backward, recompute fn inside the recording scope
    def backward(grad_out):
        replayed = fn(*args)  # re-runs forward with autograd active
        replayed.backward(grad_out)

    out._backward_fn = backward
    return out

This is exactly what torch.utils.checkpoint.checkpoint does, per the PyTorch docs: drop the intermediate tensors to free memory, then recompute them on demand during backward. Memory savings scale with the number of layers between checkpoints — typically you checkpoint every block.

Step 5 — AdamW on fp32 master weights

AdamW decouples weight decay from the gradient update, as defined in Loshchilov & Hutter 2019:

class AdamW:
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1):
        self.params = params
        self.lr = lr; self.b1, self.b2 = betas; self.eps = eps; self.wd = weight_decay
        self.t = 0
        self.m = [np.zeros_like(p.data) for p in params]
        self.v = [np.zeros_like(p.data) for p in params]

    def step(self):
        self.t += 1
        for i, p in enumerate(self.params):
            g = p.grad
            self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g
            self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (g * g)
            m_hat = self.m[i] / (1 - self.b1 ** self.t)
            v_hat = self.v[i] / (1 - self.b2 ** self.t)
            # Decoupled weight decay — applied directly, not via gradient
            p.data -= self.lr * (m_hat / (np.sqrt(v_hat) + self.eps) + self.wd * p.data)

    def zero_grad(self):
        for p in self.params:
            p.grad = None

The p.data -= self.lr * (... + self.wd * p.data) line is the entire reason AdamW exists. Without the + self.wd * p.data term, you get L2 regularization through the gradient (as in Adam), which interacts badly with the adaptive learning rate. Decoupling it gives the clean L2 regularization that BERT and GPT-2 used.

Step 6 — Putting it together

The training loop is the moment of truth:

import numpy as np
from autograd import Tensor
from amp import autocast
from checkpoint import checkpoint
from model import MiniGPT
from optimizer import AdamW

model = MiniGPT(vocab_size=4096, n_layers=6, n_heads=6, d_model=384, max_seq=256)
# Master fp32 weights + bf16 working copies
master_params = [Tensor(p.data.copy(), requires_grad=True) for p in model.params]
optim = AdamW(master_params, lr=3e-4, weight_decay=0.1)

for step, batch in enumerate(data_loader("tinyshakespeare.txt", block_size=256, batch=8)):
    tokens = Tensor(batch, requires_grad=False)
    with autocast(enabled=True, dtype="bfloat16"):
        loss = checkpoint(model.forward, tokens)  # checkpointed forward
    loss.backward()

    # Gradient clipping — standard practice, see in [PyTorch docs](https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html)
    grad_norm = np.sqrt(sum(np.sum(g**2) for g in [p.grad for p in master_params]))
    for p in master_params:
        p.grad *= min(1.0, 1.0 / grad_norm) if grad_norm > 1.0 else 1.0

    optim.step()
    optim.zero_grad()

    # Sync bf16 working copies from fp32 masters
    for w, m in zip(model.params, master_params):
        w.data = m.data.astype(np.bfloat16)

    if step % 50 == 0:
        print(f"step {step:5d} | loss {loss.data:.4f} | gnorm {grad_norm:.2f}")

Six numbered items, ~500 lines total, and you’ve exercised the entire LLM training stack at a level that demonstrates real understanding.

Running and Testing It

The project must run end-to-end. Here’s how to prove it works:

Prerequisites. PyTorch ≥ 2.1 (for native bf16 on CPU/MPS/Metal/CUDA), NumPy, and a small text corpus. Karpathy’s TinyShakespeare is ~1 MB and trains in minutes.

Quick start.

git clone https://github.com/yourhandle/mini-gpt-bf16
cd mini-gpt-bf16
pip install torch numpy
python train.py --data tinyshakespeare.txt --steps 2000 --batch 8 --block 256

Expected output on a single A100 or M-series Mac:

step     0 | loss  8.31 | gnorm 1.00
step    50 | loss  6.84 | gnorm 0.92
step   500 | loss  3.21 | gnorm 0.74
step  2000 | loss  2.07 | gnorm 0.61

Tests that prove it’s real. The repo should ship with these:

def test_autograd_matches_pytorch():
    """Numerical agreement between custom engine and torch autograd."""
    import torch
    a = Tensor(np.random.randn(3, 4), requires_grad=True)
    b = Tensor(np.random.randn(4, 5), requires_grad=True)
    out = matmul(a, b)
    out.backward(np.ones_like(out.data))
    ta = torch.tensor(a.data, requires_grad=True)
    tb = torch.tensor(b.data, requires_grad=True)
    torch.matmul(ta, tb).sum().backward()
    assert np.allclose(a.grad, ta.grad.numpy(), atol=1e-5)

def test_checkpoint_saves_memory():
    """Peak memory drops at least 30% vs uncheckpointed."""
    import tracemalloc
    tracemalloc.start()
    run_uncheckpointed()
    base = tracemalloc.get_traced_memory()[1]
    tracemalloc.reset_peak()
    run_checkpointed()
    peak = tracemalloc.get_traced_memory()[1]
    assert peak < 0.7 * base

def test_loss_decreases():
    """Smoke test: 200 steps should drop loss by ≥ 1 nat."""
    final = train(steps=200)
    assert final < initial - 1.0

Sanity checks that hiring managers actually run. They will:

  1. Read the README. Make sure it explains why bf16, why fp32 masters, why checkpoint every block.
  2. Clone and python train.py. If it crashes on import, you’re done.
  3. Skim autograd.py. If it’s 1000 lines of clever metaprogramming, that’s a red flag. If it’s 150 lines of clear ops with _backward closures, that’s a green flag.
  4. Look at the commit history. Six commits over six weeks beats one commit with everything.

Extending It: Your Roadmap to Senior-Level

The base build is the floor. Here’s how to grow it into something that reads as production-flavored to a senior reviewer. Each upgrade is small in isolation but compounds.

  1. Weight persistence with safetensors + atomic writes. Replace np.savez with safetensors.numpy.save_file and write to model.safetensors.tmp then os.replace() for crash safety. Why it matters: Hugging Face’s safetensors format is the de facto standard for model distribution; atomic writes are what every training daemon needs.

  2. Multi-process data loading with a bounded queue. Use torch.multiprocessing or concurrent.futures with a prefetch buffer of 4× batch size. Why it matters: GPU starvation from a slow data pipeline is the #1 cause of “training is slower than expected” in production.

  3. Loss/grad-norm/log-throughput to W&B or TensorBoard. Wrap the training loop in a Logger class that emits per-step metrics. Add a learning-rate finder following Leslie Smith’s 2017 paper. Why it matters: Observability is the difference between “I trained a model” and “I ran an experiment”.

  4. Fault tolerance with periodic checkpoint + resume. Save {model, optim, step, rng} every N steps; on startup, resume if a checkpoint exists. Why it matters: This is literally what AWS SageMaker and Kubernetes + Volcano provide for LLM training jobs — implementing it yourself proves you understand the contract.

  5. A benchmarking harness with nsight or PyTorch profiler. Wrap each phase (forward, backward, optimizer step, data load) in torch.cuda.profile regions and dump a Chrome trace. Compare checkpointed vs uncheckpointed throughput. Why it matters: Performance engineering is a separate discipline; even a basic profile demonstrates the instinct.

  6. Distributed training with PyTorch DDP. Add torch.distributed.init_process_group("nccl"), wrap the model in DistributedDataParallel, and shard the data loader with DistributedSampler. Why it matters: Multi-GPU training is the bread-and-butter of any serious pretraining setup; DDP is the canonical approach.

Any three of these, done well, push the project from “tutorial” to “I’m a systems engineer who knows what they’re doing.”

Key Takeaways

  • A custom autograd engine plus manual mixed precision plus gradient checkpointing is a ~500-line project that signals deep ML systems knowledge without requiring a GPU cluster.
  • bf16 is the right default on modern hardware (Ampere and later): no loss scaling, wider dynamic range than fp16, and identical training dynamics in practice for most transformer workloads.
  • Gradient checkpointing trades ~30% extra FLOPs for memory that scales as O(√N) in the number of layers — a great teaching example of resource budgeting.
  • AdamW’s decoupled weight decay is the small but essential detail that distinguishes serious training code from copy-pasted tutorials.
  • The CV signal isn’t “I built a transformer”; it’s “I understand why every component is in the loop.” The README and tests do as much work as the code.

Further Reading