Pre-training a large language model.

Learn one conditional distribution at every position.

Given a token sequence x₁ … xₜ, an autoregressive model assigns a probability to the sequence by factorizing it into next-token conditionals. Training minimizes the negative log-probability of the observed next token at every valid position.

For an input tensor X with shape [B,T], the targets are the same stream shifted by one position. The model returns logits [B,T,V], where V is vocabulary size. Cross-entropy compares every logit vector against one target ID. This simple interface should remain visible throughout the implementation; most silent failures are violations of shapes, masks, shifts, or token counts around it.

# one contiguous token stream, sampled at offset s
x = tokens[s : s + B*T].view(B, T)
y = tokens[s + 1 : s + B*T + 1].view(B, T)
logits = model(x)                 # [B, T, V]
loss = cross_entropy(logits.view(-1, V), y.view(-1))

Tokenization is the model’s first architecture decision.

A Python string is a sequence of Unicode code points, but a model vocabulary needs a finite alphabet that can represent every possible input. UTF-8 supplies that foundation: each code point becomes one to four bytes, and the 256 byte values guarantee coverage. UTF-16 and UTF-32 are valid encodings, but UTF-8 is compact for common web text and matches the byte-level design used by GPT-style tokenizers.

Byte-pair encoding begins with those 256 byte IDs. It counts consecutive pairs in the training sample, merges the most frequent pair into a new ID, rewrites the sequence, and repeats until the vocabulary reaches its target size. The training loop produces two inseparable artifacts: an ordered merge table and the byte sequence associated with every token ID.

ids = list(training_text.encode("utf-8"))
merges = {}
while 256 + len(merges) < target_vocab:
    pair_counts = count_adjacent_pairs(ids)
    pair = max(pair_counts, key=pair_counts.get)
    new_id = 256 + len(merges)
    ids = merge_every_occurrence(ids, pair, new_id)
    merges[pair] = new_id

Compression ratio is a diagnostic, not the only goal. It compares UTF-8 bytes with produced tokens and reveals whether the learned vocabulary fits the corpus. A tokenizer trained on one language or domain can fragment another into long byte sequences, consuming context and compute before the Transformer sees the content.

Encoding begins from bytes and applies only merges that exist in the learned table, always respecting merge rank. Decoding recursively expands a token ID back to its byte sequence and then uses UTF-8 decoding. A token may contain bytes that are not a valid standalone character; validity is restored only after neighboring token bytes are concatenated.

The regex, special tokens, and vocabulary size change model behavior.

Production BPE tokenizers usually split text with a regular expression before applying merges. This prevents merges across selected categories such as letters, numbers, punctuation, or whitespace. GPT-2 and later tokenizers use different split patterns, which means the same BPE idea can produce different segment boundaries and compression.

Special tokens represent boundaries or control structure: end of text, conversation roles, tool calls, document separators, or padding. They must be added through an explicit reserved-token policy. Encoding ordinary text should reject or escape them unless the caller intentionally enables them; otherwise user text can impersonate control structure.

SentencePiece commonly trains from Unicode text and can use BPE or unigram segmentation, while GPT-style libraries expose byte-level BPE with model-specific regex rules. The library name does not define the model contract; the serialized vocabulary, ordered merges, normalization rules, split pattern, and special-token policy do.

Adding new vocabulary entries after pre-training is not free. The new embedding and output rows have not learned from the original token stream, while existing text may be resegmented into unfamiliar units. Measure any proposed prompt compression against model quality, not token count alone.

Tokenizer artifacts explain several visible model quirks: spelling and character counting can be hard when a word is one token; arithmetic strings fragment irregularly; spaces may be attached to following words; capitalization can change segmentation; non-English text can consume more positions. Multimodal systems generalize the idea by quantizing image, audio, or video features into discrete or token-like representations, but the coverage and compression contract remains.

for sample in adversarial_text_suite:
    ids = tokenizer.encode(sample, allowed_special=set())
    assert tokenizer.decode(ids) == sample
    assert all(0 <= token_id < vocab_size for token_id in ids)

assert reserved_ids.isdisjoint(tokenizer.encode(user_text))

Build a data pipeline whose decisions can be audited.

The pre-training corpus determines which patterns receive gradient updates. A production pipeline normally moves through source registration, parsing, language and quality signals, exact and approximate deduplication, policy filtering, benchmark decontamination, mixture weighting, tokenization, and immutable sharding.

Keep document boundaries explicit. Concatenating documents without an end marker teaches false transitions between unrelated texts. Packing several documents into one context improves utilization, but the token stream must retain boundary tokens and the training manifest must retain source membership.

Split evaluation documents before approximate deduplication, then remove overlapping clusters from training. Otherwise a benchmark can look strong because the model saw a paraphrase, duplicate, or answer-bearing neighbor. Report mixture weights in tokens—not only documents or bytes—because tokens are what produce optimizer updates.

Small teaching corpora can be loaded as one contiguous token array. Production corpora should be stored as immutable token shards with enough metadata to reproduce every boundary and sampling decision. Advance the shard cursor deterministically and checkpoint it with the optimizer state.

Each batch is a set of shifted windows from one token stream.

For a contiguous token array, sample a start position and take T+1 IDs. The first T become the input; the same window shifted by one becomes the target. Stacking B windows creates inputs and targets with shape [B,T]. A held-out split must use different documents or shards, not adjacent windows from the same duplicated source.

window = tokens[start : start + T + 1]
x = window[:-1]       # [T]
y = window[1:]        # [T]

x_batch = stack([sample_window(train_tokens) for _ in range(B)])
y_batch = stack([sample_targets(train_tokens) for _ in range(B)])

Token embeddings produce [B,T,C]; positional information is added or applied inside attention; a stack of blocks preserves [B,T,C]; the language-model head projects to [B,T,V].

The 124M GPT-2 configuration is a useful reference point: 12 blocks, width 768, 12 attention heads, context 1024, and vocabulary 50,257. Assert every shape and calculate the parameter count from the configuration at startup. A wrong bias, untied output matrix, or incorrect feed-forward width can add millions of parameters while leaving the code visually plausible.

Causal attention is learned, content-addressed communication.

A bigram language model is the minimum baseline: the current token indexes a row of logits for the next token. It cannot use an earlier subject, indentation level, or open delimiter because no context crosses positions. Averaging previous embeddings adds context but treats every previous position identically. Self-attention makes the aggregation depend on the content at both ends of the connection.

For hidden states X ∈ ℝT×C, learned projections create queries, keys, and values. Query-key dot products measure which previous positions matter to each current position. The scores are scaled, future positions are masked, a softmax produces weights, and the weights mix value vectors.

q, k, v = project(x).chunk(3, dim=-1)
q, k, v = split_heads(q), split_heads(k), split_heads(v)
scores = (q @ k.transpose(-2, -1)) / sqrt(head_dim)
scores = scores.masked_fill(causal_mask == 0, -inf)
weights = softmax(scores, dim=-1)
y = merge_heads(weights @ v)

Multiple heads perform the same operation in smaller subspaces and concatenate their results. The attention weights move information only within a sequence; the batch dimension contains independent examples. Because attention sees a set unless position is added, token and positional embeddings must meet before the first block or position must be injected into the attention calculation.

With dense attention, the score matrix grows as T². Longer context therefore increases both memory and computation sharply. Fused attention kernels avoid writing the full score matrix to high-bandwidth memory, which changes feasibility without changing the mathematical result.

Match the reference architecture before optimizing it.

GPT-2 124M uses a token embedding, learned absolute position embedding, 12 Transformer blocks, a final layer normalization, and a vocabulary projection. Each block contains causal multi-head attention and a four-times wider MLP with GELU. The attention implementation commonly stores the query, key, and value projections in one combined matrix, then splits the result.

x = token_embedding(ids) + position_embedding(positions)
for block in blocks:
    x = x + block.attn(block.ln_1(x))
    x = x + block.mlp(block.ln_2(x))
logits = lm_head(final_norm(x))

The model width must divide evenly across attention heads. GPT-2’s width 768 and 12 heads produce a head width of 64. The MLP expands 768 → 3072 → 768. Combined projection weights from another framework may require transposition when copied into a local Linear layer; compare every parameter name and shape before loading.

The token embedding and language-model head share the same weight matrix. Weight tying saves parameters and makes input and output token geometry use one table. Confirm object identity or shared storage, not merely equal initialization.

The model ends with logits [B,T,V]. During training, every position contributes a target. During sampling, only the final position chooses the next token, which is appended before repeating the forward loop.

Initialize for depth, then prove parity with a known checkpoint.

GPT-2 initializes most linear and embedding weights from a normal distribution with standard deviation near 0.02, with biases at zero and normalization scales at one. Residual branches accumulate through depth, so the output projections of attention and the MLP receive an additional depth-dependent scale.

A strong architecture test is to load the released GPT-2 124M weights and compare logits for a fixed token prefix. Resolve transposed Conv1D-style matrices, tied embeddings, and parameter-name mapping once. If the logits match within tolerance, the module hierarchy, shapes, norms, attention mask, and output head are jointly validated.

Define the optimizer in global tokens, not device-local batches.

AdamW maintains first and second moments for every trainable tensor and applies decoupled weight decay. Matrix weights receive decay; biases, normalization parameters, and often embeddings are placed in explicit non-decay groups. Fused AdamW can combine many elementwise updates into fewer kernels when the device supports it.

A GPT-style schedule usually warms the learning rate linearly, then decays it with a cosine or similar schedule toward a lower floor. Warmup prevents early updates from overwhelming uncalibrated activations and optimizer moments. Global gradient clipping limits rare spikes before the optimizer step.

optimizer.zero_grad(set_to_none=True)
for micro_step in range(accumulation_steps):
    x, y = next_batch()
    with autocast(dtype=bf16):
        loss = model(x, y) / accumulation_steps
    loss.backward()

grad_norm = clip_grad_norm_(model.parameters(), max_norm)
optimizer.step()
scheduler.step()

Gradient accumulation recreates a large logical batch from several micro-batches. Divide the loss before backpropagation so the accumulated gradient is an average. When using distributed data parallelism, synchronize gradients only on the final micro-step; earlier micro-steps can accumulate locally.

Record optimizer hyperparameters with the checkpoint: learning-rate schedule, warmup steps, decay floor, betas, epsilon, weight decay, clipping threshold, batch tokens, and total trained tokens. A weight file without these cannot explain the run.

Speed comes from reducing memory movement and idle work.

The GPT-2 reproduction lecture improves one implementation step by step: enable TensorFloat-32 for suitable matrix multiplications, use bfloat16 autocasting to reduce traffic while preserving range, compile the graph to remove Python overhead and fuse kernels, use fused causal attention, choose matrix-friendly dimensions, and fuse optimizer work. Each change is timed against the same batch.

Padding GPT-2’s vocabulary projection from 50,257 to a hardware-friendly multiple such as 50,304 can make the largest matrix multiplication more efficient, even though it adds unused rows. Inputs and targets must still use only real tokenizer IDs.

Distributed data parallelism places one complete model replica on each accelerator. Rank r consumes a disjoint token range, computes local gradients, and participates in all-reduce. Divide the desired global batch by world_size × B × T to obtain the accumulation count, and assert that the division is exact.

tokens_per_microstep = B * T * world_size
assert global_batch_tokens % tokens_per_microstep == 0
accumulation_steps = global_batch_tokens // tokens_per_microstep

loader_offset = rank * B * T
loader_stride = world_size * B * T

The checkpoint must capture the data cursor, rank-independent random state policy, model, optimizer, scheduler, global token count, and scaler state when applicable. On resume, compare the next batch IDs and next loss with an uninterrupted run.