Build GPT from Scratch: A Production-Grained Walkthrough

Building a GPT from scratch is a core skill for any ML engineer working with language models — not just an academic exercise. Understanding the internals lets you fine-tune, debug, and scale models beyond what off-the-shelf APIs offer. The canonical nanoGPT repo by Andrej Karpathy provides a clean, minimal implementation that reproduces GPT-2 (124M) on a single node, making it the perfect starting point for serious developers.

This article walks through every component of a GPT: tokenization, embedding, multi-head self-attention, feed-forward networks, layer normalization, and the training loop. We'll reference nanoGPT's ~300-line model.py and train.py, but go deeper into production considerations like memory profiling, gradient accumulation, and checkpointing strategies.

You'll learn not just how to code a GPT, but how to debug it when training diverges, optimize throughput on GPUs, and avoid common pitfalls that waste compute. By the end, you'll have a mental model that scales from a character-level Shakespeare model to a 1.3B parameter GPT-2.

This is not a beginner tutorial. You need working knowledge of PyTorch, transformers, and basic deep learning. We'll assume you've read the Attention Is All You Need paper and understand backpropagation. If not, start with Karpathy's 'Neural Networks: Zero to Hero' series first.

Introduction: Why Build GPT from Scratch

By 2026, GPTs are commodity infrastructure. You don't build one to beat OpenAI — you build one to own your stack, control your data, and ship models that fit in a single GPU for under $500. The era of 'just call the API' is over for production systems that need predictable latency, zero data leakage, and custom tokenization for domain-specific corpora like legal documents, medical records, or codebases with proprietary syntax. Building from scratch gives you surgical control over every parameter, from embedding dimension to attention head count, and lets you deploy on edge devices or air-gapped environments where no cloud API reaches.

This walkthrough implements a decoder-only transformer that mirrors GPT-2's architecture at 124M parameters — the smallest viable model that exhibits coherent long-range dependencies. We use PyTorch 2.x with compile, Flash Attention kernels, and the tiktoken BPE tokenizer. The final model trains on OpenWebText in under 4 days on a single 8x A100 node, reproducing GPT-2's loss curve. But more importantly, you'll understand every line of code: tokenization, embeddings, causal self-attention, feedforward blocks, layer normalization, weight tying, and the training loop with cosine decay and gradient clipping.

Why 2026 specifically? Because hardware has shifted: consumer GPUs now have 24-48GB VRAM, Flash Attention is standard in PyTorch, and quantization (FP8, INT4) is trivial. The barrier to training a 124M model from scratch is a weekend project on a single RTX 4090. The knowledge you gain transfers directly to scaling laws, mixture-of-experts, and multi-modal architectures. If you can build GPT from scratch, you can debug any transformer-based system in production.

This is not a tutorial for beginners. You need working knowledge of PyTorch, backpropagation, and basic NLP. We skip the 'what is attention' hand-waving and go straight to tensor shapes, masking logic, and numerical stability. Every code block is runnable and tested against PyTorch 2.5. Let's build.

import torch, tiktoken, sys
print(f'PyTorch {torch.__version__}, CUDA {torch.version.cuda}')
print(f'tiktoken {tiktoken.__version__}')
enc = tiktoken.get_encoding('gpt2')
tokens = enc.encode('Hello, world!')
print(f'Tokens: {tokens}')
assert torch.cuda.is_available(), 'Need CUDA for training'
print('Setup OK')

Embedding Layer: Token and Positional Embeddings

The embedding layer converts token IDs (integers) into dense vectors. GPT-2 uses a learned token embedding matrix of shape (vocab_size, n_embd) where n_embd = 768 for the 124M model. Each token ID indexes a row in this matrix, producing a vector of size 768. This is a simple lookup: no computation, just memory access. The embedding matrix is shared with the output projection layer (weight tying) to reduce parameters and improve training stability.

Positional embeddings are also learned, not sinusoidal. GPT-2 uses a separate learned embedding of shape (block_size, n_embd) where block_size = 1024. The position index (0 to 1023) is added to the token embedding element-wise. This gives the model a sense of order without any inductive bias. The sum of token and positional embeddings is then passed through layer normalization before the first transformer block.

In code, we implement a combined Embedding module that stores both token and position embeddings. The forward pass takes token IDs of shape (batch, seq_len) and returns embeddings of shape (batch, seq_len, n_embd). We use PyTorch's nn.Embedding with padding_idx=None (no padding token in GPT-2). The position indices are generated on the fly as torch.arange(seq_len, device=x.device).

Weight tying is implemented by setting the output linear layer's weight equal to the token embedding weight. This is done after model initialization: model.lm_head.weight = model.transformer.wte.weight. This halves the embedding parameter count and empirically improves convergence. For the 124M model, the embedding layer accounts for 50,257 * 768 ≈ 38.6M parameters, about 31% of total parameters.

import torch
import torch.nn as nn

class GPTEmbeddings(nn.Module):
    def __init__(self, vocab_size: int, n_embd: int, block_size: int):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        self.position_embedding = nn.Embedding(block_size, n_embd)
        self.block_size = block_size

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        batch, seq_len = input_ids.shape
        assert seq_len <= self.block_size, f'Sequence length {seq_len} exceeds block size {self.block_size}'
        positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch, -1)
        token_emb = self.token_embedding(input_ids)  # (batch, seq, n_embd)
        pos_emb = self.position_embedding(positions)  # (batch, seq, n_embd)
        return token_emb + pos_emb

# Example usage
vocab_size = 50257
n_embd = 768
block_size = 1024
emb = GPTEmbeddings(vocab_size, n_embd, block_size)
x = torch.randint(0, vocab_size, (2, 128))  # batch=2, seq=128
out = emb(x)
print(f'Input shape: {x.shape}, Output shape: {out.shape}')
print(f'Token embedding params: {sum(p.numel() for p in emb.token_embedding.parameters())}')

Multi-Head Self-Attention: Implementation with Causal Masking

Multi-head self-attention is the core of the transformer. For each head, we compute queries (Q), keys (K), and values (V) from the input via learned linear projections. The attention scores are Q @ K^T / sqrt(d_k) where d_k = n_embd / n_head. For GPT-2 124M, n_embd=768 and n_head=12, so d_k=64. The scores are masked to prevent attending to future tokens (causal masking) by setting upper-triangular entries to -inf before softmax. The softmax output is then multiplied by V to produce the head's output. All heads are concatenated and projected back to n_embd.

Causal masking is implemented as a boolean mask of shape (1, 1, seq_len, seq_len) where mask[i,j] = 0 if i >= j else -inf. We add this mask to the attention scores before softmax. In practice, we use torch.triu with diagonal=1 to create the mask. For efficiency, we use Flash Attention (torch.nn.functional.scaled_dot_product_attention with is_causal=True) which fuses the QKV projections, masking, and softmax into a single kernel, reducing memory from O(n^2) to O(n).

The attention mechanism has O(n^2 * d_k) complexity per head, where n is sequence length. For GPT-2's block_size=1024, this is manageable. But for longer sequences (e.g., 8k), Flash Attention is essential. Our implementation falls back to manual attention for clarity but includes a flag to use Flash Attention when available.

In code, we implement a single attention head as a module, then combine multiple heads in MultiHeadAttention. The forward pass: (1) project input to Q, K, V for all heads simultaneously using a single linear layer, (2) split into heads, (3) compute attention with causal mask, (4) concatenate heads, (5) final projection. We include dropout on attention weights and output for regularization.

import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd: int, n_head: int, dropout: float = 0.1):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.n_embd = n_embd
        self.head_dim = n_embd // n_head
        
        self.c_attn = nn.Linear(n_embd, 3 * n_embd)  # Q, K, V projections
        self.c_proj = nn.Linear(n_embd, n_embd)      # output projection
        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)
        
        # Causal mask: (1, 1, block_size, block_size)
        self.register_buffer('mask', torch.triu(torch.full((1024, 1024), float('-inf')), diagonal=1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape  # batch, seq_len, n_embd
        qkv = self.c_attn(x)  # (B, T, 3*C)
        q, k, v = qkv.chunk(3, dim=-1)  # each (B, T, C)
        
        # Reshape for multi-head: (B, n_head, T, head_dim)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        
        # Flash Attention (preferred) or manual
        if hasattr(F, 'scaled_dot_product_attention'):
            y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.attn_dropout.p if self.training else 0.0, is_causal=True)
        else:
            att = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)  # (B, n_head, T, T)
            att = att + self.mask[:, :, :T, :T]  # causal mask
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = att @ v  # (B, n_head, T, head_dim)
        
        y = y.transpose(1, 2).contiguous().view(B, T, C)  # concatenate heads
        y = self.resid_dropout(self.c_proj(y))
        return y

# Example
attn = CausalSelfAttention(n_embd=768, n_head=12, dropout=0.1)
x = torch.randn(2, 128, 768)
out = attn(x)
print(f'Input: {x.shape}, Output: {out.shape}')
print(f'Parameters: {sum(p.numel() for p in attn.parameters())}')

Transformer Block: Attention, Feed-Forward, Layer Norm, and Residuals

The transformer block is the fundamental building unit of GPT. Each block consists of two sub-layers: multi-head causal self-attention and a position-wise feed-forward network (FFN). Both sub-layers are wrapped with residual connections and preceded by layer normalization (pre-norm). The pre-norm formulation, where LayerNorm is applied before the sub-layer rather than after, has become standard in GPT-style models because it stabilizes training at depth. The residual path allows gradients to flow directly through the stack, mitigating vanishing gradient problems even with 12, 24, or 48 blocks.

Multi-head attention splits the embedding dimension into h heads, each of dimension d_k = d_model / h. For each head, we compute queries Q, keys K, and values V via learned linear projections. The attention scores are computed as softmax(QK^T / sqrt(d_k) + M), where M is a causal mask that sets all future positions to -inf. This ensures position i can only attend to positions j ≤ i. The mask is typically implemented as a lower-triangular matrix filled with 0s in the lower triangle and -inf in the upper triangle. After computing attention, the heads are concatenated and projected back to d_model.

The feed-forward network is a simple two-layer MLP with a GELU activation in between. The typical GPT-2 configuration uses an inner dimension of 4 d_model. For d_model=768, the FFN expands to 3072 and then projects back to 768. This expansion-contraction pattern allows the model to learn complex non-linear transformations. The GELU activation is approximated as 0.5 x (1 + tanh(sqrt(2/pi) (x + 0.044715 * x^3))), though modern implementations often use the exact erf-based version.

Residual connections are critical: each sub-layer's output is added to its input. If we denote the input to a block as x, the output is x + Attention(LayerNorm(x)) + FFN(LayerNorm(x + Attention(LayerNorm(x)))). This additive structure means the model can learn to ignore sub-layers by learning near-zero weights, effectively reducing depth if needed. In practice, we initialize the output projection of each sub-layer with a small weight (e.g., N(0, 0.02)) and often use a scaling factor of 1/sqrt(2 * num_layers) to keep activations stable.

Layer normalization computes mean and variance across the feature dimension (not the sequence dimension). For an input x of shape (batch, seq_len, d_model), LayerNorm computes μ = mean(x, dim=-1) and σ = std(x, dim=-1), then outputs γ * (x - μ) / (σ + ε) + β, where γ and β are learnable parameters of size d_model. The epsilon (typically 1e-5) prevents division by zero. This normalization is crucial for training stability, especially when using FP16 or BF16 mixed precision.

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_dim = config.n_embd // config.n_head
        
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.c_proj.NANOGPT_SCALE_INIT = 1  # flag for scaled init
        
        # causal mask: (1, 1, block_size, block_size)
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(config.block_size, config.block_size))
            .view(1, 1, config.block_size, config.block_size)
        )

    def forward(self, x):
        B, T, C = x.shape  # batch, seq_len, embedding_dim
        qkv = self.c_attn(x)  # (B, T, 3*C)
        q, k, v = qkv.split(self.n_embd, dim=2)
        
        # reshape to (B, n_head, T, head_dim)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        
        # scaled dot-product attention with causal mask
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        y = att @ v  # (B, n_head, T, head_dim)
        
        # reassemble all head outputs
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.c_proj(y)
        return y

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.c_proj.NANOGPT_SCALE_INIT = 1

    def forward(self, x):
        x = self.c_fc(x)
        x = F.gelu(x, approximate='tanh')  # exact GELU is default, tanh approx for GPT-2 compat
        x = self.c_proj(x)
        return x

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

The Full GPT Model: Stacking Blocks and the Language Modeling Head

The full GPT model is a stack of N transformer blocks (typically 12 for GPT-2 small, 24 for medium, 36 for large, 48 for XL) followed by a language modeling head. The input pipeline starts with token embeddings and position embeddings, which are summed to produce the initial hidden state. There is no segment embedding (unlike BERT) because GPT is a unidirectional decoder-only model. The token embeddings are a learned lookup table of size vocab_size × n_embd, and the position embeddings are a learned lookup table of size block_size × n_embd.

After the embedding layer, the hidden state passes through the stack of transformer blocks. Each block maintains the same dimensionality (n_embd) throughout. After the final block, a layer normalization is applied, followed by a linear projection (the LM head) that maps from n_embd to vocab_size. This produces logits of shape (batch, seq_len, vocab_size). During training, we compute cross-entropy loss between these logits and the target tokens (shifted by one position). During inference, we sample from the logits to generate the next token.

The weight tying trick, popularized by the original Transformer paper and used in GPT-2, shares the weight matrix between the token embedding layer and the LM head. This reduces the number of parameters by vocab_size × n_embd (e.g., ~38M for GPT-2 small with vocab_size=50257 and n_embd=768). The shared weights are typically scaled by sqrt(n_embd) in the embedding layer to keep the variance of the summed embeddings consistent with the residual stream.

The model also includes dropout layers for regularization. In the original GPT-2, dropout is applied to the embedding layer (with rate 0.1) and to the output of each attention sub-layer (also 0.1). However, many modern implementations (including nanoGPT) set dropout to 0 during pretraining on large datasets, as the regularization from large-scale data is sufficient. Dropout is more commonly used during fine-tuning on smaller datasets.

For the GPT-2 124M parameter configuration, the architecture is: vocab_size=50257, block_size=1024, n_embd=768, n_head=12, n_layer=12, bias=True. The total parameter count is approximately 124M, which includes embeddings (50257 × 768 ≈ 38.6M), transformer blocks (12 × (attention: 4 × 768² + MLP: 2 × 768 × 3072) ≈ 85M), and layer norms (12 × 2 × 768 × 2 ≈ 36K). The bias terms add a small fraction.

import torch
import torch.nn as nn
from .transformer_block import TransformerBlock

class GPTConfig:
    def __init__(self, vocab_size=50257, block_size=1024, n_layer=12, n_head=12, n_embd=768, bias=True, dropout=0.0):
        self.vocab_size = vocab_size
        self.block_size = block_size
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_embd = n_embd
        self.bias = bias
        self.dropout = dropout

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),
            drop=nn.Dropout(config.dropout),
            h=nn.ModuleList([TransformerBlock(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 tying: share weights between token embedding and LM head
        self.transformer.wte.weight = self.lm_head.weight
        
        # initialize weights
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            std = 0.02
            if hasattr(module, 'NANOGPT_SCALE_INIT'):
                std *= (2 * self.config.n_layer) ** -0.5
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        assert T <= self.config.block_size, f"Sequence length {T} exceeds block size {self.config.block_size}"
        
        # token + position embeddings
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device).unsqueeze(0)  # (1, T)
        tok_emb = self.transformer.wte(idx)  # (B, T, n_embd)
        pos_emb = self.transformer.wpe(pos)  # (1, T, n_embd)
        x = self.transformer.drop(tok_emb + pos_emb)
        
        # forward through transformer blocks
        for block in self.transformer.h:
            x = block(x)
        
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)  # (B, T, vocab_size)
        
        loss = None
        if targets is not None:
            loss = nn.functional.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1
            )
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            # crop context to block_size
            idx_cond = idx[:, -self.config.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')
            
            probs = nn.functional.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

Training Loop: Data Loading, Loss Computation, Backpropagation, and Optimization

The training loop for GPT follows the standard autoregressive language modeling setup. Data is preprocessed into a flat array of token IDs (typically using tiktoken for BPE or a simple character-level encoding for small experiments). The data loader samples random contiguous chunks of length block_size from this array, creating input-target pairs where the target is the input shifted by one position. For example, if the input sequence is [t0, t1, ..., t_{n-1}], the target is [t1, t2, ..., t_n]. This is implemented efficiently by memory-mapping the token array and using random offsets to avoid loading the entire dataset into memory.

Loss computation uses cross-entropy between the predicted logits and the target tokens. The loss is averaged over all non-padding tokens (padding tokens are masked with ignore_index=-1 in the cross-entropy function). For a batch of B sequences of length T, the loss is: L = -1/(BT) Σ_b Σ_t log P(t_{b,t+1} | t_{b,≤t}). This is equivalent to the negative log-likelihood of the next token given all previous tokens. The perplexity, often reported as a metric, is exp(L).

Backpropagation computes gradients of the loss with respect to all model parameters. The AdamW optimizer is the standard choice for training GPTs. AdamW decouples weight decay from the adaptive learning rate, applying L2 regularization only to the weights (not biases or layer norms). The typical configuration is: learning_rate=3e-4 for 124M model, β1=0.9, β2=0.95, weight_decay=0.1, and epsilon=1e-8. A cosine learning rate schedule with linear warmup is used: the learning rate linearly increases from 0 to max_lr over the first few thousand steps (e.g., 2000), then follows a cosine decay to a minimum value (typically 10% of max_lr).

Gradient clipping is essential to prevent exploding gradients. The typical threshold is max_grad_norm=1.0. After computing gradients via loss.backward(), we call torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm). This scales down gradients whose L2 norm exceeds the threshold, ensuring stable training. Without clipping, a single outlier batch can destabilize the entire training run.

The evaluation loop runs periodically (e.g., every 1000 iterations) on a held-out validation set. It computes the loss without gradient computation (torch.no_grad()) and reports validation loss/perplexity. This is used for checkpoint selection: we save the model whenever validation loss improves. The training loop also logs metrics (loss, learning rate, gradient norm) to a dashboard like Weights & Biases or TensorBoard for monitoring.

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
import math
import os

class TokenDataset(Dataset):
    def __init__(self, data_path, block_size):
        self.data = np.memmap(data_path, dtype=np.uint16, mode='r')
        self.block_size = block_size

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

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

def get_batch(dataloader):
    return next(iter(dataloader))

def train_step(model, optimizer, scheduler, scaler, x, y, max_grad_norm=1.0):
    model.train()
    optimizer.zero_grad()
    
    with torch.amp.autocast(device_type='cuda', dtype=torch.float16):
        logits, loss = model(x, targets=y)
    
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
    scaler.step(optimizer)
    scaler.update()
    scheduler.step()
    
    return loss.item()

@torch.no_grad()
def evaluate(model, dataloader, num_batches=100):
    model.eval()
    total_loss = 0.0
    for i, (x, y) in enumerate(dataloader):
        if i >= num_batches:
            break
        x, y = x.cuda(), y.cuda()
        logits, loss = model(x, targets=y)
        total_loss += loss.item()
    return total_loss / min(num_batches, len(dataloader))

# Example training loop (simplified)
def train(model, train_loader, val_loader, config):
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config.learning_rate,
        betas=(0.9, 0.95),
        weight_decay=0.1
    )
    
    # Cosine schedule with warmup
    def lr_lambda(step):
        warmup_steps = config.warmup_steps
        if step < warmup_steps:
            return step / warmup_steps
        else:
            progress = (step - warmup_steps) / (config.max_steps - warmup_steps)
            return 0.5 * (1.0 + math.cos(math.pi * progress))
    
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
    scaler = torch.cuda.amp.GradScaler()
    
    best_val_loss = float('inf')
    for step in range(config.max_steps):
        x, y = get_batch(train_loader)
        x, y = x.cuda(), y.cuda()
        
        loss = train_step(model, optimizer, scheduler, scaler, x, y)
        
        if step % config.eval_interval == 0:
            val_loss = evaluate(model, val_loader)
            print(f"Step {step}: train_loss={loss:.4f}, val_loss={val_loss:.4f}, lr={scheduler.get_last_lr()[0]:.6f}")
            
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                torch.save(model.state_dict(), os.path.join(config.out_dir, 'best_model.pt'))

Production Considerations: Mixed Precision, Gradient Accumulation, Checkpointing, and Debugging

Mixed precision training (FP16 or BF16) is essential for training large GPT models efficiently. Modern GPUs (A100, H100) have dedicated Tensor Cores that provide 2-4x throughput for FP16/BF16 operations compared to FP32. The standard approach uses torch.cuda.amp (automatic mixed precision) with a GradScaler to prevent underflow in the loss during backpropagation. The scaler multiplies the loss by a scale factor before backward, then divides the gradients by the same factor after. If gradients overflow (become inf/nan), the scaler skips the step and reduces the scale. BF16 is preferred over FP16 when available because it has the same exponent range as FP32, eliminating the need for loss scaling in many cases.

Gradient accumulation allows training with effective batch sizes larger than what fits in GPU memory. Instead of computing the gradient over one large batch, we accumulate gradients over multiple micro-batches. For example, to achieve an effective batch size of 512 with micro-batch size 16, we accumulate gradients over 32 steps. The loss for each micro-batch is divided by the number of accumulation steps to keep the gradient magnitude consistent. This is implemented by calling loss.backward() on each micro-batch without zeroing gradients, then calling optimizer.step() after the accumulation is complete. Gradient accumulation is transparent to the optimizer: it sees the sum of gradients, which is equivalent to the gradient of the full batch.

Checkpointing strategy is critical for long training runs (days to weeks). Save checkpoints at regular intervals (e.g., every 1000 steps) and always keep the best model based on validation loss. A checkpoint should include: model state_dict, optimizer state_dict, scheduler state_dict, current step, and best validation loss. This allows resuming training from any checkpoint. Use a naming convention like 'ckpt_{step}_{val_loss:.4f}.pt' and implement a retention policy (e.g., keep last 5 checkpoints plus best). For distributed training, only save from rank 0 to avoid file corruption.

Debugging training issues requires systematic monitoring. Log these metrics every step: training loss, learning rate, gradient norm, and scale factor (for AMP). Watch for these red flags: (1) Loss not decreasing after 1000 steps → check learning rate, data loading, or model initialization. (2) Gradient norm suddenly spiking to >10 → reduce learning rate or increase gradient clipping. (3) Loss going to NaN → check for numerical instability in attention (use torch.nn.functional.scaled_dot_product_attention which is numerically stable), or reduce learning rate. (4) Validation loss diverging from training loss → overfitting; increase dropout or reduce model size.

For distributed training across multiple GPUs, use PyTorch's DistributedDataParallel (DDP). The key is to split the batch across GPUs and synchronize gradients during backward. With gradient accumulation, each GPU processes micro-batches independently and gradients are synchronized only at the optimizer step. The effective batch size is micro_batch_size × gradient_accumulation_steps × num_gpus. For example, with micro_batch_size=8, accumulation_steps=4, and 8 GPUs, the effective batch size is 256. DDP adds communication overhead, but for models up to 1.5B parameters on 8 GPUs, the overhead is negligible compared to compute time.

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import os
import math

def train_production(model, train_loader, val_loader, config):
    # Setup DDP
    local_rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend='nccl')
    model = model.cuda()
    model = DDP(model, device_ids=[local_rank])
    
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config.learning_rate,
        betas=(0.9, 0.95),
        weight_decay=0.1
    )
    
    # Cosine schedule with warmup
    def lr_lambda(step):
        warmup_steps = config.warmup_steps
        if step < warmup_steps:
            return step / warmup_steps
        else:
            progress = (step - warmup_steps) / (config.max_steps - warmup_steps)
            return 0.5 * (1.0 + math.cos(math.pi * progress))
    
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
    scaler = torch.cuda.amp.GradScaler()
    
    # Gradient accumulation
    accumulation_steps = config.effective_batch_size // (config.micro_batch_size * dist.get_world_size())
    
    best_val_loss = float('inf')
    model.train()
    optimizer.zero_grad()
    
    for step in range(config.max_steps):
        for micro_step in range(accumulation_steps):
            x, y = get_batch(train_loader)
            x, y = x.cuda(), y.cuda()
            
            with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
                logits, loss = model(x, targets=y)
                loss = loss / accumulation_steps  # normalize for accumulation
            
            scaler.scale(loss).backward()
        
        # Gradient clipping (unscale first)
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
        
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()
        optimizer.zero_grad()
        
        # Logging (rank 0 only)
        if local_rank == 0 and step % config.log_interval == 0:
            print(f"Step {step}: loss={loss.item() * accumulation_steps:.4f}, lr={scheduler.get_last_lr()[0]:.6f}")
        
        # Evaluation and checkpointing (rank 0 only)
        if local_rank == 0 and step % config.eval_interval == 0:
            val_loss = evaluate(model.module, val_loader)
            print(f"Validation loss: {val_loss:.4f}")
            
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                checkpoint = {
                    'model': model.module.state_dict(),
                    'optimizer': optimizer.state_dict(),
                    'scheduler': scheduler.state_dict(),
                    'step': step,
                    'val_loss': val_loss,
                    'config': config
                }
                torch.save(checkpoint, os.path.join(config.out_dir, 'best_model.pt'))
    
    dist.destroy_process_group()