Training a Language Model from Scratch (Part 1: Building Blocks)

This is for future me when I need to review the building blocks of an LM.

While working through CS336 Assignment 1, I kept running into places where I thought I understood the concept, but the implementation details exposed a missing piece. RoPE was not just “add position somehow”; the tokenizer was not the same thing as the embedding layer; cross entropy was not just “take softmax and index the answer”; gradient clipping was not per-parameter clipping.

This post is a cleaned-up version of those notes and aha moments. It is not meant to be a full assignment solution dump. It is the map I want future me to have when I need to reload how the pieces connect: tokenizer, embeddings, RoPE, attention, normalization, feed-forward layers, loss, optimization, and decoding.

Background refresher

Before the implementation details, I want to keep a small shelf of background notes I can come back to. Some of these are notes I created while course staff, and some are materials Sarah Pohland created during the semester. I like them because they connect the Transformer machinery back to a few intuitive ideas: similarity as an inner product, dimensions, attention masks, and training vs inference.

The map of the building blocks

I. Tokenizer: byte-level BPE

The tokenizer is not learning semantic vectors. It is deciding how raw text is split into chunks, and mapping each chunk to an integer ID.

The byte-level BPE tokenizer starts with a universal fallback:

For English, many characters are one byte:

For Korean, one visible character may require multiple UTF-8 bytes:

So initially, "한" might be represented by three byte-level token IDs. If it appears often enough in the training corpus, BPE may learn a merged token for those bytes, letting "한" become one token ID.

BPE training

BPE training is where the vocabulary grows. The sequence gets more compact, but the vocabulary gets larger.

The high-level loop is:

1. Start with the 256 byte tokens, plus special tokens like <|endoftext|>.
2. Pre-tokenize raw text into rough chunks.
3. Count repeated pre-tokens.
4. Convert each pre-token to UTF-8 byte pieces.
5. Count adjacent byte/token pairs, weighted by pre-token frequency.
6. Merge the most frequent adjacent pair globally.
7. Repeat until the target vocab size is reached.

Example:

My rough tokenizer sketch: byte fallback, pre-token frequency, UTF-8 byte pieces, and adjacent pair counts.

The BPE loop keeps merging the most frequent adjacent pair until the target vocabulary size is reached.

BPE encoding

Encoding is different from training.

During encoding, the merges are already learned. We do not count frequencies again. For each pre-token:

1. Start from single-byte pieces.
2. List adjacent pairs currently present.
3. Keep only pairs that exist in the learned merge_ranks.
4. Merge the earliest learned pair.
5. Repeat until no learned pair applies.
6. Map final byte pieces to token IDs.

pieces = [bytes([b]) for b in token_bytes]

while True:
    pairs = [(pieces[i], pieces[i + 1]) for i in range(len(pieces) - 1)]
    valid_pairs = [pair for pair in pairs if pair in merge_ranks]

    if not valid_pairs:
        break

    best_pair = min(valid_pairs, key=lambda pair: merge_ranks[pair])
    pieces = merge_everywhere(pieces, best_pair)

Special tokens add one more rule: split them out first. If the text contains <|endoftext|>, we encode that whole string as one ID, then run normal BPE only on the surrounding text.

II. Embeddings and PyTorch shape conventions

After tokenization, the model sees integer IDs:

token_ids = [[41, 93, 17, 10]]

The embedding layer is just a learned lookup table:

where $V$ is vocab size and $D$ is model width. If token ID 41 appears, the model takes row 41 of E.

class Embedding(torch.nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.empty(vocab_size, d_model))

    def forward(self, token_ids):
        return self.weight[token_ids]

So the flow is:

The other important shape convention is linear layers. Papers often write $Wx$ using column vectors. PyTorch usually stores activations as row/batched vectors:

So we store:

and compute:

y = x @ W.T

III. The Transformer block

The block I implemented is a pre-norm decoder block:

Pre-norm means we normalize before the attention/FFN sublayer, then add the residual connection.

Transformer LM overview and pre-norm Transformer block.

RMSNorm

LayerNorm subtracts the mean and rescales. RMSNorm only rescales by the root mean square:

The learnable scale $g$ has shape [d_model]. It broadcasts across [batch, sequence, d_model].

Multi-head attention shape flow

The packed projection gives Q/K/V in model width:

Then we split the final feature dimension into heads:

Q = q_proj(x).reshape(B, L, H, d_k)
K = k_proj(x).reshape(B, L, H, d_k)
V = v_proj(x).reshape(B, L, H, d_k)

Then transpose so attention runs over sequence length inside each head:

Q = Q.transpose(-3, -2)  # [B, H, L, d_k]
K = K.transpose(-3, -2)
V = V.transpose(-3, -2)

The attention matrix is [B, H, L, L], where each row says which previous tokens that position reads from.

After attention, we reverse the shape path:

out = attention(Q, K, V)          # [B, H, L, d_k]
out = out.transpose(-3, -2)       # [B, L, H, d_k]
out = out.reshape(B, L, D)        # [B, L, D]

SwiGLU

SwiGLU is the feed-forward block:

In PyTorch row-vector shape:

gate = silu(x @ W1.T)    # [..., d_ff]
value = x @ W3.T         # [..., d_ff]
hidden = gate * value    # [..., d_ff]
out = hidden @ W2.T      # [..., d_model]

The SiLU part is:

So the W1 branch is not just another plain projection. It makes a data-dependent gate for every token and every hidden feature. Large positive gate inputs mostly pass through, large negative gate inputs get pushed close to zero, and values near zero are softened smoothly. Then gate * value says: for this token, which expanded features from the W3 branch should be emphasized or muted before projecting back down?

IV. RoPE: position through rotation

Self-attention compares token vectors with dot products. Without positional information, the model knows which tokens exist, but not where they occur. Absolute position embeddings attach a learned vector to index 0, 1, 2, and so on, but that has two limitations:

• positions beyond the trained context length have no learned vector,
• relative distance has to be learned indirectly.

RoPE takes a different route. It rotates query and key vectors by their token position.

One thing that confused me at first: this is not merging or rotating adjacent tokens. RoPE rotates adjacent feature dimensions inside one token’s Q/K vector.

If a single head has $d_k = 8$, then one query vector at position $i$ looks like:

RoPE groups adjacent feature dimensions inside that vector: $(0,1)$, $(2,3)$, $(4,5)$, and $(6,7)$. The token-to-token interaction happens later, when the rotated queries and keys are used in $QK^\top$. RoPE itself is a within-vector rotation.

Important detail:

For a 2D pair, the rotation matrix is:

The key identity is:

So if token $m$ gets rotation $R(m\theta)$, and token $n$ gets rotation $R(n\theta)$, their attention dot product becomes:

That is the core RoPE idea. The attention score now depends on the relative offset $n-m$, not just the two content vectors. If both positions shift by the same amount $k$, the relative spacing stays the same.

Full matrix idea

For a head dimension $d_k$, RoPE conceptually builds a block-diagonal matrix:

Each block rotates one pair of dimensions. Building this full matrix would be wasteful.

For one pair $k$ at token position $i$, the actual operation is:

So the first row of the rotation matrix gives the new even coordinate:

and the second row gives the new odd coordinate:

Instead, precompute:

for every position $i$ and pair index $k$, then apply the 2D rotation directly to even and odd coordinates.

The precomputed RoPE tables store one row per token position and one column per even/odd feature pair.

old_even = x[..., 0::2]
old_odd = x[..., 1::2]

new_even = old_even * cos - old_odd * sin
new_odd = old_even * sin + old_odd * cos

out[..., 0::2] = new_even
out[..., 1::2] = new_odd

For each feature pair, the first row of the 2D rotation matrix produces the new even coordinate, and the second row produces the new odd coordinate.

Generalizing the trick: split all even and odd feature columns, select the matching position rows from the precomputed sin/cos tables, then apply the same row-wise rotation everywhere.

The sin/cos values are not learned parameters. They are fixed buffers, reused across batches and layers:

self.register_buffer("cos_angles", cos_table, persistent=False)
self.register_buffer("sin_angles", sin_table, persistent=False)

A more complete version of the implementation looks like:

class RotaryPositionalEmbedding(torch.nn.Module):
    def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None):
        super().__init__()

        positions = torch.arange(max_seq_len, device=device)      # [max_seq_len]
        pair_idx = torch.arange(d_k // 2, device=device)          # [d_k / 2]

        inv_freq = 1.0 / (theta ** (2 * pair_idx / d_k))
        angles = positions[:, None] * inv_freq[None, :]           # [max_seq_len, d_k / 2]

        self.register_buffer("cos_angles", torch.cos(angles), persistent=False)
        self.register_buffer("sin_angles", torch.sin(angles), persistent=False)

    def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor:
        # x is usually Q or K with shape [B, H, L, d_k]
        # token_positions is usually [L], e.g. [0, 1, 2, ..., L-1]

        old_even = x[..., 0::2]  # [B, H, L, d_k / 2]
        old_odd = x[..., 1::2]   # [B, H, L, d_k / 2]

        cos = self.cos_angles[token_positions]  # [L, d_k / 2]
        sin = self.sin_angles[token_positions]  # [L, d_k / 2]

        out = torch.empty_like(x)
        out[..., 0::2] = old_even * cos - old_odd * sin
        out[..., 1::2] = old_even * sin + old_odd * cos
        return out

The reason this broadcasts is that cos and sin have shape [L, d_k / 2], which lines up with the last two dimensions of old_even and old_odd. The batch and head dimensions are leading dimensions, so PyTorch reuses the same position-angle table across batches and heads.

Then in attention:

Q = self.rope(Q, token_positions)
K = self.rope(K, token_positions)
# V is not rotated.

V. Training: many next-token predictions at once

The language modeling dataset is usually treated as one long stream of token IDs. Every step samples random windows:

x = tokens[start : start + L]
y = tokens[start + 1 : start + L + 1]

For a toy chunk:

x text:    the   cat   sat   on
y text:    cat   sat   on    mat

The Transformer keeps the sequence dimension:

Each position gets its own next-token logits.

The LM head is position-wise:

logits[:, i, :] = lm_head(hidden[:, i, :])

It is not cheating because the hidden state at position $i$ only contains information from tokens $\leq i$. The future blocking already happened inside causal attention.

Cross entropy

The model outputs logits, not probabilities:

For loss, flatten batch and sequence positions:

loss = cross_entropy(logits.reshape(-1, V), targets.reshape(-1))

That converts:

and:

The stable cross entropy identity I want to remember is:

and with the max trick:

The reason subtracting $m$ is legal is that softmax does not change when every logit is shifted by the same constant:

because both numerator and denominator are multiplied by the same factor $e^{-m}$:

Then cross entropy is just negative log of the correct-token probability:

So this is mathematically the same softmax probability. The exponentials are safer because the max shift guarantees:

Perplexity is the same information on a more intuitive scale:

Across many tokens, it is the inverse geometric mean probability assigned to the correct next token.

VI. Optimizer, schedule, clipping, checkpointing

The training-loop part I want to remember is how the scheduler and global gradient clipping show up in code:

for it in range(start_iter, args.max_iters):
    lr = get_lr_cosine_schedule(
        it=it,
        max_learning_rate=args.max_lr,
        min_learning_rate=args.min_lr,
        warmup_iters=args.warmup_iters,
        cosine_cycle_iters=args.cosine_cycle_iters,
    )

    for group in optimizer.param_groups:
        group["lr"] = lr

    x, y = get_batch(train_ids, args.batch_size, args.context_length, device)

    optimizer.zero_grad()
    logits = model(x)  # [B, L, V]
    loss = cross_entropy(logits.reshape(-1, vocab_size), y.reshape(-1))
    loss.backward()

    gradient_clipping(model.parameters(), args.max_grad_norm)
    optimizer.step()

For the cosine learning-rate schedule, the useful shape trick is:

This maps the decay interval $[T_w, T_c]$ into progress $s \in [0,1]$. Then:

smoothly moves from 1 to 0, mapping $\alpha_{\max}$ to $\alpha_{\min}$.

Gradient clipping has to come after loss.backward() for a simple reason: backward() is the step that computes and stores the current gradients in each parameter’s .grad field. Before that, there is nothing from the current loss to clip.

The important detail is that this clipping is global. It calculates one L2 norm across the entire set of gradients:

Then it applies one shared scale factor to every gradient tensor if the norm is too large, meaning it goes beyond the chosen max_l2_norm.

def gradient_clipping(
    parameters: Iterable[torch.nn.Parameter],
    max_l2_norm: float,
) -> None:
    """Clip gradients in-place by global L2 norm."""
    eps = 1e-6

    grads = [p.grad for p in parameters if p.grad is not None]
    if len(grads) == 0:
        return

    total_sq_sum = sum(torch.sum(g**2) for g in grads)
    l2_norm = torch.sqrt(total_sq_sum)

    if l2_norm >= max_l2_norm:
        scale = max_l2_norm / (l2_norm + eps)
        for g in grads:
            g.mul_(scale)

AdamW then uses those clipped gradients inside optimizer.step(). AdamW still keeps per-parameter moment state like $m$, $v$, and $t$, but the clipping decision happened globally before the optimizer update.

AdamW also decouples weight decay from the gradient:

Finally, checkpointing is simply saving:

• model state,
• optimizer state,
• current iteration.

This matters because AdamW’s moment estimates are part of training state. Restarting with only the model weights is not the same as resuming training.

VII. Decoding: use the last row

Training predicts next-token logits at every position. Decoding only needs the final position.

Given prompt token IDs:

context_ids = generated_ids[-context_length:]
logits = model(context_ids)[..., -1, :]

For a prompt like Hi I am going, the model still returns one logits row per input position:

input tokens:   Hi        I        am        going
output rows:    next(Hi)  next(I)  next(am)  next(going)
decode uses:                                  ^ this row

During training, that last row would be compared against the actual next token from the dataset. During decoding, that next token is unknown, so we sample it from next(going), append it, and repeat.

Then decode one token:

if temperature <= 0:
    next_id = torch.argmax(logits)
else:
    probs = torch.softmax(logits / temperature, dim=-1)
    next_id = torch.multinomial(probs, num_samples=1)

Top-p sampling first sorts tokens by probability, keeps the smallest prefix whose probability mass exceeds $p$, renormalizes, and samples from that smaller set.

If the sampled token is <|endoftext|>, generation can stop. If it does not appear, generation stops at the chosen maximum token limit.

Final mental picture

The pieces feel much less mysterious when I separate the levels:

1. Tokenizer: defines the discrete vocabulary and maps text to IDs.
2. Embedding: gives each ID a learned vector.
3. RoPE + attention: lets each position read previous positions with relative-position geometry.
4. Transformer block: repeatedly mixes and transforms prefix information.
5. LM head: maps every hidden row to vocab logits.
6. Cross entropy: pushes up the correct next-token logit.
7. Optimizer/training loop: turns those losses into stable parameter updates.
8. Decoder: repeatedly reads the last-position logits and appends one sampled ID.

That is the building-block view I want to keep.