There’s an old saying in engineering: “You don’t really understand something until you can build it.” This has never been more true than in the era of LLMs. While we’ve previously explored the foundational concepts in my post on the Transformer architecture explained here, true understanding comes from implementation. That’s why today, we’re building a GPT-style model (the 124M variant) from scratch in PyTorch.
This post pairs well with my Transformer Architecture deep dive: read that first if attention mechanisms are new to you.
This post has a different focus than my last “from scratch” endeavor, where I built an entire deep learning framework to grasp the low-level mechanics of autograd and tensor ops. Here, we’ll use PyTorch’s primitives to focus on what makes GPT special: multi-head attention, positional encodings, and the specific architectural decisions that enable language understanding.
GPT (Generative Pre-trained Transformer) models, developed by OpenAI, represent a breakthrough in natural language processing. GPT-2, released in 2019, demonstrated that a transformer-based model trained on vast amounts of text could generate remarkably coherent and contextually relevant content. GPT-3, its successor, scaled this approach to 175 billion parameters and showed emergent capabilities like few-shot learning and complex reasoning. Both models share the same fundamental architecture: stacked transformer decoder blocks that predict the next token in a sequence, trained on the simple objective of minimizing prediction error across massive text corpora.
The 124M variant has 12 layers, 12 heads, and 768-dim embeddings. GPT-3 scaled to 96 layers, 96 heads, and 12288-dim: same architecture, 1000x parameters.
The 124M parameter version we’ll be building captures the essential architecture while remaining computationally tractable for individual developers, though even at this “small” scale, you’ll quickly discover why the ML community spends so much time optimizing both training efficiency and model performance.
By the end of this journey, you won’t just know how transformers work; you’ll have built the critical components with your own hands, optimized the training loop, and watched your model evolve from random noise to coherent text generation. Let’s begin.
Beyond the architecture itself, we use a few optimizations that make a large difference to training throughput:
TensorFloat32 (TF32): NVIDIA’s precision format that keeps FP32’s 8-bit exponent but truncates the mantissa from 23 bits to 10. Matrix multiplications then run on tensor cores at roughly 2-3x the throughput of true FP32, with no measurable quality loss. One line of code enables it.
BFloat16 with Autocast: Mixed precision training using brain floating-point format, which maintains the same exponent range as FP32 but reduces mantissa precision. Combined with automatic mixed precision (AMP), this halves memory usage and speeds up training.
torch.compile: PyTorch 2.0’s just-in-time compilation, which traces the model into a graph, fuses operations, and generates kernels that avoid needless round trips to memory. It is a single line to enable, at the cost of a slow first step while compilation happens.
Flash Attention: An algorithmic improvement that computes attention without ever materializing the full \(T \times T\) score matrix. It streams the softmax, carrying a running maximum and normalizer for each row so blocks of keys can be folded in as they are read, which drops attention memory from \(O(n^2)\) to \(O(n)\).
Fused AdamW: A single-kernel implementation of the AdamW optimizer that reduces memory reads/writes by computing all parameter updates in one pass, cutting the time of the optimizer step by up to half.
Annealed Learning Rate: A linear warmup followed by cosine decay, the schedule that has become standard for transformers. We’ll derive why the warmup is there instead of treating it as convention.
Weight Decay Only on Matrices: Weight decay applies to the weight matrices in Linear and Embedding layers, and not to biases or layer normalization parameters. Those 1D parameters set scale and offset rather than mixing inputs, so pulling them toward zero fights the layer instead of regularizing it.
Vocabulary Padding: Rounding the vocabulary from 50257 up to 50304 so the dimension divides evenly by 64. We add embedding rows that can never be emitted, and every matmul touching the vocab dimension gets faster.
Distributed Data Parallelism (DDP): Scaling training across multiple GPUs using PyTorch’s DDP, including gradient synchronization, correct data sharding, and keeping model state consistent across devices.
TF32, bfloat16, flash attention, fused AdamW, and vocabulary padding look like five unrelated tricks, but they are one trick applied five times. A modern GPU can do arithmetic far faster than it can move operands to and from memory, so most training steps are memory-bound, not compute-bound. Every optimization here either moves fewer bytes (lower precision, padded dimensions that tile cleanly) or moves them fewer times (fusing operations so intermediates never leave fast memory). Flash attention performs more arithmetic than the naive implementation and is faster anyway, because it stops writing the \(T \times T\) score matrix to memory.
The GPT-2 paper omits certain architectural details and hyperparameter specifications, so we refer to the GPT-3 paper to fill these gaps. The core architecture is consistent between the two models, which is what makes the substitution safe.
We start with imports and two small helpers. listify normalizes arguments, and the annealer decorator lets us compose the warmup-then-cosine learning rate schedule further down.
import inspect
import math
import os
import time
from dataclasses import dataclass
from functools import partial, wraps
from typing import Callable, Iterable
import tiktoken
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as opt
from torch.distributed import destroy_process_group, init_process_group
from torch.nn.parallel import DistributedDataParallel as DDPdef listify(obj):
if obj is None:
return []
elif isinstance(obj, str):
return [obj]
elif isinstance(obj, list):
return obj
elif isinstance(obj, Iterable):
return list(obj)
else:
return [obj]def annealer(func: Callable):
@wraps(func)
def annealer_wrapper(*args, **kwargs):
return partial(func, *args, **kwargs)
return annealer_wrapper
@annealer
def lin_sched(start, end, pos):
"""Linear scheduler."""
return start + (end - start) * pos
@annealer
def cos_sched(start, end, pos):
"""Cosine scheduler."""
return start + (1 + math.cos(math.pi * (1 - pos))) * (end - start) / 2
def combine_scheds(pcts, scheds):
"""
Combine multiple schedulers, each run for a given percentage of the
training process.
"""
assert len(pcts) == len(scheds), "Each scheduler should have its `pct`."
assert sum(pcts) == 1.0, "Sum of the `pcts` should be equal to 1."
pcts = torch.tensor([0] + listify(pcts))
assert (pcts >= 0).all(), "All percentages should be non-negative."
pcts = torch.cumsum(pcts, 0)
def _inner(pos):
idx = (pos >= pcts).nonzero().max()
actual_pos = (pos - pcts[idx]) / (pcts[idx + 1] - pcts[idx])
return scheds[idx](actual_pos)
return _innerOne dataclass holds every architectural and training hyperparameter. The values are GPT-2 124M: 12 layers, 12 heads, 768-dimensional embeddings, and a 1024-token context.
@dataclass
class GPTConfig:
block_sz: int = 1024 # Sequence length
# Originally 50000 BPE merges + 256 byte tokens + 1 for <|endoftext|> token
# which will delimit different documents. This token's index is 50256
# However, we found that using 50257 as the vocab size is not a multiple of
# 64 and we could improve efficiency and performance (through better
# occupancy) if we round up to the closest multiple of 64, which is 50304.
vocab_sz: int = 50257
n_layer: int = 12 # Number of layers
n_embd: int = 768 # Embedding dimension
n_head: int = 12 # Number of attention heads
lr: float = 3e-4 # Good for big models
batch_sz: int = 4
dropout: float = 0.0
bias: bool = TrueThe MLP runs on each token independently, with no interaction between positions. Because it projects out to 4 x n_embd and back, it holds roughly two-thirds of the model’s parameters.
class MLP(nn.Module):
def __init__(self, config: GPTConfig):
# Point-wise feed-forward network that applies non-linearity
# on every token separately. THERE IS NO INTERACTION BETWEEN TOKENS.
# This is where almost all the capacity and non-linearities of the
# model come from especially when we project it to 4 x n_embd
super().__init__()
self.c_fc = nn.Linear(config.n_embd, config.n_embd * 4)
# Found to be better than ReLU in terms of gradient saturation
self.gelu = nn.GELU(approximate="tanh")
self.c_proj = nn.Linear(config.n_embd * 4, config.n_embd)
self.dropout = nn.Dropout(config.dropout)
self.c_proj.NANOGPT_SCALE_INIT = 1
def forward(self, x):
return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))Attention is the one operation that lets tokens exchange information. All three projections live in a single fused c_attn matrix that we split after the matmul, and F.scaled_dot_product_attention supplies both Flash Attention and the causal mask.
class CausalSelfAttention(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
self.n_head = config.n_head
self.n_embd = config.n_embd
# For optimization purposes, QKV matrics will be in one big matrix
# that will be split after we do matmul with X
self.c_attn = nn.Linear(config.n_embd, config.n_embd * 3, bias=config.bias)
# Main job of this projection is for the interaction between heads
# because they worked in isolation -> cross-head communication
# that would lead to suppressing or amplify the combination
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
self.attn_dropout = nn.Dropout(config.dropout)
self.resid_dropout = nn.Dropout(config.dropout)
self.dropout = config.dropout
self.c_proj.NANOGPT_SCALE_INIT = 1
# NOTE: Mask is not needed when we use Pytorch's Flash attention
# self.register_buffer(
# "mask",
# torch.tril(torch.ones(config.block_sz, config.block_sz))
# )
def forward(self, x):
B, T, _ = x.shape
qkv = self.c_attn(x)
# q/k/v is B x T x n_embd each
q, k, v = torch.split(qkv, self.n_embd, dim=-1)
# Reshape q/k/v to B x n_head x T x (n_embd / n_head)
# So each head would be learning different kind of relationships
q = q.view(B, T, self.n_head, self.n_embd // self.n_head).transpose(1, 2)
k = k.view(B, T, self.n_head, self.n_embd // self.n_head).transpose(1, 2)
v = v.view(B, T, self.n_head, self.n_embd // self.n_head).transpose(1, 2)
# attn is B x T x T
# attn = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.shape[-1]))
# Mask out future tokens
# attn = attn.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
# attn = self.attn_dropout(F.softmax(attn, dim=-1))
# y is B x T x n_embd
# y = attn @ v
# Uses Flash attention that never materialize attention matrices for
# each head and is aware of the memory hierarchy and tries to reduce
# read/writes with more FLOPs -> Speed up since we're memory bound
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
y = y.transpose(1, 2).contiguous().view(B, T, self.n_embd)
return self.resid_dropout(self.c_proj(y))A block wraps attention and the MLP in residual connections. The normalization sits before each sublayer instead of after, which departs from Attention Is All You Need. That ordering keeps the residual path an unobstructed identity from input to output, so gradients reach the early layers without crossing a normalization at every step. Deep stacks are much easier to train this way.
class Block(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
self.ln_1 = nn.LayerNorm(config.n_embd)
self.ln_2 = nn.LayerNorm(config.n_embd)
self.mlp = MLP(config)
self.attn = CausalSelfAttention(config)
def forward(self, x):
# Use Pre-layer normalization which deviates from the
# transformer original paper that uses post-layer normalization.
# This should help stabilize training
x = x + self.attn(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return xThis class assembles the embeddings, the stack of blocks, and the language modeling head. It also carries weight tying between the token embedding and the output head, which accounts for 38.6M of the model’s 124M parameters, the residual-aware initialization scaling derived below, and weight decay applied only to matrices. One thing to notice in forward: it hands raw logits to F.cross_entropy and never applies a softmax of its own, for the reasons in Why Your Final Layer Shouldn’t Have Softmax.
class GPT2(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
self.config = config
self.transformer = nn.ModuleDict(
dict(
wte=nn.Embedding(config.vocab_sz, config.n_embd),
# Attention operation is a permutation equivariant, this means that
# if we permute the input then the corresponding output will be
# permuted in exactly the same way. In other words, attention mechanism
# is not aware of the relative ordering of the tokens. Therefore, we
# need some way to encode the positions of the tokens in each sequence.
# This is where positional encoding comes into play.
# Here we use a simple positional encoding that is a simple
# embedding of the position of the token in the sequence.
wpe=nn.Embedding(config.block_sz, config.n_embd),
h=nn.ModuleList(
[Block(config) for _ in range(config.n_layer)]
),
# Final layer norm after all transformer layers
ln_f=nn.LayerNorm(config.n_embd),
)
)
# Sometimes called classifier
self.lm_head = nn.Linear(config.n_embd, config.vocab_sz, bias=False)
# Weight sharing between the token embedding layer and
# last linear layer (LM head classifier). The rationale is
# that tokens that are semantically similar to each other in
# the embedding space should have similar probabilities in the
# softmax of the LM head layer.
# This is also one of the biggest matrices in the model: 50257 * 768 =
# 38.6M parameters, ~31% of the model's 124M total. Without tying, the
# same model would carry that matrix twice and weigh in at ~163M.
self.transformer.wte.weight = self.lm_head.weight
self.apply(self._init_weights)
def _init_weights(self, module):
# The following initialization comes from gpt2 src code
# NOTE: Because token embedding and classifier weights are shared,
# our initialization logic will initialize the weight matrix twice
# but shouldn't be an issue since they're being initialized with the
# same std and mean
if isinstance(module, nn.Linear):
std = 0.02
# We're changing std because residual path affect std
# by increasing it on every layer so we need to adjust
# it so we still have the same std = 0.02
if hasattr(module, "NANOGPT_SCALE_INIT"):
# `2` here because every layer has two blocks:
# - Attention block
# - MLP block
# `N` is the number of layers in the model (n_layer)
# Since they are independent, variance of the sum of the two
# blocks is the sum of the variances
std *= (2 * self.config.n_layer) ** -0.5
nn.init.normal_(module.weight, std=std)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
# We're initializing the token and positional embeddings
# with the same std but the paper initialized the positional
# embedding with std = 0.01
nn.init.normal_(module.weight, std=0.02)
def forward(self, x, targets=None):
T = x.shape[-1]
assert (
T <= self.config.block_sz
), f"Sequence length must be <= {self.config.block_sz}, got {T}"
pos_emb = self.transformer.wpe(
torch.arange(0, T, dtype=torch.long, device=x.device)
)
tok_emb = self.transformer.wte(x)
x = pos_emb + tok_emb
for block in self.transformer.h:
x = block(x)
x = self.transformer.ln_f(x)
# logits is B x T x vocab_sz
logits = self.lm_head(x)
loss = None
if targets is not None:
# F.cross_entropy expects the 2nd dimension to be probabilities
loss = F.cross_entropy(
logits.view(-1, self.config.vocab_sz), targets.view(-1)
)
return logits, loss
def configure_optimizer(self, weight_decay, lr, device):
params_dict = {
pn: p for pn, p in self.named_parameters() if p.requires_grad
}
# We're not applying weight decay to bias and layer norm parameters
# And any 1D parameters. Therefore, we are ONLY applying weight decay
# to the weight matrices in Embedding and Linear layers
decay_params = [p for p in params_dict.values() if p.ndim >= 2]
nondecay_params = [p for p in params_dict.values() if p.ndim < 2]
params_groups = [
{"params": decay_params, "weight_decay": weight_decay},
{"params": nondecay_params, "weight_decay": 0.0},
]
# Fused AdamW is available for PyTorch 2.0+
fused_available = "fused" in inspect.signature(opt.AdamW).parameters
use_fused = fused_available and "cuda" in device
return opt.AdamW(
params_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8, fused=use_fused
)
@torch.no_grad
def generate(self, idxs: torch.Tensor, max_tokens: int = 5):
for i in range(max_tokens):
# x would be B x T x vocab_sz (At most we would have
# block_sz tokens since we're using fixed block_sz for the
# positional embedding
idxs = idxs[:, -self.config.block_sz :]
logits, _ = self(idxs)
# Get probs for last token to predict next token
# This would be B x vocab_sz
logits = logits[:, -1, :]
# Apply softmax to get probabilities
probs = F.softmax(logits, dim=-1)
# Pick top 50 prob -> we would never pick tokens with
# very small probs (right tails) -> B x 50
# probs/idxs are sorted in descending order
topk_probs, topk_idxs = torch.topk(probs, 50, dim=-1)
# Sample 1 token from the top 50 tokens -> idx is B x 1
idx = torch.multinomial(topk_probs, 1)
# Get the vocab idx as `multinomial` returns only indices that
# corresponds to the given array
idx = torch.gather(topk_idxs, -1, idx)
# TODO: We should check for end_of_text token and break out of
# the loop (stop generation) even if we have not reached max_tokens
idxs = torch.cat([idxs, idx], dim=1)
return idxsThe line std *= (2 * self.config.n_layer) ** -0.5 in _init_weights looks like a tuned constant, but it falls out of how the residual stream accumulates.
A block writes into the stream twice, in sequence, exactly as Block.forward does:
\[x_n' = x_{n-1} + a_n, \qquad x_n = x_n' + m_n\]
with \(a_n = \mathrm{attn}(\mathrm{LN}_1(x_{n-1}))\) and \(m_n = \mathrm{MLP}(\mathrm{LN}_2(x_n'))\). These are not independent draws, since \(m_n\) is computed from \(x_n'\), which already contains \(a_n\). What lets their variances add is the output projection: each sublayer ends in a c_proj initialized zero-mean and independent of its input, so \(\mathbb{E}[m_n \mid x_n'] = 0\) and \(\mathrm{Cov}(x_n', m_n) = 0\). A write is uncorrelated with the stream it lands on even though it was computed from it. Pre-normalization then fixes each write’s size: attn and MLP see unit variance however wide the stream has grown, so \(\mathrm{Var}(a_n) = \mathrm{Var}(m_n) = \sigma^2\) at every depth. Applying this at each step and unrolling over \(N\) blocks, from embeddings at \(\sigma_0^2\):
\[\mathrm{Var}(x_N) = \sigma_0^2 + 2N\sigma^2\]
So the standard deviation grows like \(\sqrt{2N}\), set by depth alone. With \(\sigma = \sigma_0 = 0.02\) and \(N = 12\), the stream leaves the stack at \(0.02\sqrt{25} = 0.1\), five times the width it entered with. Choosing \(\sigma = \sigma_0(2N)^{-1/2}\) instead makes the accumulated term \(2N \cdot \sigma_0^2/(2N) = \sigma_0^2\), so \(\mathrm{Var}(x_N) = 2\sigma_0^2\) and total growth is \(\sqrt2\) at any depth. The 2 counts the writes per block and the \(-1/2\) undoes the square root. It is also why NANOGPT_SCALE_INIT is set on c_proj in MLP and CausalSelfAttention and nowhere else: those are the two projections that write into the stream.
The loader tokenizes the corpus once with tiktoken’s GPT-2 BPE encoding (which I build from scratch in Byte Pair Encoding, and place among the alternatives in Tokenization Strategies), then hands out contiguous windows. Sharding is where this gets subtle. Each process must take a disjoint slice, and all processes must wrap around together; otherwise the ranks drift out of step and train on overlapping data with nothing to signal it.
class DataLoaderLight:
def __init__(
self,
file_path: str,
batch_sz: int,
block_sz: int,
process_rank: int = 0,
number_processes: int = 1,
) -> None:
self.batch_sz = batch_sz
self.block_sz = block_sz
self.process_rank = process_rank
self.number_processes = number_processes
with open(file_path, "r") as f:
text = f.read()
encoder = tiktoken.get_encoding("gpt2")
tokens = torch.tensor(encoder.encode(text), dtype=torch.long)
# We can truncate the tokens list/tensor based on a multiple of
# `batch_sz x block_sz x number_processes`. This is useful
# for multi-node training and mimics the behavior of DataLoader's
# `drop_last` parameter. Think of it as for every iteration we sample
# batch_sz x block_sz x number_processes input tokens where each
# process would get batch_sz x block_sz tokens -> Each iteration would
# lead to all processes running the same load and no process is idling
self.tokens_per_process = self.batch_sz * self.block_sz
self.tokens_per_iteration = (
self.tokens_per_process * self.number_processes
)
# Every global iteration predicts `tokens_per_iteration` labels, so
# it needs one additional token to label the final input token.
self.num_batches = (
len(tokens) - 1
) // self.tokens_per_iteration
if self.num_batches < 1:
raise ValueError(
"dataset must contain at least "
f"{self.tokens_per_iteration + 1} tokens"
)
self.tokens = tokens[:self.num_batches * self.tokens_per_iteration + 1]
print(f"Loaded {len(self.tokens)} tokens")
print(f"1 epoch = {len(self)} batches")
self.current_batch = 0
def __len__(self):
return self.num_batches
def next_batch(self):
global_start = self.current_batch * self.tokens_per_iteration
process_start = (
global_start + self.process_rank * self.tokens_per_process
)
process_end = process_start + self.tokens_per_process + 1
buf = self.tokens[process_start:process_end]
x = buf[:-1].view(self.batch_sz, self.block_sz)
y = buf[1:].view(self.batch_sz, self.block_sz)
# Each process will process batch_sz x block_sz tokens in each
# iteration -> with number_processes processes, total tokens processed
# in each iteration is batch_sz x block_sz x number_processes. In the
# case of one process, total tokens would be batch_sz x block_sz
# Similar to DataLoader's `drop_last` parameter, we drop the last
# batch if it's not a multiple of batch_sz x block_sz x number_processes.
# Tracking the global batch directly makes every process reset together
# after consuming the same number of complete global batches.
self.current_batch = (self.current_batch + 1) % self.num_batches
return x, ytorchrun sets RANK, LOCAL_RANK, and WORLD_SIZE in the environment. This cell reads them to initialize the process group and pin each process to its own device.
###########
## Distributed Data Parallel
###########
# Distributed Data Parallel let us run the same model (replica) on different GPUs,
# where each GPU would work on a different slice of data. After we do the backward
# pass, we average the gradients across all processes (GPUs) and synchronize all
# parameters across all devices. We use allReduce op to do this and communicate the
# updates with all processes.
# Each process would go through the same code from top to bottom not aware there
# are other processes running the same thing on other devices
#
# torchrun command sets the following environment variables:
# RANK: Id of the process in the process group. It is an int 0 - WORLD_SIZE
# LOCAL_RANK: In the case of multi-nodes, LOCAL_RANK is the id of
# the process in the same node. Example: If we have a node
# with 4 GPUs, the first process will have LOCAL_RANK = 0
# but RANK of this process may not be 0 if we are running
# on multiple nodes.
# This is useful when we have multiple nodes and we want to
# run the processes on different GPUs in the same node.
# In this case, we can set the LOCAL_RANK to the GPU id in the
# node.
# WORLD_SIZE: Total number of processes
ddp = int(os.getenv("RANK", -1)) != -1 # Check if it is a ddp run
if ddp:
# DDP requires CUDA so we need to set the device for each process
# so only one process can run per device
assert torch.cuda.is_available(), "DDP requires CUDA"
init_process_group(backend="nccl")
ddp_rank = int(os.getenv("RANK"))
ddp_local_rank = int(os.getenv("LOCAL_RANK"))
ddp_world_size = int(os.getenv("WORLD_SIZE"))
device = f"cuda:{ddp_local_rank}"
torch.cuda.set_device(device)
# Master process will do more things such as checkpointing and logging
# while other processes would assist in the computations.
# It always has RANK = 0
master_process = ddp_rank == 0
else:
ddp_rank = 0
ddp_local_rank = 0
ddp_world_size = 1
master_process = True
if torch.cuda.is_available():
device = "cuda"
elif torch.mps.is_available():
device = "mps"
torch.mps.manual_seed(1337)
else:
device = "cpu"
print(device)torch.manual_seed(1337)
if torch.cuda.is_available():
torch.cuda.manual_seed(1337)Here we instantiate the model with the padded vocabulary, then build the fused AdamW optimizer and the learning rate schedule it will follow.
##########
## Initialize model and optimizer
##########
# GPU kernels tile in powers of two, so dimensions that divide evenly by
# 64 or 128 make far better use of:
# • Tensor Cores
# • Memory coalescing
# • Shared memory bank alignment
# • Warp scheduling
# Here we round vocab_sz from 50257 up to 50304, the next multiple of 64.
# 50304 is not itself a power of two (it is 64 x 786); what matters is that
# the dimension divides evenly, so the matmul tiles without a ragged
# remainder. We pay for a few hundred embedding rows that can never be
# emitted and get faster kernels everywhere the vocab dimension appears.
model = GPT2(GPTConfig(vocab_sz=50304)).to(device)
# Speed up model by building static graph that analyzes all ops
# and optimizes them such as fusing some of them to avoid unnecessary
# trips to memory. Uncomment to enable: it costs a slow first step while
# the graph is traced and compiled, which is why it is off by default here.
# model = torch.compile(model)
if ddp:
model = DDP(model, device_ids=[ddp_local_rank])
raw_model = model.module if ddp else model
max_lr = 3e-4
min_lr = max_lr * 0.1
warmup_steps = 10
max_steps = 50
sched = combine_scheds(
[warmup_steps / max_steps, 1 - (warmup_steps / max_steps)],
[lin_sched(min_lr, max_lr), cos_sched(max_lr, min_lr)],
)
optimizer = raw_model.configure_optimizer(
weight_decay=0.1, lr=max_lr, device=device
)The loop is where the optimizations come together: gradient accumulation to reach a 0.5M-token batch, bfloat16 autocast, TF32 matmuls, gradient clipping, and the learning rate schedule. Two details in this loop are easy to get wrong in ways that produce no error and no obviously bad number.
##########
## Run training loop
#########
# NOTE: In order to run 0.5M (from GPT3 paper) tokens per fwd/bwd iteration,
# we need to use gradient accumulation because we can't fit it in almost
# any commodity GPU -> We only do backward after we loop through ~0.5M tokens.
total_batch_sz = 2**19 # closest number to 0.5M
assert total_batch_sz % (GPTConfig.batch_sz * GPTConfig.block_sz * ddp_world_size) == 0, "total batch size must be divisible by micro batch_sz x block_sz x ddp_world_size"
grad_accum_steps = total_batch_sz // (
GPTConfig.batch_sz * GPTConfig.block_sz * ddp_world_size
)
if master_process:
print(f"Total desired batch size: {total_batch_sz}")
print(f"Calculated gradient accumulation steps: {grad_accum_steps}")
train_dl = DataLoaderLight(
"tinyshakespeare.txt",
batch_sz=GPTConfig.batch_sz,
block_sz=GPTConfig.block_sz,
process_rank=ddp_rank,
number_processes=ddp_world_size
)
# Pytorch will use TensorFloat32 if available, else use FP32
# But the weights will still be stored using FP32 with less precision
# (10 bits for mantissa instead of 23). It is just the
# operations would be executed as TF32 if available
torch.set_float32_matmul_precision("high")
# `autocast` wants a device TYPE ("cuda"), not a full device string. DDP
# assigns each process "cuda:{local_rank}", so comparing against "cuda"
# directly would silently skip mixed precision on every multi-GPU run.
device_type = "cuda" if device.startswith("cuda") else device
for step in range(max_steps):
start = time.time()
optimizer.zero_grad()
loss_accum = 0.0
for macro_step in range(grad_accum_steps):
# Draw a fresh micro-batch on every accumulation step. Fetching it
# outside this loop would add up `grad_accum_steps` copies of the
# same gradient and leave the effective batch at batch_sz x block_sz
x, y = train_dl.next_batch()
x = x.to(device)
y = y.to(device)
if device_type == "cuda":
# Tensors that will be greatly affected by less precision such
# as loss, layernorm would still be in FP32 while others such
# as attention weights would be in BF16
with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
logits, loss = model(x, y)
else:
logits, loss = model(x, y)
# Just accumulating gradients yield to summation of objective but
# we want mean so we weight each loss by 1/grad_accum_steps
loss /= grad_accum_steps
loss_accum += loss.detach()
# To avoid syncing the gradients between the processes after every
# macro step, we disable it and only allows the sync up of
# gradients after we finish all gradient accumulation in each
# process
if ddp:
model.require_backward_grad_sync = (
macro_step == grad_accum_steps - 1
)
loss.backward()
# Each process would have its own loss_accum tensor, so to get the
# average loss_accum across all processes, we want to compute the
# average of all loss_accum in all processes
if ddp:
dist.all_reduce(loss_accum, op=dist.ReduceOp.AVG)
# Clips gradient to global norm. It is very useful to avoid having a
# very high loss for some (bad) batch(es) which would lead to very high
# gradients and huge updates
# NOTE: In the beginning of training it is normal to have high norms
# as the model is initialized randomly
norm = nn.utils.clip_grad_norm_(model.parameters(), 1.0)
# TODO: Use ParamScheduler from `cmn_ai`
lr = sched(step / max_steps)
for pg in optimizer.param_groups:
pg["lr"] = lr
optimizer.step()
# CUDA launches kernels asynchronously, so without this barrier we would
# be timing how long it takes to QUEUE the work, not to finish it
if device_type == "cuda":
torch.cuda.synchronize()
end = time.time()
elapsed_time = end - start
token_per_sec = (
GPTConfig.batch_sz
* GPTConfig.block_sz
* grad_accum_steps
* ddp_world_size
) / (elapsed_time)
print(
f"step {step}, loss: {loss_accum.item():.4f}, lr {lr:.4e}, norm: {norm:.2f}, time: {elapsed_time:.2f}s, tok/sec: {token_per_sec:.2f}"
)
if ddp:
# Kills all processes
destroy_process_group()combine_scheds runs a linear warmup over the first 20% of steps, then hands off to cosine decay for the remaining 80%.
lin_sched/cos_sched/combine_scheds code defined above. Linear warmup from 3e-5 to the 3e-4 peak over 10 of 50 steps, then cosine decay back to a tenth of peak.Warmup exists because AdamW divides each update by a running estimate of the gradient’s second moment. For the first few steps that denominator is built from almost no samples, so a full-size learning rate applies a large update scaled by a quantity nobody has measured yet. Starting small buys time for the estimate to stabilize. The cosine decay at the other end shrinks the step size as the model approaches a minimum, so the last updates are small enough to settle into it.
With training finished, we can sample from the model. Generation reuses the same forward pass, run one token at a time, with each sampled token appended to the context and fed back in.
# Sample from the trained model. `generate` crops the context to block_sz,
# takes the logits at the final position, and samples from the top 50 tokens
# so we never draw from the long tail of near-zero probabilities.
enc = tiktoken.get_encoding("gpt2")
prompt = torch.tensor(enc.encode("Hello, I'm a language model,"), dtype=torch.long)
# Generate 4 independent continuations from the same prompt
prompt = prompt.unsqueeze(0).repeat(4, 1).to(device)
raw_model.eval()
out = raw_model.generate(prompt, max_tokens=30)
for row in out:
print(enc.decode(row.tolist()))
print("-" * 40)We’ve come a long way in this journey, from implementing the core transformer architecture with multi-head attention and positional encodings, to building an efficient training pipeline complete with modern optimizations like flash attention, mixed precision training, and distributed parallelism. We’ve debugged exploding gradients, optimized memory usage, and watched our model evolve from producing random gibberish to generating coherent text.
If you want to go deeper on the optimization side, see my DL Tips & Tricks reference, which covers many of the heuristics used here.
Along the way, we’ve gained deep insights into why each component exists and how they work together to create these remarkable language models.
I hope this deep dive has been as illuminating for you as it has been for me. Writing this implementation forced me to confront gaps in my own understanding and solidified concepts that previously felt abstract. There’s something uniquely satisfying about seeing your hand-built transformer successfully predict its first coherent sentence: a moment where theory becomes understanding.
If you’ve made it this far, thank you for joining me on this journey. I’d love to hear about your experiences implementing transformers, any bugs you’ve encountered, optimizations you’ve discovered, or questions this post might have raised. Feel free to reach out with feedback, corrections, or insights: the best part of sharing these implementations is learning from the community’s collective wisdom. Happy building!
The loop above trains on tinyshakespeare for 50 steps, which is enough to watch the loss fall and nothing more. To turn it into a real run:
.npy files, so the loader memory-maps shards instead of holding the corpus in RAM.torch.compile, then re-measure. Take the throughput number after the first step, since that one pays for compilation.