Building Large Language Models from Scratch: Tokenization to Training

In 2022, Andrej Karpathy published nanoGPT — a from-scratch reimplementation of GPT-2 in roughly 300 lines of PyTorch, small enough to read in an afternoon and complete enough to train a working language model on a single GPU. It strips away the abstraction layers that hide what a transformer language model actually is: a stack of matrix multiplications with a specific shape, initialized with specific numbers, updated by a specific optimizer rule. A stream of token ids doesn't just appear, though — it's produced by a tokenizer that itself has to be built and trained, learning its vocabulary from the same kind of text the model will later be trained on. Once that stream of token ids exists, everything from there to a trained model is architecture code and a training loop. This chapter builds both: the byte-pair-encoding tokenizer that turns text into ids, and the transformer that turns those ids into a trained model — tracing its shapes and parameter counts by hand, and running the training loop step by step, so that "training a language model" stops being a phrase you use and becomes a sequence of operations you can predict the output of.

Byte-pair encoding: building the tokenizer's vocabulary

The tokenizer is not a fixed lookup table handed down from nowhere — it is trained, on a corpus of text, before the language model ever sees a token id. GPT-2 and everything downstream of it uses byte-pair encoding (BPE), introduced for subword tokenization by Sennrich, Haddow, and Birch ("Neural Machine Translation of Rare Words with Subword Units," ACL 2016) and adapted to operate on raw bytes rather than Unicode characters in GPT-2 itself. The algorithm starts from the smallest possible units — bytes — and greedily merges whichever adjacent pair occurs most often in the training corpus, one merge at a time, until the vocabulary reaches a target size. Frequent chunks of text — common words, common suffixes — end up as single tokens; rare or unseen strings fall back to shorter pieces, down to individual bytes in the worst case, so nothing is ever out-of-vocabulary.

Trace the first three merges by hand on a small corpus (word frequency in a training set): low ×5, lower ×2, newest ×6, widest ×3 — the same toy example Sennrich et al. use. Each word starts as a sequence of characters with an end-of-word marker </w>, so the corpus begins as:

Counting every adjacent symbol pair across the corpus, weighted by word frequency, the top pairs are (e,s): 6+3=9, (s,t): 6+3=9, (t,</w>): 6+3=9, (w,e): 2+6=8, (l,o): 5+2=7, (o,w): 5+2=7. Ties are broken by a fixed rule (e.g. first-seen); take (e,s) first:

• Merge 1: e+s → es (count 9). newest becomes n e w es t </w>, widest becomes w i d es t </w>.
• Merge 2: es+t → est (count 9, recomputed after merge 1 — (es,t) and (t,</w>) are now tied at 9, and (es,t) is taken first). Both words now end ... est </w>.
• Merge 3: est+</w> → est</w> (count 9). newest is now n e w est</w>, widest is w i d est</w> — the suffix "est" plus word boundary has become one indivisible token.

Repeating this — recount all pairs, merge the winner, repeat — thousands of times against a real corpus is exactly how GPT-2's actual vocabulary (50,257 byte-level tokens) or the vocab_size=1000 used in the worked example later in this chapter is built. In code:

from collections import Counter

def get_pair_counts(corpus):
    # corpus: dict mapping tuple-of-symbols -> frequency
    pairs = Counter()
    for symbols, freq in corpus.items():
        for i in range(len(symbols) - 1):
            pairs[(symbols[i], symbols[i + 1])] += freq
    return pairs

def merge_pair(pair, corpus):
    a, b = pair
    merged = a + b
    new_corpus = {}
    for symbols, freq in corpus.items():
        new_symbols, i = [], 0
        while i < len(symbols):
            if i < len(symbols) - 1 and symbols[i] == a and symbols[i + 1] == b:
                new_symbols.append(merged)
                i += 2
            else:
                new_symbols.append(symbols[i])
                i += 1
        new_corpus[tuple(new_symbols)] = freq
    return new_corpus

def train_bpe(corpus, num_merges):
    merges = []
    for _ in range(num_merges):
        pairs = get_pair_counts(corpus)
        if not pairs:
            break
        best = max(pairs, key=pairs.get)          # ties: first found, in dict order
        merges.append(best)
        corpus = merge_pair(best, corpus)
    return merges, corpus

corpus = {
    ('l', 'o', 'w', ''): 5,
    ('l', 'o', 'w', 'e', 'r', ''): 2,
    ('n', 'e', 'w', 'e', 's', 't', ''): 6,
    ('w', 'i', 'd', 'e', 's', 't', ''): 3,
}
merges, final_corpus = train_bpe(corpus, num_merges=3)
# merges == [('e', 's'), ('es', 't'), ('est', '')]

Once training produces an ordered list of merges, applying the tokenizer to new text is deterministic: split into bytes (or characters), then apply the learned merges in the order they were learned, repeatedly, until no merge in the list applies. The vocabulary is the set of all symbols that appear anywhere in that process — the original bytes plus every merged unit created along the way — each assigned an integer id. Encoding and decoding are then inverses of each other:

def encode(text, merges, token_to_id):
    symbols = list(text) + ['']
    for a, b in merges:                       # apply merges in learned order
        i = 0
        while i < len(symbols) - 1:
            if symbols[i] == a and symbols[i + 1] == b:
                symbols[i:i + 2] = [a + b]
            else:
                i += 1
    return [token_to_id[s] for s in symbols]

def decode(ids, id_to_token):
    text = ''.join(id_to_token[i] for i in ids)
    return text.replace('', '')

# vocab built from the toy corpus's base characters plus every learned merge
vocab = sorted({c for w in corpus for c in w} | {a + b for a, b in merges})
token_to_id = {tok: i for i, tok in enumerate(vocab)}
id_to_token = {i: tok for tok, i in token_to_id.items()}

ids = encode('newest', merges, token_to_id)
assert decode(ids, id_to_token) == 'newest'   # round trip holds

This is the boundary the rest of the chapter builds on: encode turns text into the integer sequence the model consumes, decode turns the model's output ids back into text, and everything the model itself does happens strictly in between.

Where tokenization ends and the model begins

Assume the tokenizer has already done its job: a document has become a sequence of integers, each in [0, vocab_size), each indexing a row of a vocabulary table. The model never sees characters or words again — only integers. Its first job is to turn each integer into a vector the network can compute with, and its last job is to turn a vector back into a probability distribution over the same integer space. Everything in between is what makes it a language model rather than a lookup table: it has to use the vectors at earlier positions to shape the distribution at the current position, respecting the fact that position 5 may not see position 8 (that would be cheating — position 8 is the answer key).

The architecture, layer by layer

A decoder-only transformer — the GPT family's architecture — is five ideas stacked on top of each other:

Token embedding. A lookup table of shape (vocab_size, d_model). Row t is the learned vector for token id t. This table starts as random noise and only acquires meaning through training — a point the misconception section below makes precise.

Positional embedding. Self-attention (below) computes a weighted average over positions with no notion of order baked in — swap the order of the input vectors and, absent something extra, attention produces the same set of outputs just reshuffled. GPT-2 (Radford et al., "Language Models are Unsupervised Multitask Learners," 2019) fixed this with a second lookup table of shape (max_seq_len, d_model), one learned vector per position, added directly to the token embedding. The original Transformer paper (Vaswani et al., "Attention Is All You Need," NeurIPS 2017) used a fixed sinusoidal formula instead of a learned table; GPT-2's version is simpler and is what we build here.

Causal self-attention. At each position, the model computes a query vector and compares it against the key vectors of every position up to and including itself, turning the comparison scores into attention weights via softmax, and using those weights to blend the corresponding value vectors. The comparison is scaled by 1/√d_head — without this, the dot products grow with the dimension of the vectors being compared, pushing softmax into regions with near-zero gradient. "Causal" means positions after the current one are masked to -∞ before the softmax, so they receive exactly zero weight: a training-time engineering constraint, not a modeling nicety, since without it the model would learn to peek at the token it is being asked to predict.

The MLP. After attention mixes information across positions, a two-layer feed-forward network processes each position independently, expanding to 4×d_model and back down, with a GELU nonlinearity (Hendrycks and Gimpel, "Gaussian Error Linear Units," 2016) in between. This is where most of a transformer's parameters and most of its per-token computation live.

Residual connections and pre-LayerNorm. Both attention and the MLP are wrapped as x = x + sublayer(LN(x)) rather than replacing x outright. Two consequences: the "residual stream" that a 96-layer model like GPT-3 carries end to end is a running sum, not a relay race where information must survive being fully overwritten at every layer; and normalizing the input to each sublayer (pre-LN) rather than its output (post-LN, the original 2017 design) gives well-behaved gradients at initialization without the layer count-dependent warmup schedule post-LN needs — analyzed formally by Xiong et al., "On Layer Normalization in the Transformer Architecture," ICML 2020. GPT-2 adopted pre-LN; we do too.

Here is the full stack as code — small enough to read end to end, structurally identical to a real GPT:

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

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        assert d_model % n_head == 0
        self.n_head = n_head
        self.d_head = d_model // n_head
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, T, C = x.shape
        qkv = self.qkv(x)                                            # (B, T, 3C)
        q, k, v = qkv.split(C, dim=2)                                 # each (B, T, C)
        q = q.view(B, T, self.n_head, self.d_head).transpose(1, 2)    # (B, nh, T, hd)
        k = k.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)      # (B, nh, T, T)
        mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T)
        att = att.masked_fill(mask == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        out = att @ v                                                 # (B, nh, T, hd)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out)

class MLP(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.fc1 = nn.Linear(d_model, 4 * d_model)
        self.fc2 = nn.Linear(4 * d_model, d_model)

    def forward(self, x):
        return self.fc2(F.gelu(self.fc1(x)))

class Block(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_head)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = MLP(d_model)

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

class TinyGPT(nn.Module):
    def __init__(self, vocab_size, d_model, n_layer, n_head, seq_len):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(seq_len, d_model)
        self.blocks = nn.ModuleList([Block(d_model, n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        self.apply(self._init_weights)
        for name, p in self.named_parameters():
            if name.endswith('proj.weight') or name.endswith('fc2.weight'):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * n_layer))
        self.head.weight = self.tok_emb.weight   # weight tying, applied after init

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

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        logits = self.head(x)                                        # (B, T, vocab_size)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

Two lines deserve a second look. First, self.head.weight = self.tok_emb.weight: an nn.Embedding(vocab_size, d_model) stores its table as a (vocab_size, d_model) matrix, and an nn.Linear(d_model, vocab_size, bias=False) stores its weight as (vocab_size, d_model) too — literally the same shape, so the assignment isn't a hack, it's pointing two names at one tensor. Every gradient step that updates a token's input embedding also updates that token's output-logit direction. This is weight tying, introduced by Press and Wolf ("Using the Output Embedding to Improve Language Models," EACL 2017), and every subsequent OpenAI GPT model uses it. Second, the loop that rescales proj.weight and fc2.weight — the two matrices that write directly back into the residual stream — by 0.02/√(2·n_layer): this is GPT-2's specific initialization trick, and the next section derives why it's necessary.

Initialization: why the first forward pass must not blow up

Every weight starts as N(0, 0.02²) noise. Consider the residual stream's variance as it passes through the network: each block adds the output of an attention sublayer and an MLP sublayer to the running sum. If those additions were left at unscaled N(0, 0.02²) initialization, the variance of the residual stream would grow roughly linearly with the number of blocks that have added to it — by block 48 in a GPT-2-scale model, activations would be dozens of times larger than at block 1, pushing LayerNorm and softmax into numerically unstable regions before a single gradient step has been taken. GPT-2's fix scales the standard deviation of exactly the two matrices that write into the residual stream (attention's output projection and the MLP's second layer) by 1/√(2·n_layer), so that the total variance contributed across all 2·n_layer such writes (two per block: one from attention, one from the MLP) stays bounded regardless of depth. It is a targeted fix, not a global one — it deliberately leaves the token embedding, positional embedding, QKV projection, and MLP's first layer at the unscaled 0.02, because those don't accumulate across depth the same way.

Worked example: every parameter in a 172,288-parameter GPT

Take vocab_size=1000, d_model=64, n_layer=2, n_head=4, seq_len=128 — small enough to count by hand, structurally identical to the code above. nn.LayerNorm(d) contributes 2d parameters (a learned scale and shift per dimension); an nn.Linear(a, b) contributes a·b weights plus b biases.

Tracing shapes through forward with a batch of B=2 sequences of length T=4: idx enters as (2, 4) of dtype int64. tok_emb(idx) gives (2, 4, 64); pos_emb(arange(4)) gives (4, 64), which broadcasts against the batch dimension when added, leaving x at (2, 4, 64). Inside CausalSelfAttention, qkv(x) produces (2, 4, 192), split into three (2, 4, 64) tensors, reshaped to (2, 4, 4, 16) and transposed to (2, 4, 4, 16) in (B, n_head, T, d_head) order (4 heads × 16 dims = 64, recovering d_model). The attention score tensor is (2, 4, 4, 4) — batch, heads, query positions, key positions — masked and softmaxed over the last axis, then combined with v back to (2, 4, 4, 16), reassembled to (2, 4, 64). Every block preserves the (2, 4, 64) shape by construction (residual addition requires it). head(x) maps the final (2, 4, 64) to logits of shape (2, 4, 1000) — one distribution over the 1,000-token vocabulary at every one of the 8 (batch × sequence) positions.

Common misconception: "more parameters means the model starts out knowing something"

A 172,288-parameter network sounds substantial, and it's tempting to think a freshly initialized model already has some baseline competence proportional to its size. It doesn't. At initialization every weight is independent random noise, so the model's output distribution at any position is, to a first approximation, uniform over the vocabulary — it hasn't seen a single example of what follows what. The cross-entropy loss for a uniform distribution over V classes is exactly -log(1/V) = log(V). For vocab_size = 1000, that's ln(1000) = 6.9078 nats — a number you can compute before writing a single training step, and a number you should see on the very first logged loss when you actually run one. Equivalently, since perplexity is defined as exp(loss), the model's perplexity at step 0 is exp(6.9078) = 1000 — literally as confused as if it were rolling a fair 1,000-sided die for every token. Training is the process of driving that loss down from ln(V) toward whatever floor the data's true entropy allows; parameter count determines how low the model is capable of getting, not where it starts.

The training loop: what happens on the GPU each step

Given the model above, one training step is five operations in a fixed order: compute the loss, clear old gradients, backpropagate, clip the gradient, step the optimizer. GPT-3 (Brown et al., "Language Models are Few-Shot Learners," NeurIPS 2020) settled on a specific recipe that has become the default for GPT-style pretraining: AdamW with β₁=0.9, β₂=0.95 (a lower β₂ than Adam's usual 0.999, making the second-moment estimate react faster), weight decay 0.1, global gradient-norm clipping at 1.0, and a cosine learning-rate decay preceded by a short linear warmup.

model = TinyGPT(vocab_size=1000, d_model=64, n_layer=2, n_head=4, seq_len=128)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)

def get_batch(split):
    ...  # (assumed helper, not shown) returns (idx, targets), each (B, T) int64

warmup_steps, max_steps, peak_lr = 100, 5000, 3e-4

def lr_at(step):
    if step < warmup_steps:
        return peak_lr * step / warmup_steps
    progress = (step - warmup_steps) / max(1, max_steps - warmup_steps)
    return 0.5 * peak_lr * (1 + math.cos(math.pi * progress))

for step in range(max_steps):
    idx, targets = get_batch('train')
    logits, loss = model(idx, targets)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    for group in optimizer.param_groups:
        group['lr'] = lr_at(step)
    optimizer.step()

clip_grad_norm_ computes one global L2 norm across every parameter's gradient combined — ‖g‖ = √(Σₚ ‖gₚ‖²) — and if that exceeds max_norm, rescales every gradient tensor by the same factor max_norm / ‖g‖, so the whole gradient vector's norm becomes exactly 1.0 without changing its direction. This has to run after backward() (gradients don't exist before then) and before optimizer.step() (otherwise the update uses the un-clipped values). Without clipping, a single batch containing an unusually long repeated sequence or an outlier target can produce a gradient spike that AdamW's per-parameter adaptive step size doesn't fully protect against, corrupting the moment estimates for many subsequent steps — the practical reason every large-scale GPT training run clips.

Diagram: forward pass, block internals, and where the parameters live

Active recall

Attempt each question before reading its answer.

1. Weight tying requires tok_emb.weight and head.weight to have the same shape. What shape, and why does it hold automatically for any GPT-style model?
Both must be (vocab_size, d_model). It holds automatically because nn.Embedding(vocab_size, d_model) stores its table as (vocab_size, d_model), and nn.Linear(d_model, vocab_size) stores its weight as (out_features, in_features) = (vocab_size, d_model) — the same shape for any choice of vocab_size and d_model, so tying never requires a reshape or transpose.

2. Compute the total parameter count for the 172,288-parameter model if the output head is untied (given its own independent weight matrix instead).
Untying adds one full (vocab_size, d_model) = (1000, 64) matrix with no bias (the code uses bias=False): 1000 × 64 = 64,000 new parameters. New total: 172,288 + 64,000 = 236,288 — a 37% increase for a component that, tied, cost nothing extra.

3. Change n_head from 4 to 2, keeping d_model=64 fixed. Does the total parameter count change? Trace the full effect, not just the attention block.
No — the total stays at 172,288. n_head only determines d_head = d_model / n_head (32 instead of 16), which affects how the existing qkv and proj weight matrices — still shaped (d_model, 3·d_model) and (d_model, d_model) regardless of head count — are reshaped and split during the forward pass. It changes the shape of intermediate attention-score tensors (from (B, 4, T, T) to (B, 2, T, T)) and how compute is distributed across heads, but it adds or removes no learnable weight. The common error is assuming "more heads" means "more parameters"; heads partition an existing matrix, they don't create a new one.

4. Increase seq_len from 128 to 256. Which specific components change parameter count, and by how much? Which stay fixed?
Only the positional embedding changes: 256 × 64 = 16,384 versus 128 × 64 = 8,192, a +8,192 increase, bringing the total to 172,288 + 8,192 = 180,480. The token embedding, every block's weights, and the final LayerNorm are untouched — none of their matrix shapes reference seq_len. The one other thing that grows is the causal mask itself, from 128×128 to 256×256 — but that's a non-learnable buffer built fresh each forward call, not a parameter.

5. Derive the expected cross-entropy loss and perplexity of this model at initialization, for vocab_size = 1000.
An untrained network's output is approximately uniform over the vocabulary, so loss ≈ -log(1/V) = log(V) = ln(1000) = 6.9078 nats. Perplexity is exp(loss), which for a uniform distribution over V classes recovers exp(ln V) = V = 1000 exactly — the model is, on average, as uncertain as an even guess among all 1,000 tokens.

6. In the training loop, clip_grad_norm_ runs after loss.backward() but before optimizer.step(). What does it compute, and why must it sit exactly there?
It computes the single global L2 norm across every parameter's gradient combined, ‖g‖ = √(Σₚ ‖gₚ‖²), and if that exceeds max_norm, rescales every gradient tensor by max_norm / ‖g‖ so the combined norm becomes exactly max_norm, preserving direction. It cannot run before backward() because gradients don't exist until then, and it must run before step() because the optimizer applies whatever values the gradients hold at the moment step() is called — clipping afterward would have no effect on the update already made.

Think About It

Think about this: How would you explain building large language models from scratch: tokenization to training to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where building large language models from scratch: tokenization to training is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting building large language models from scratch: tokenization to training to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind building large language models from scratch: tokenization to training, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.