The Foundations of LLMs, Part I: Pretraining and Midtraining

by David Oniani

Posted on

Table of Contents

I had been planning to make a YouTube video about this for quite some time. However, just as I was preparing to release it, Andrej Karpathy released an excellent video on the topic, and it quickly went viral. After that, I decided to hold off on recording my version. I am now turning the material into a written guide on the autoregressive transformer architecture, where autoregressive means that the model generates text one token at a time by feeding each new prediction back into its own input to predict the next one. We will implement a GPT-like Large Language Model (LLM) from scratch in Python using PyTorch as the only dependency. I assume basic familiarity with PyTorch, and for those new to it, Introduction to PyTorch Code Examples from Stanford provides a helpful starting point before we build and train our model end to end.

The complete code for the GPT-like autoregressive transformer implemented in this article is available in this GitHub Gist. In 222 lines of self-contained Python, the script automatically downloads Mary Shelley's Frankenstein, builds a character-level tokenizer, pretrains the transformer on a GPU or Apple Silicon chip, and samples new text from the trained weights.

Note that the model implemented and trained in this article, also referred to as the base model, is a sentence completer rather than a chatbot. Converting it into a chatbot involves midtraining, Supervised Fine-Tuning (SFT), and Reinforcement Learning (RL), which are beyond the scope of this article. Still, a pretrained generative LLM is a powerful tool for sampling knowledge.

Tokenization

Given input text, our objective is to generate output text conditioned on that sequence. Neural networks only operate on numbers, so in traditional transformer models, excluding byte-latent transformers that group raw UTF-8 bytes into dynamic patches without a fixed vocabulary, the first stage is always tokenization: converting a text string into a sequence of integer IDs. This applies whether the architecture is encoder-only (such as BERT, which reads an entire passage bidirectionally to classify or search text), encoder-decoder (such as T5, which encodes an input sequence and decodes a separate output sequence for tasks like translation), or decoder-only (such as GPT-2, which predicts the next token from left to right). The simplest method is character-level tokenization, where each distinct character in the text corpus receives a unique integer ID:

import string

class CharacterLevelTokenizer:
    """A character-level tokenizer that treats individual characters as tokens."""

    def __init__(self, text: str = string.printable) -> None:
        """Builds the vocabulary from the given text."""

        self.vocab: list[str] = sorted(set(text))
        self.vocab_size: int = len(self.vocab)
        self.char_to_token: dict[str, int] = {char: idx for idx, char in enumerate(self.vocab)}

    def encode(self, text: str) -> list[int]:
        """Converts a string into a list of token IDs."""

        return [self.char_to_token[char] for char in text]

    def decode(self, token_ids: list[int]) -> str:
        """Converts a list of token IDs back into a string."""

        return "".join(self.vocab[token_id] for token_id in token_ids)

tokenizer = CharacterLevelTokenizer()
enc = tokenizer.encode("robot")    # -> [87, 84, 71, 84, 89]
dec = tokenizer.decode(enc)        # -> "robot"

As shown above, CharacterLevelTokenizer extracts the sorted set of unique characters from the input text to build its vocabulary list self.vocab of length self.vocab_size. Calling encode looks up each character of "robot" in self.char_to_token to produce [87, 84, 71, 84, 89], and calling decode indexes back into self.vocab to reconstruct "robot" without losing information.

While character-level tokenization keeps the vocabulary tiny and simple, production LLMs use subword tokenizers so the model can read common words as single tokens rather than spelling every word letter by letter. Byte-Pair Encoding (BPE) starts with individual characters and repeatedly merges the most frequent adjacent pair in the corpus (for example, merging "t" and "h" into "th", and then "th" and "e" into "the"), whereas WordPiece merges pairs that maximize the likelihood of the training corpus under a unigram language model. In this article, we use CharacterLevelTokenizer so the entire pipeline runs with zero external libraries.

Architecture

Before examining attention and feed-forward layers in isolation, let us inspect the top-level Config dataclass and Transformer module. The Config object stores seven hyperparameters: vocab_size (number of unique tokens), block_size (maximum sequence length \(T\) the model can read at once), n_layer (number of stacked transformer blocks), n_head (number of parallel attention heads per block), n_embd (width of each token's feature vector), dropout (the probability \(p\) of randomly zeroing neuron outputs during training so the network does not memorize fixed pathways), and bias (whether linear and normalization layers add a learnable offset vector \(\mathbf{b}\) in \(\mathbf{y} = \mathbf{x}W^\top + \mathbf{b}\), which modern LLMs disable to save memory and improve stability):

from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F

@dataclass(frozen=True)  # We like our dataclasses frozen!
class Config:
    """Transformer config."""

    vocab_size: int  # Tokenizer vocabulary size
    block_size: int  # Max sequence length (context window)
    n_layer: int     # Number of transformer layers
    n_head: int      # Attention heads per layer
    n_embd: int      # Embedding dimension (must be divisible by `n_head`)
    dropout: float   # Dropout probability
    bias: bool       # Whether `nn.Linear` and `nn.LayerNorm` use bias

    @property
    def head_size(self) -> int:
        """Returns the per-head dimension (embedding is split evenly across attention heads)."""

        return self.n_embd // self.n_head

class Transformer(nn.Module):
    """Autoregressive transformer language model."""

    def __init__(self, cfg: Config) -> None:
        """Initializes the building blocks of the transformer."""

        super().__init__()
        self.block_size = cfg.block_size

        self.tok_emb_table = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.pos_emb_table = nn.Embedding(cfg.block_size, cfg.n_embd)

        self.blocks = nn.Sequential(*[Block(cfg) for _ in range(cfg.n_layer)])
        self.ln = nn.LayerNorm(cfg.n_embd, bias=cfg.bias)
        self.proj = nn.Linear(cfg.n_embd, cfg.vocab_size)

        # Weight tying: reduces the total number of parameters without degrading accuracy
        # Reference: https://arxiv.org/abs/1608.05859
        self.proj.weight = self.tok_emb_table.weight

    def forward(self, token_ids: torch.Tensor, targets: torch.Tensor | None = None) -> tuple:
        """Computes logits and optional loss."""

        B, T = token_ids.shape
        tok_emb = self.tok_emb_table(token_ids)
        pos_emb = self.pos_emb_table(torch.arange(T, device=token_ids.device))
        emb = tok_emb + pos_emb

        out = self.ln(self.blocks(emb))
        if targets is None:
            logits = self.proj(out[:, [-1], :])  # Last-token projection with time dimension
            loss = None
        else:
            logits = self.proj(out)
            loss = F.cross_entropy(logits.view(B * T, -1), targets.view(-1))

        return logits, loss

    @torch.no_grad()
    def generate(
        self,
        token_ids: torch.Tensor,
        max_new_token_ids: int,
        temperature: float = 0.7,
        top_k: int | None = None,
    ) -> torch.Tensor:
        """Generates token IDs autoregressively."""

        self.eval()
        for _ in range(max_new_token_ids):
            logits, _ = self(token_ids[:, -self.block_size :])
            logits = logits[:, -1, :] / temperature

            if top_k is not None:
                k = min(top_k, logits.size(-1))
                threshold = torch.topk(logits, k).values[:, [-1]]
                logits[logits < threshold] = -float("inf")

            probs = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            token_ids = torch.cat((token_ids, next_token), dim=1)
        return token_ids

Let us trace forward step by step. Its input token_ids is a 2D integer tensor of shape (B, T), where B is the batch size (how many independent text sequences we process in parallel) and T is the sequence length (how many token positions are in each sequence, up to block_size). Looking up token_ids in self.tok_emb_table produces tok_emb of shape (B, T, n_embd), which assigns each token ID a learnable vector representing its identity. Because token embeddings alone do not record word order, self.pos_emb_table maps the position indices 0, 1, ..., T - 1 (created by torch.arange(T, device=token_ids.device)) into a positional tensor pos_emb of shape (T, n_embd) rather than using fixed sinusoidal encodings. Adding tok_emb + pos_emb uses PyTorch broadcasting to stretch pos_emb across all B batch items, yielding a combined representation emb of shape (B, T, n_embd) that encodes both what each token is and where it sits in the sequence.

Learned positional embeddings are simple to implement and work well within a fixed context window, but they cannot extrapolate to positions \(t \ge \text{block_size}\) because those rows in pos_emb_table are never trained. Modern frontier architectures instead use Rotary Position Embedding (RoPE), which encodes position by rotating query and key vectors in 2D frequency planes at angles proportional to their token index \(t\). Because the dot product between two rotated vectors depends only on their relative distance \(t_1 - t_2\), RoPE generalizes more naturally to longer sequences. We use learned positional embeddings here to keep the code minimal.

Next, out = self.ln(self.blocks(emb)) passes the embeddings through self.blocks, a stack of transformer blocks that mix information across tokens via Multi-Head Attention (MHA) and refine each token's features via a Multilayer Perceptron (MLP). Afterward, self.ln applies layer normalization, which shifts and scales each token's n_embd feature vector to have mean 0.0 and variance 1.0 so activations stay well-behaved across deep stacks. Finally, self.proj projects the normalized features into logits, which are raw, unnormalized scores for every token in the vocabulary. Notice self.proj.weight = self.tok_emb_table.weight: because both nn.Embedding(vocab_size, n_embd) and nn.Linear(n_embd, vocab_size) store their weight matrix with shape (vocab_size, n_embd), they can share the exact same underlying tensor in memory (weight tying), ensuring that tokens with similar input embeddings receive similar output logit scores while reducing parameter count.

At the end of forward, execution branches based on whether ground-truth targets are supplied. During inference (targets is None), we only need to predict the single token that comes after the current sequence, so out[:, [-1], :] slices only the final time step (using a single-element list [-1] to preserve the 3D shape (B, 1, n_embd)) before projecting it to (B, 1, vocab_size) logits, saving unnecessary matrix multiplications on earlier positions. During training, self.proj(out) computes logits of shape (B, T, vocab_size) for every position. Because PyTorch's F.cross_entropy expects a 2D prediction matrix of shape (N, vocab_size) and a 1D label vector of shape (N,) where \(N = B \times T\), we reshape the tensors with logits.view(B * T, -1) (where -1 automatically infers vocab_size) and targets.view(-1).

This cross-entropy calculation implements the core objective of causal pretraining: next-token prediction. By maximizing the probability assigned to the true next token at every position in the training corpus, the model learns the statistical structure of language. Formally, we minimize the negative log-likelihood of the sequence:

$$\mathcal{L}_{\text{pretrain}}(\theta) = -\sum_{t=1}^T \log P_{\theta}(x_{t} \mid x_{\lt t})$$

Here, \(P_{\theta}\) represents our transformer parameterized by weights \(\theta\), and \(x_{t}\) is the true target token at position \(t\) given all preceding tokens \(x_{\lt t} = (x_1, \dots, x_{t-1})\). Intuitively, if the model assigns a probability near \(1.0\) to the true next character, \(-\log P_{\theta}\) is close to \(0.0\) and the loss penalty is tiny. If the model assigns a tiny probability to the true character, \(-\log P_{\theta}\) grows very large, creating a steep gradient signal that adjusts \(\theta\) to raise that character's score on future steps.

Multi-Head Attention and Transformer Blocks

While the final projection and loss calculation define the training objective, the model's ability to understand context comes from the repeated attention and feed-forward layers inside self.blocks. To see how tokens exchange information and build contextual meaning across a sequence, let us unpack self.blocks(emb) and examine AttentionHead, MultiHeadAttention, and Block at the tensor level:

class AttentionHead(nn.Module):
    """A single causal self-attention head."""

    def __init__(self, cfg: Config) -> None:
        """Initializes QKV projection and dropout, and caches causal mask to avoid recomputing it."""

        super().__init__()
        self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.head_size, bias=cfg.bias)
        self.dropout = nn.Dropout(cfg.dropout)
        self.register_buffer("mask", torch.tril(torch.ones(cfg.block_size, cfg.block_size)))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Computes masked single-head self-attention for the input tensor."""

        _, T, D = x.shape
        q, k, v = self.qkv(x).split(self.qkv.out_features // 3, dim=-1)

        attn_scores = q @ k.transpose(-2, -1)
        attn_scores = attn_scores * D**-0.5  # Prevent softmax from blowing up
        attn_scores = attn_scores.masked_fill(self.mask[:T, :T] == 0, float("-inf"))  # Mask future tokens

        attn_weights = torch.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        return attn_weights @ v

class MultiHeadAttention(nn.Module):
    """A Multi-Head Attention (MHA)."""

    def __init__(self, cfg: Config) -> None:
        """Initializes multi-head self-attention with output projection and dropout."""

        super().__init__()
        self.heads = nn.ModuleList([AttentionHead(cfg) for _ in range(cfg.n_head)])
        self.proj = nn.Linear(cfg.n_embd, cfg.n_embd)
        self.dropout = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Computes masked multi-head self-attention for the input tensor."""

        out = torch.cat([head(x) for head in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

class Block(nn.Module):
    """Transformer block."""

    def __init__(self, cfg: Config) -> None:
        """Initializes a transformer block."""

        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.n_embd, bias=cfg.bias)
        self.mha = MultiHeadAttention(cfg)
        self.ln2 = nn.LayerNorm(cfg.n_embd, bias=cfg.bias)
        self.mlp = nn.Sequential(
            nn.Linear(cfg.n_embd, 2 * cfg.n_embd, bias=cfg.bias),  # Flexible: can change to e.g. `4 * cfg.n_embd`
            nn.GELU(),
            nn.Linear(2 * cfg.n_embd, cfg.n_embd, bias=cfg.bias),
            nn.Dropout(cfg.dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Performs a forward pass: attention + MLP with residuals and layer norms."""

        x = x + self.mha(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

The AttentionHead class implements a single causal self-attention head. Attention, introduced in the Attention Is All You Need paper, is a directional communication mechanism that lets each token in a sequence gather relevant context from earlier tokens. Rather than running three separate matrix multiplications, self.qkv(x) projects each token's n_embd vector into a combined vector of width 3 * cfg.head_size, and .split(self.qkv.out_features // 3, dim=-1) slices that last axis into three tensors q, k, and v of shape (B, T, head_size):

To compute all pairwise query-key dot products in one operation, k.transpose(-2, -1) swaps the last two axes of k from (B, T, head_size) to (B, head_size, T), so the batched matrix multiplication q @ k.transpose(-2, -1) produces an (B, T, T) grid of raw similarity scores attn_scores. Next, we scale attn_scores by D**-0.5 (note that in this compact script D is unpacked from x.shape as cfg.n_embd, whereas standard multi-head formulations scale by cfg.head_size**-0.5). To enforce causality so token \(i\) can never peek at future tokens \(j \gt i\), __init__ pre-builds a lower-triangular matrix of ones via torch.tril(torch.ones(cfg.block_size, cfg.block_size)) and stores it with self.register_buffer("mask", ...), which tells PyTorch to move the tensor to the GPU alongside the model without treating it as a trainable parameter. During forward, masked_fill(self.mask[:T, :T] == 0, float("-inf")) replaces all upper-triangular future entries (\(j \gt i\)) with \(-\infty\). Formally, for sequence length \(n\) and scaling dimension \(d\), attention computes:

$$\text{Attention}(Q, K, V) = \text{softmax}\big(\frac{Q K^\top}{\sqrt{d}}\big) V$$

Dividing by \(\sqrt{d}\) keeps the variance of the dot products bounded as embedding dimensions grow, preventing softmax from saturating into extreme \(0.0\) or \(1.0\) probabilities where gradients vanish. Because \(\exp(-\infty) = 0\), torch.softmax(attn_scores, dim=-1) turns every masked future position into an exact 0.0 weight while normalizing the allowed past and present positions in each row to sum to 1.0. Multiplying attn_weights @ v then aggregates the value vectors into an output of shape (B, T, head_size).

In MultiHeadAttention, we run cfg.n_head independent AttentionHead modules in parallel, concatenate their (B, T, head_size) outputs along the last axis (dim=-1) to recover shape (B, T, n_embd), and pass the result through the output projection self.proj. Running multiple heads in parallel lets the model attend to several types of relationships simultaneously, with one head tracking local punctuation or syntax while another tracks long-range character names. Applying self.dropout to the projected output regularizes the layer so individual heads cannot co-adapt and learn redundant patterns.

The Block class pairs cross-token communication (self.mha) with per-token computation (self.mlp). Before each sublayer, layer normalization (self.ln1 and self.ln2) standardizes the input features, and after each sublayer, a residual connection (x = x + ...) adds the sublayer's input directly to its output. Because the derivative of \(x + f(x)\) with respect to \(x\) is \(1 + f'(x)\), this residual addition creates an unimpeded highway that carries gradients straight from the loss back to the earliest layers without vanishing. While self.mha mixes information across time steps, self.mlp processes each token position independently: it projects the n_embd vector up to a wider hidden dimension (2 * cfg.n_embd, chosen instead of the standard 4 * cfg.n_embd to keep our miniature model compact at 140,127 parameters), applies a smooth nonlinear GELU (Gaussian Error Linear Unit) activation so the two linear layers do not collapse into a single linear matrix multiply, and projects back down to cfg.n_embd.

Autoregressive Generation

During training, the model processes all \(T\) tokens of a sequence in parallel using teacher forcing, meaning the model always receives the true ground-truth tokens from the dataset as context at every position even if its own prediction at the previous position was wrong. At inference, however, future ground-truth tokens do not exist, so Transformer.generate runs sequentially: it first calls self.eval() to disable random dropout layers (note that if you call generate in the middle of a training loop, you should call model.train() afterward to re-enable dropout), crops the running sequence to the trailing self.block_size tokens via token_ids[:, -self.block_size :] so it never exceeds pos_emb_table, and extracts the 2D slice of final-position logits logits[:, -1, :] of shape (B, vocab_size) representing raw scores \(\mathbf{z} \in \mathbb{R}^{\lvert \mathcal{V} \rvert}\) across the vocabulary \(\mathcal{V}\). To balance coherence and creativity, we divide the logits by a temperature parameter \(\tau \gt 0\) and optionally restrict sampling to the \(k\) highest-scoring candidate tokens \(\mathcal{V}_k \subset \mathcal{V}\):

$$P(x_{t} = v \mid x_{\lt t}) = \frac{\exp(z_{v} / \tau)}{\sum_{j \in \mathcal{V}_k} \exp(z_{j} / \tau)}$$

Effect of temperature on next-token probabilities
Effect of temperature on next-token probabilities for candidate continuations

The diagram above illustrates how temperature \(\tau\) reshapes the probability distribution for the prompt "Artificial intelligence will transform the ...". Dividing logits by a small temperature (\(\tau = 0.2\)) magnifies score differences fivefold before exponentiation, concentrating 82% of the probability mass on the top candidate ("future") and driving tail options near zero. This sharp distribution suits deterministic tasks like coding or math, whereas moderate temperatures (\(\tau = 0.7\) and \(\tau = 1.0\)) give meaningful weight to plausible alternatives like "world" (22%) and "industry" (16%), and high temperatures (\(\tau = 2.0\)) compress logit gaps to flatten the distribution across candidates.

When top_k is enabled in generate, torch.topk(logits, k).values[:, [-1]] finds the \(k\)-th largest logit score in each row as a cutoff threshold, and logits[logits < threshold] = -float("inf") sets every token outside the top \(k\) to \(-\infty\) so torch.softmax assigns them zero probability. Rather than always picking the single highest-scoring token with torch.argmax (greedy decoding, which frequently gets stuck in repetitive loops), torch.multinomial(probs, num_samples=1) samples an index from the categorical distribution like spinning a weighted roulette wheel where each token's slice of the wheel equals its probability:

$$x_{t} \sim \text{Categorical}(P(x_{t} \mid x_{\lt t}))$$

High-probability tokens are selected most often, while lower-probability tokens inside the top-\(k\) set retain a controlled chance of being chosen, allowing natural variation across runs. Each sampled token index next_token of shape (B, 1) is appended to token_ids along the time dimension via torch.cat((token_ids, next_token), dim=1), and the loop repeats until max_new_token_ids tokens have been generated.

Pretraining Loop

With the architecture in place, we can pretrain our model on Mary Shelley's Frankenstein. In the script below, sample_batch picks batch_size random starting indices in the tokenized book, slices chunks of length block_size as input token_ids, and slices the same chunks shifted right by one character (idx + 1 : idx + block_size + 1) as targets. Notice how we also split the model's parameters by tensor dimension (param.ndim < 2): 1D tensors (such as LayerNorm scale vectors and the default bias vectors on self.proj) go to AdamW, while 2D weight and embedding matrices go to Muon:

import urllib.request

def sample_batch(data: torch.Tensor, batch_size: int, block_size: int) -> tuple:
    """Randomly samples training sequences for Next-Token Prediction (NTP)."""

    idxs = torch.randint(len(data) - block_size, (batch_size,))                    # Random starting positions
    token_ids = torch.stack([data[idx : idx + block_size] for idx in idxs])        # Input sequences
    targets = torch.stack([data[idx + 1 : idx + block_size + 1] for idx in idxs])  # Same sequences shifted by +1 (NTP)
    return token_ids, targets

with urllib.request.urlopen("https://www.gutenberg.org/cache/epub/84/pg84.txt") as f:  # Read Frankenstein
    text = f.read().decode("utf-8")
tokenizer = CharacterLevelTokenizer(text)

device = torch.device("cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
cfg = Config(vocab_size=tokenizer.vocab_size, block_size=32, n_layer=4, n_head=4, n_embd=64, dropout=0.0, bias=False)
max_steps, log_interval, batch_size, learning_rate = 2_000, 10, 256, 1e-3

data = torch.tensor(tokenizer.encode(text), device=device)
model = Transformer(cfg).to(device).train()

n_params, adamw_params, muon_params = 0, [], []
for param in model.parameters():
    n_params += param.numel()
    (adamw_params if param.ndim < 2 else muon_params).append(param)
print(f"Model parameters: {n_params:,}\n")

adamw = torch.optim.AdamW(adamw_params, lr=3e-4, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.01)
muon = torch.optim.Muon(muon_params, lr=0.02, momentum=0.95, weight_decay=0.1)
for step in range(1, max_steps + 1):
    token_ids, targets = sample_batch(data, batch_size, cfg.block_size)

    adamw.zero_grad()
    muon.zero_grad()
    _, loss = model(token_ids, targets)
    loss.backward()
    adamw.step()
    muon.step()

    if step % log_interval == 0:
        print(f"\rStep {step}/{max_steps} ({step / max_steps:.2%}) | Loss: {loss.item():.4f} ", end="", flush=True)

# Seed the model with the prompt "I am" to kick off text generation
token_ids = torch.tensor(tokenizer.encode("I am"), device=device).reshape(1, -1)
output = model.generate(token_ids, max_new_token_ids=512)[0]
print(f"\n\nOUTPUT:\n\n{tokenizer.decode(output.tolist())}")

Why do we pair two different optimizers with different learning rates (lr=3e-4 for AdamW and lr=0.02 for Muon, overriding the baseline learning_rate = 1e-3 constant)? AdamW updates each parameter coordinate independently by dividing a running average of the gradient (betas[0]=0.9) by the square root of a running average of squared gradients (betas[1]=0.95, plus eps=1e-8 to prevent division by zero) while gently shrinking weights toward zero at each step (weight_decay=0.01). For 2D weight matrices, however, raw gradients are often dominated by a few large singular directions. Muon accumulates gradient momentum (momentum=0.95) and applies Newton-Schulz matrix iterations to orthogonalize the 2D update matrix, equalizing the step size across all singular directions so the spectral norm is normalized and a larger learning rate (lr=0.02) can be taken safely. Inside each step of the loop, zero_grad() clears the previous step's gradients (since PyTorch accumulates gradients by addition into .grad buffers by default), loss.backward() uses the chain rule to compute gradients across the network, and step() updates the weights.

Pretraining Results

Running the 2,000-step pretraining loop takes under 40 seconds on a MacBook Pro with an Apple M4 Pro chip. Over those 2,000 steps (processing \(2{,}000 \times 256 \times 32 \approx 16.4\text{M}\) characters), the cross-entropy loss drops to 1.3831, and our 140,127-parameter character model learns English word boundaries, punctuation, and Shelley's vocabulary from raw character statistics:

Model parameters: 140,127

Step 2000/2000 (100.00%) | Loss: 1.3831

OUTPUT:

I am the forth which were on his stranger. The light of many our reasun alive in the enter of my
cottagers, there in the murderer, when I walked in
my enemy sat much of spring in its removed. I had not for my friend, and the cottager of the lack of
the wint on the lovely of cestant, and that the old man, but I was entered that
the eviner of my hands. My journey is many beauty will of compassins of the more spully appeared
with meaning a man work was not rease to be macking my own and a
gentle understand even

Midtraining

In production pipelines, base pretraining on web-scale text provides broad world knowledge, but raw web crawls are noisy and low in reasoning density. As analyzed in a recent survey on midtraining, frontier LLMs bridge this gap through midtraining, also known as continual pretraining. Midtraining optimizes the same autoregressive next-token prediction objective over a higher-quality data mixture and an extended sequence length across three pillars:

$$\mathcal{L}_{\text{midtrain}}(\theta) = -\sum_{t=1}^T \log P_{\theta}(x_{t} \mid x_{\lt t})$$

  1. Targeted Data Distributions: Upsampling high-utility reasoning domains (such as code, STEM mathematics, synthetic textbooks, step-by-step chain-of-thought traces, and fill-in-the-middle data where a middle span of code is moved to the end so the model learns to infill missing code given both prefix and suffix context) while retaining a replay buffer (typically 10% to 30%) of general web tokens to prevent catastrophic forgetting (overwriting previously learned general knowledge when adapting to a narrower distribution).

  2. Learning Rate Annealing: Decaying the learning rate along a cooldown schedule such as cosine decay (smoothly lowering the step size along a half-cosine curve) or warmup-stable-decay (holding the learning rate constant during base pretraining and rapidly decaying it during midtraining). Smaller gradient steps reduce mini-batch gradient noise variance and let the weights settle into a sharper, lower-loss minimum on curated reasoning data.

  3. Context Length Extension: Progressively expanding the context window block_size (often from 4K or 8K tokens up to 32K or 128K) by adapting positional encodings on a mixture of long and short documents. Because RoPE represents token positions as rotation angles at different frequencies (like clock hands moving at different speeds), rescaling those frequencies slows the rotations so longer sequences fit inside the trained angular range, while YaRN interpolates only the slow, low-frequency dimensions so sharp local token order is preserved. Training on short sequences before expanding to long sequences acts as curriculum learning (starting with easier examples before harder ones), which accelerates convergence and strengthens in-context learning (the model's ability to learn a new task purely from examples inside the prompt at inference time).

Conclusion

In this article, we implemented and pretrained a GPT-like autoregressive transformer from scratch using PyTorch, covering character-level tokenization, causal multi-head self-attention, residual connections, next-token prediction with AdamW and Muon, and the foundations of midtraining. While our 140K-parameter model produces only stylistic echoes of Shelley rather than fully coherent prose, the underlying algorithmic blueprint is identical to the one powering frontier systems. At scale, pretraining and midtraining establish broad world knowledge and structured reasoning, yielding capable base models. However, raw base models remain open-ended text completers. Transforming a next-token predictor into a helpful conversational assistant requires post-training, beginning with Supervised Fine-Tuning (SFT), which we explore in Part II.