Building a Transformer Language Model in PyTorch

This article connects a decoder-only Transformer’s embeddings, rotary positions, causal blocks, next-token loss, and sampling loop. Run the Python blocks in order. The small character-level example supplies its own text and vocabulary so the model can actually be trained and sampled without downloading a tokenizer. It provides a small implementation check.

Configuration

Use PyTorch 2.0 or newer for the attention API below. Input IDs have shape (N,T): N sequences, each with T tokens. The embedding turns them into (N,T,d), where d=n_embd; n_head splits that width into attention heads, n_layer counts stacked blocks, and vocab_size=V sets the final (N,T,V) logits. block_size is the maximum context processed in one call. The runnable example later overrides the defaults with a much smaller model.

import math
from dataclasses import dataclass

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

@dataclass
class GPTConfig:
    vocab_size: int = 50257
    block_size: int = 256        # context length
    n_layer: int = 6
    n_head: int = 6
    n_embd: int = 384
    dropout: float = 0.1

    def __post_init__(self):
        if min(self.vocab_size, self.block_size, self.n_layer, self.n_head, self.n_embd) < 1:
            raise ValueError("configuration sizes must be positive")
        if self.n_embd % self.n_head or (self.n_embd // self.n_head) % 2:
            raise ValueError("model width must divide into heads of positive even width")
        if not 0 <= self.dropout < 1:
            raise ValueError("dropout must be in [0,1)")

Self-attention, masking, and the full block are derived in Transformers from Scratch: Self-Attention to Encoder–Decoder.

Rotary position embedding

RoPE rotates coordinate pairs of queries and keys, leaving values unchanged. For fixed content vectors, \((R_mq)^T(R_nk)=q^TR_{n-m}k\), so the explicit rotation factor depends on relative position. The score still depends on q and k themselves. Adjacent-coordinate pairing here requires an even head width; matching another checkpoint also requires matching its pairing convention, rotation base, and scaling settings.

def build_rope_cache(seq_len, head_dim, base=10000.0, device=None):
    if seq_len < 1 or head_dim < 2 or head_dim % 2 or base <= 0:
        raise ValueError("positive sequence length, even head width, and positive base required")
    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2,
                                            device=device).float() / head_dim))
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(t, inv_freq)              # (T, head_dim/2)
    return freqs.cos(), freqs.sin()

def apply_rope(x, cos, sin):
    """x: (N, h, T, head_dim)"""
    x1, x2 = x[..., 0::2], x[..., 1::2]
    cos, sin = cos[None, None, :, :], sin[None, None, :, :]
    out = torch.stack([x1 * cos - x2 * sin,
                       x1 * sin + x2 * cos], dim=-1)
    return out.flatten(-2).to(x.dtype)

The cache depends on the position range, head width, rotation base, and numerical settings. This model precomputes it for the configured context length as nonpersistent buffers; loading weights requires reconstructing the same configuration. Apply rotations at the intended positions, then return to the attention tensor’s dtype. Computing angles in float32 avoids unnecessarily coarse trigonometric values during mixed-precision execution.

Causal self-attention

class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.n_head = cfg.n_head
        self.head_dim = cfg.n_embd // cfg.n_head
        self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=False)
        self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
        self.resid_drop = nn.Dropout(cfg.dropout)
        self.dropout = cfg.dropout

    def forward(self, x, cos, sin):
        N, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        q, k, v = (t.view(N, T, self.n_head, self.head_dim).transpose(1, 2)
                   for t in (q, k, v))

        q = apply_rope(q, cos[:T], sin[:T])
        k = apply_rope(k, cos[:T], sin[:T])

        y = F.scaled_dot_product_attention(
            q, k, v, is_causal=True,
            dropout_p=self.dropout if self.training else 0.0)

        y = y.transpose(1, 2).contiguous().view(N, T, C)
        return self.resid_drop(self.proj(y))

The scaled-dot-product attention API chooses an available backend based on device, dtype, shape, and other constraints. Eligible fused kernels can avoid storing the full attention matrix; CPU or unsupported cases may use another implementation. Check the selected backend and measure performance instead of assuming this call always runs FlashAttention. The example uses full equal-length token windows, so it needs only a causal mask, not a padding mask.

dropout_p is conditioned on self.training. The functional API does not consult the module’s mode, so forgetting this applies dropout during evaluation.

Block and model

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.n_embd)
        self.attn = CausalSelfAttention(cfg)
        self.ln2 = nn.LayerNorm(cfg.n_embd)
        self.mlp = nn.Sequential(
            nn.Linear(cfg.n_embd, 4 * cfg.n_embd, bias=False),
            nn.GELU(),
            nn.Linear(4 * cfg.n_embd, cfg.n_embd, bias=False),
            nn.Dropout(cfg.dropout))

    def forward(self, x, cos, sin):
        x = x + self.attn(self.ln1(x), cos, sin)
        x = x + self.mlp(self.ln2(x))
        return x


class GPT(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.ln_f = nn.LayerNorm(cfg.n_embd)
        self.head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)


        cos, sin = build_rope_cache(cfg.block_size, cfg.n_embd // cfg.n_head)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

        self.apply(self._init_weights)
        self.head.weight = self.tok_emb.weight          # tie after initialization
        for name, p in self.named_parameters():
            if name.endswith("mlp.2.weight") or name.endswith("proj.weight"):
                nn.init.normal_(p, mean=0.0,
                                std=0.02 / math.sqrt(2 * cfg.n_layer))

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

    def forward(self, idx, targets=None):
        if idx.ndim != 2 or idx.dtype != torch.long or idx.size(0) == 0:
            raise ValueError("idx must be a nonempty batch of integer token IDs")
        T = idx.size(1)
        if not 1 <= T <= self.cfg.block_size:
            raise ValueError("sequence length is outside the configured context")
        x = self.drop(self.tok_emb(idx))
        for block in self.blocks:
            x = block(x, self.cos, self.sin)
        logits = self.head(self.ln_f(x))

        loss = None
        if targets is not None:
            if targets.shape != idx.shape or targets.dtype != torch.long:
                raise ValueError("targets must have the same shape and integer dtype as idx")
            if not (targets != -1).any():
                raise ValueError("at least one target must contribute to the loss")
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
                                   targets.reshape(-1), ignore_index=-1)
        return logits, loss

Weight tying shares a \(V\times d\) matrix: an input ID selects one row, and each output logit is the final hidden vector’s dot product with a row. Sharing saves one such matrix compared with an otherwise identical untied model. At the default settings this saves 19.3 million parameters; the tied model has about 29.9 million. Tie after initialization so traversal does not initialize the shared tensor twice. Quality effects depend on the model and data.

Scaled residual initialization reduces the standard deviation of the two residual-output projections per block by \(\sqrt{2L}\), where L is the number of blocks. This controls the initial size of the added branches. Linear growth of variance would require assumptions about their sizes and correlations; instability does not follow from depth alone. The final LayerNorm normalizes the stream before the vocabulary projection.

Training

def get_batch(data, block_size, batch_size, device):
    if data.ndim != 1 or data.dtype != torch.long or data.device.type != "cpu":
        raise ValueError("data must be a 1-D CPU tensor of integer token IDs")
    if block_size < 1 or batch_size < 1 or len(data) < block_size + 1:
        raise ValueError("need positive batch/context sizes and at least context+1 tokens")
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i + block_size] for i in ix])
    y = torch.stack([data[i + 1:i + 1 + block_size] for i in ix])
    return x.to(device), y.to(device)          # y is x shifted by one


def configure_optimizer(model, lr, weight_decay=0.1):
    decay = [p for n, p in model.named_parameters()
             if p.dim() >= 2 and p.requires_grad]
    no_decay = [p for n, p in model.named_parameters()
                if p.dim() < 2 and p.requires_grad]
    return torch.optim.AdamW(
        [{"params": decay, "weight_decay": weight_decay},
         {"params": no_decay, "weight_decay": 0.0}],
        lr=lr, betas=(0.9, 0.95))


def train(model, data, steps=2000, batch_size=32, lr=3e-4, device="cpu"):
    if steps < 1 or lr <= 0:
        raise ValueError("steps and learning rate must be positive")
    model.to(device).train()
    opt = configure_optimizer(model, lr)
    warmup = max(1, steps // 50)

    for step in range(steps):
        scale = ((step + 1) / warmup if step < warmup
                 else 0.5 * (1 + math.cos(math.pi * (step - warmup)
                                          / max(1, steps - warmup))))
        for g in opt.param_groups:
            g["lr"] = lr * max(scale, 0.05)

        x, y = get_batch(data, model.cfg.block_size, batch_size, device)
        _, loss = model(x, y)

        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()

        if step % 200 == 0:
            print(f"step {step}  loss {loss.item():.4f}  "
                  f"minibatch_ppl {math.exp(loss.item()):.1f}")

For tokens [2, 5, 7, 1] and context length 3, the batch is x=[2, 5, 7] and y=[5, 7, 1]. Logit position t uses input positions up to and including t to predict token t+1. The diagonal of the causal mask is therefore allowed. The model flattens \((N,T,V)\) logits and \((N,T)\) targets for cross-entropy; it does not shift them again. A target of -1 excludes that position from the loss, but does not mask an input from attention.

Uniform predictions have cross-entropy \(\log V\), about 10.8 for V=50,257. This is a reference value, not a required initialization result: random tied embeddings need not produce uniform probabilities. The logged perplexity is the exponential of one training minibatch’s mean loss, not held-out corpus perplexity. The learning rate warms up and then follows a cosine schedule with a floor at 5% of the requested rate. AdamW decays matrix parameters, including the shared embedding, once; LayerNorm vectors are excluded.

Sampling

def filter_logits(logits, top_k=None, top_p=None):
    V = logits.size(-1)
    if top_k is not None and (not isinstance(top_k, int) or not 1 <= top_k <= V):
        raise ValueError("top_k must be an integer between 1 and vocabulary size")
    if top_p is not None and not 0 < top_p <= 1:
        raise ValueError("top_p must be in (0,1]")
    if top_k is not None:
        values, indices = logits.topk(top_k, dim=-1)
        logits = torch.full_like(logits, -float("inf")).scatter(-1, indices, values)
    if top_p is not None and top_p < 1:
        values, indices = logits.sort(dim=-1, descending=True)
        cumulative = values.softmax(dim=-1).cumsum(dim=-1)
        remove = torch.zeros_like(cumulative, dtype=torch.bool)
        remove[..., 1:] = cumulative[..., :-1] >= top_p
        values = values.masked_fill(remove, -float("inf"))
        logits = torch.full_like(logits, -float("inf")).scatter(-1, indices, values)
    return logits


@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None):
    if idx.ndim != 2 or min(idx.shape) < 1 or idx.dtype != torch.long:
        raise ValueError("provide a nonempty batch of token-ID prompts")
    if not isinstance(max_new_tokens, int) or max_new_tokens < 0:
        raise ValueError("max_new_tokens must be a nonnegative integer")
    if not math.isfinite(temperature) or temperature < 0:
        raise ValueError("temperature must be finite and nonnegative")
    filter_logits(torch.zeros(1, model.cfg.vocab_size), top_k, top_p)
    was_training = model.training
    model.eval()
    try:
        for _ in range(max_new_tokens):
            logits, _ = model(idx[:, -model.cfg.block_size:])
            logits = logits[:, -1, :]
            if not torch.isfinite(logits).all():
                raise ValueError("model produced nonfinite logits")
            if temperature == 0:
                next_id = logits.argmax(dim=-1, keepdim=True)
            else:
                logits = filter_logits(logits / temperature, top_k, top_p)
                next_id = torch.multinomial(logits.softmax(dim=-1), 1)
            idx = torch.cat([idx, next_id], dim=1)
        return idx
    finally:
        model.train(was_training)

Temperature divides logits before softmax: a positive value below 1 concentrates probability on higher-scoring tokens, while a value above 1 spreads it out. Repetition and coherence also depend on the model and prompt. Here temperature=0 selects the largest logit directly and bypasses filtering. Top-k retains exactly k candidates (ties may be resolved arbitrarily); top-p retains the smallest sorted prefix reaching probability p. If both are supplied, top-p uses probabilities renormalized within the top-k set. This simple generator emits a fixed number of tokens and has no end-of-sequence stopping rule. Prompt tensors must be on the model’s device.

Run a small character model

The text below is deliberately repetitive so a short CPU run can exercise the whole path. Split the token stream before drawing windows, fit the vocabulary on training text, and keep validation windows inside the held-out suffix. The repeated phrases make this a mechanics check; held-out loss on this text does not assess general language ability. Real data also needs document or source separation where related passages could leak across the split.

torch.manual_seed(7)
text = "the cat sat. the dog ran.\n" * 80
cut = int(0.8 * len(text))
train_text, valid_text = text[:cut], text[cut:]
chars = sorted(set(train_text))
stoi = {ch: i for i, ch in enumerate(chars)}
encode = lambda s: torch.tensor([stoi[ch] for ch in s], dtype=torch.long)
decode = lambda ids: "".join(chars[i] for i in ids)
train_ids, valid_ids = encode(train_text), encode(valid_text)
cfg = GPTConfig(vocab_size=len(chars), block_size=16, n_layer=2,
                n_head=2, n_embd=32, dropout=0.0)
model = GPT(cfg)
train(model, train_ids, steps=40, batch_size=8, lr=0.003)

@torch.no_grad()
def validation_loss(model, ids):
    was_training = model.training
    model.eval()
    total, count = 0.0, 0
    B = model.cfg.block_size
    device = next(model.parameters()).device
    try:
        for start in range(0, len(ids) - 1, B):
            window = ids[start:start + B + 1].to(device)
            x, y = window[:-1][None], window[1:][None]
            _, loss = model(x, y)
            total += loss.item() * y.numel()
            count += y.numel()
        if count == 0:
            raise ValueError("validation needs at least two tokens")
        return total / count
    finally:
        model.train(was_training)

held_out = validation_loss(model, valid_ids)
print("validation loss", round(held_out, 3), "ppl", round(math.exp(held_out), 2))
prompt = encode("the ")[None]
print(decode(generate(model, prompt, 40, temperature=0)[0].tolist()))

On a CPU run with PyTorch 2.2.2, this 40-update example reached validation loss 1.246 (perplexity 3.48); its greedy continuation still combined fragments such as “rathe”. A decreasing loss does not mean the generated text is already fluent. The vocabulary maps characters to integers; decoding maps generated integers back to characters. Each validation window predicts every next token once within that window. Perplexity uses their token-weighted mean loss with dropout disabled. Exact losses and sampled text depend on the backend and random state. For a real tokenizer, save its vocabulary and normalization rules with the configuration and weights; unseen characters in this tiny dictionary raise an error. Keep model parameters and RoPE buffers in float32 when adding autocast; converting the whole model with half() also rounds those buffers.

KV caching

This loop recomputes the retained prefix at each step. Without a context cap, dense attention costs \(O(T^2)\) per step for a prefix of length T and \(O(T^3)\) across growing prefixes. Here the crop caps that attention work at \(O(B^2)\) per step, where B is block_size. KV caching reuses previous keys and values and makes the attention part linear in the retained context per new token; projections, the feed-forward layers, and the vocabulary head still have costs.

Reuse is valid for an unchanged prefix in evaluation mode. This generator instead rebuilds cropped contexts and restarts RoPE positions at zero. A uniform positional shift preserves RoPE’s relative rotation factors, but dropping old context can change the hidden states of retained tokens. Simply evicting old cache entries would not reproduce this recomputation. A cached implementation must choose its window and position policy explicitly.

The inference-time cost of this, and how to reduce it, is covered in LLM Inference Optimization: KV Cache, FlashAttention, and Quantization.

Bringing it up

1. Fit a fixed, small pattern with dropout off and inspect whether loss decreases. Conflicting next tokens for identical available contexts prevent zero loss; poor optimization can also obstruct fitting.
2. Compare the starting loss with the uniform-prediction reference \(\log V\).
3. Verify causality: change the last input token and confirm that the logits at earlier positions are unchanged.
4. Run the character example, inspect held-out loss, and read its samples before increasing model size.
5. Try compilation or a supported autocast dtype after checking correctness; measure speed and memory on your device.

The default configuration has about 29.9 million unique parameters, but training also needs memory for activations, gradients, and optimizer state. Larger language models often change normalization, feed-forward activations, attention head sharing, and other details. This implementation establishes the train-to-sample path on which those changes can be tested.

How text becomes the integers a model actually sees is covered in Tokenization Explained: BPE, WordPiece, and SentencePiece.

Exercises

1. Count the saving from weight tying. Use this article’s architecture with vocabulary 50,257, width 768, and 12 layers. Include LayerNorm scale and bias parameters; RoPE has no trainable parameters. Count tied and untied models and express the saving as a fraction of the untied total.

The output projection duplicates the embedding’s size only in the untied model. The Transformer blocks remain in both totals.

V, d, L = 50257, 768, 12
emb = V * d
body = L * (12 * d * d + 4 * d) + 2 * d
print("embedding", emb)
print("body", body)
print("tied", emb + body)
print("untied", 2 * emb + body)
print("saving fraction", round(emb / (2 * emb + body), 3))
# embedding 38597376
# body 84973056
# tied 123570432
# untied 162167808
# saving fraction 0.238

2. Flattening a sliced tensor. Take a batch of targets, slice off the first column, and call .view(-1) on the result. Report what happens and which method to use instead.

The printed exception type identifies the failed operation; inspect the skipped columns to explain why flattening cannot use the existing strides.

import torch
targets = torch.arange(12).reshape(3, 4)
shifted = targets[:, 1:]                 # demonstrate a sliced layout
print(shifted.is_contiguous())
try:
    shifted.view(-1)
except RuntimeError as e:
    print(type(e).__name__)
print(shifted.reshape(-1).tolist())
# False
# RuntimeError
# [1, 2, 3, 5, 6, 7, 9, 10, 11]

3. Pre-norm versus post-norm. Compare 24 independently parameterized residual feed-forward blocks, using identical starting weights and input in the two orderings. Record selected activation norms. Can forward norms alone tell you whether learning-rate warmup is necessary?

At initialization, post-norm’s LayerNorm constrains the output scale. Pre-norm does not normalize the residual sum at each block.

from copy import deepcopy
torch.manual_seed(0)
d, L = 32, 24
layers = nn.ModuleList([
    nn.ModuleDict({"ln": nn.LayerNorm(d),
                   "ff": nn.Sequential(nn.Linear(d, 4*d), nn.GELU(),
                                       nn.Linear(4*d, d))})
    for _ in range(L)])
x0 = torch.randn(8, d)
with torch.no_grad():
    for mode in ("pre", "post"):
        stack, x = deepcopy(layers), x0.clone()
        for depth, layer in enumerate(stack, 1):
            ln, ff = layer["ln"], layer["ff"]
            x = x + ff(ln(x)) if mode == "pre" else ln(x + ff(x))
            if depth in (1, 12, 24):
                print(mode, depth, round(x.norm(dim=-1).mean().item(), 2))
# pre 1 5.71
# pre 12 7.25
# pre 24 8.48
# post 1 5.66
# post 12 5.66
# post 24 5.66

References

• Press and Wolf (2017). Using the Output Embedding to Improve Language Models.
• Vaswani et al. (2017). Attention Is All You Need.
• Radford et al. (2019). Language Models Are Unsupervised Multitask Learners. OpenAI technical report.
• Holtzman, Buys, Du, Forbes, and Choi (2020). The Curious Case of Neural Text Degeneration. ICLR.
• Dao, Fu, Ermon, Rudra, and Re (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS.
• Xiong et al. (2020). On Layer Normalization in the Transformer Architecture.
• PyTorch: scaled_dot_product_attention, dropout behavior and backend selection.
• Su et al. (2024). RoFormer: Enhanced Transformer with Rotary Position Embedding. Neurocomputing.