Tag: LLM

DL0078 LLM Training Stages

What are the standard stages of LLM training?

Answer

A modern LLM is built in a sequence of stages that share one architecture but differ in data, objective, and scale: pretraining on a web-scale corpus with next-token prediction, an increasingly explicit mid-training phase that reweights the data mix and extends the context window, supervised fine-tuning (SFT) on curated prompt-response pairs, and preference optimization that turns human or programmatic judgments into a training signal (reward model plus PPO, DPO, or reinforcement learning with verifiable rewards). The split is not cosmetic: pretraining consumes well over 90% of the total tokens and FLOPs and is where almost all knowledge and reasoning capacity is acquired, while post-training uses a tiny fraction of the compute to select and expose behaviors the base model already has. Each stage produces a named artifact that the next stage consumes, so the usual lineage is base model → long-context base → instruct model → aligned model. Practical pipelines also interleave evaluation, safety filtering, distillation, and quantization, but those are packaging steps rather than new learning objectives.

(1) Pretraining: self-supervised next-token loss on trillions of filtered web, code, and book tokens, run once at enormous cost; this stage fixes the tokenizer, the parameter count, and the knowledge cutoff.
(2) Mid-Training: continued pretraining on a higher-quality mix with upsampled math, code, and long documents, plus context-length extension and learning-rate annealing, using roughly 1-10% of the pretraining tokens.
(3) Supervised Fine-Tuning: the same cross-entropy loss but computed only on response tokens with the prompt masked out, teaching format, instruction following, and tool-call syntax from $10^4$ to $10^6$ curated examples.
(4) Preference Optimization: optimizes a reward (learned from pairwise comparisons, or computed by a verifier) under a KL penalty toward the SFT reference policy, which is what shifts a model from plausible to preferred.
(5) Compute Asymmetry: post-training is cheap, so it is where iteration happens; anything that requires new knowledge or a longer context has to go back to a pretraining-style stage.

Four-stage LLM training pipeline: filtered web text feeds pretraining with a next-token loss producing a base model, a high-quality long-document mix feeds mid-training producing a long-context base, prompt-response pairs feed SFT with the prompt masked producing an instruct model, and preference pairs or verifiable tasks feed preference RL with a reward plus KL penalty, yielding a deployed assistant

Figure 1: The four canonical stages with their input data, objective, and output artifact. Only the data distribution and loss mask change between stages 1 to 3; stage 4 replaces likelihood with a reward under a KL constraint.

Mathematical Formulation:
$\mathcal{L}_{\mathrm{PT}}(\theta) = -\sum_{t=1}^{T} \log p_{\theta}(x_t \mid x_{1:t-1})$
$\mathcal{L}_{\mathrm{SFT}}(\theta) = -\sum_{t \in \mathcal{A}} \log p_{\theta}(y_t \mid x, y_{1:t-1})$
$\mathcal{L}_{\mathrm{RM}}(\phi) = -\log \sigma\big(r_{\phi}(x,y_w) - r_{\phi}(x,y_l)\big)$
$J(\theta) = \mathbb{E}\big[r_{\phi}(x,y)\big] - \beta\,\mathrm{KL}(\pi_{\theta}\,\|\,\pi_{\mathrm{ref}})$

Where:

$x_{1:t-1}$ is the preceding context and $x_t$ the target token, so pretraining averages this loss over every position of every document.
$x$ is the prompt, $y$ the response, and $\mathcal{A}$ the set of assistant token positions; positions outside $\mathcal{A}$ are masked, which is the only structural difference between SFT and pretraining.
$y_w$ and $y_l$ are the preferred and rejected responses for the same prompt, $r_{\phi}$ is the reward model, and $\sigma$ is the logistic function, giving the Bradley-Terry pairwise objective.
$\pi_{\theta}$ is the policy being trained and $\pi_{\mathrm{ref}}$ the frozen SFT checkpoint; the KL term is what keeps generations fluent instead of collapsing onto reward-model artifacts.
$\beta > 0$ sets the strength of that anchor: small $\beta$ permits reward hacking, large $\beta$ leaves the model barely changed from SFT.

Log-scale bar chart of training tokens per stage: pretraining 15 trillion tokens at 94.9 percent, mid-training 0.8 trillion at 5.1 percent, SFT 0.2 billion at 0.001 percent, and preference RL 4 billion at 0.03 percent

Figure 2: Illustrative token budget for a frontier-scale run. Pretraining and mid-training together see about 99.97% of all tokens, so post-training cannot add missing knowledge, only elicit and shape what is already in the weights.

Post-Training Option	Reward Model + PPO	DPO	Verifiable Rewards (GRPO)
Signal required	Pairwise human preferences, then a learned scalar reward	Pairwise preferences used directly, no reward model	A programmatic checker: unit tests, math answer match, schema validity
Models held in memory	Four: policy, reference, reward, critic	Two: policy and frozen reference	Two plus a sampler: no critic, group baseline replaces it
Online generation	Yes, rollouts dominate wall-clock time	No, offline on a fixed pair dataset	Yes, several samples per prompt
Main failure mode	Reward hacking and length inflation once the reward model is over-optimized	Off-policy drift: pairs stop reflecting the current policy, likelihood of both responses can fall	Only works where correctness is checkable; gaming the checker instead of the task
Typical use	Broad helpfulness and safety at frontier labs with annotation pipelines	Small teams, quick style and tone alignment on limited GPUs	Math, code, and reasoning models where long chains of thought pay off

August 15, 2026

DL0077 LLM Prediction Objective

What is a large language model, and what is its prediction objective?

Answer

A large language model (LLM) is a neural network, in practice a decoder-only Transformer with billions of parameters, trained on trillions of tokens of text to model the probability distribution over token sequences. Its prediction objective is next-token prediction: factor the sequence probability with the chain rule and maximize the log-likelihood of each observed token given its left context, which is identical to minimizing cross-entropy over the vocabulary at every position. The training signal is self-supervised, since the label at position $t$ is simply the token at position $t+1$ , so no human annotation is required and one sequence of length $T$ yields $T-1$ supervised examples in a single forward pass, or a full $T$ when a leading BOS token supplies the context for the first real token. Everything an LLM appears to know (syntax, factual associations, translation, arithmetic, code structure) is a byproduct of driving that one loss down, because accurately predicting the next token over a broad corpus requires modeling the processes that generated the text. Base pretraining optimizes likelihood only; instruction following and refusal behavior come later, from supervised fine-tuning and preference optimization, which change the objective rather than the architecture.

(1) What Makes It “Large”: parameter count in the billions, a training corpus in the trillions of tokens, and a compute budget allocated according to scaling laws, so capability improvements are largely predictable from loss curves rather than architectural novelty.
(2) The Objective Is Autoregressive: the chain-rule factorization plus a causal mask means position $t$ may only attend to positions $1$ through $t$ , which is what makes single-pass training and left-to-right generation consistent with each other.
(3) Loss Is Cross-Entropy Over the Vocabulary: the final layer projects the hidden state to $V$ logits, softmax turns them into a distribution, and the loss is the negative log-probability of the single correct token.
(4) Perplexity Is the Reported Metric: $\mathrm{PPL} = \exp(\mathcal{L})$ is the effective branching factor, so a mean loss of $2.3$ nats means the model is about as uncertain as choosing uniformly among 10 tokens; perplexity depends on the tokenizer and is not comparable across different vocabularies.
(5) Likelihood Is Not Helpfulness: maximizing the likelihood of scraped text also rewards reproducing its errors and its style, and teacher forcing never trains the model to recover from its own mistakes, which is the origin of exposure bias.

Bar chart of a next-token distribution after the prefix The capital of France is: Paris receives 0.83, the receives 0.04, a receives 0.021, now 0.015, located 0.012, and the remaining roughly 50000 tokens share 0.082

Figure 1: Illustrative output of one position: a full distribution over the whole vocabulary. Because the observed continuation is “Paris” with probability 0.83, this position contributes $-\ln 0.83 \approx 0.19$ nats to the loss.

Two properties of this objective explain most of how LLMs are trained and evaluated. First, supervision is dense: the causal mask lets one forward pass score every shifted position simultaneously, so a 2048-token document produces 2047 gradient contributions (2048 with a BOS prepended), which is why raw web text is enough to train a model with hundreds of billions of parameters. Second, the loss is a strictly proper scoring rule on the token distribution, so the gradient pushes probability mass toward the observed token and away from every competitor at once; the model is never told which wrong token was “almost right”. That makes the loss an excellent optimization target and a poor proxy for downstream quality, which is why teams track perplexity for training health but use task benchmarks and human preference for release decisions.

Mathematical Formulation:
$p_{\theta}(x_1,\ldots,x_T) = \prod_{t=1}^{T} p_{\theta}(x_t \mid x_{1:t-1})$
$p_{\theta}(\cdot \mid x_{1:t-1}) = \mathrm{softmax}(W h_t + b)$
$\mathcal{L}(\theta) = -\frac{1}{T}\sum_{t=1}^{T} \log p_{\theta}(x_t \mid x_{1:t-1})$
$\mathrm{PPL} = \exp\big(\mathcal{L}(\theta)\big)$

Where:

$x_1,\ldots,x_T$ is the token sequence produced by the tokenizer, and $x_{1:t-1}$ is the left context of position $t$ ; the first line is the chain-rule factorization that defines an autoregressive model.
The $t=1$ term assumes a leading BOS token as context; without it the empty-context term is dropped and the sum runs over $T-1$ positions with the mean taken over $T-1$ terms.
$h_t \in \mathbb{R}^{d}$ is the final-layer hidden state at position $t$ , and $W \in \mathbb{R}^{V \times d}$ with bias $b$ is the unembedding head, frequently weight-tied to the input embedding.
$V$ is the vocabulary size (roughly 128k for Llama 3, about 50k for GPT-2) and $d$ the model width; $\theta$ collects all parameters.
$\mathcal{L}$ is the mean cross-entropy in nats per token, equal to the negative log-likelihood because the target is a one-hot distribution; dividing by $\ln 2$ converts it to bits per token.
$\mathrm{PPL}$ ranges from 1 (a perfect model) up to $V$ for a uniform untrained model, so $\mathcal{L} \approx \ln V$ is the loss value at initialization and a useful sanity check.

Diagram of teacher forcing: a row of input tokens The capital of France is, arrows down to the shifted target row capital of France is Paris, and arrows down to per-position loss boxes 2.10, 0.85, 0.42, 1.31 and 0.19 nats, with the average 0.97 nats and perplexity 2.6 noted below

Figure 2: Teacher forcing on one short sequence: the target row is the input row shifted by one, so a 6-token sequence yields 5 scored positions, every position gets its own cross-entropy term, and the sequence loss is their mean, here 0.97 nats, giving a perplexity of about 2.6.

Property	Causal LM (next token)	Masked LM (BERT style)	Fill-in-the-Middle
What is predicted	Token t+1 from tokens 1 to t	Masked tokens from both sides	A missing span, given prefix and suffix
Supervision per pass	All T-1 shifted positions (T with a BOS)	Only the ~15% masked positions	All positions, on reordered documents
Free-form generation	Native, sample left to right	Not native, needs iterative decoding	Native, plus infilling at a cursor
Best fit	Chat, completion, reasoning, agents	Classification and retrieval embeddings	Code completion inside existing files

August 15, 2026

DL0076 GQA: Grouped-Query Attention

What is grouped-query attention (GQA), and why is it introduced?

Answer

Grouped-query attention splits the $H$ query heads of a layer into $H_{kv}$ groups and gives each group a single shared key head and value head, so the number of distinct K/V projections drops from $H$ to $H_{kv}$ while every query head keeps its own $W_Q$ . It is the interpolation between two known extremes: $H_{kv} = H$ is ordinary multi-head attention (MHA), and $H_{kv} = 1$ is multi-query attention (MQA). It exists because autoregressive decoding is memory-bandwidth bound, not compute bound: generating one token requires re-reading the entire KV cache from HBM, and that cache scales linearly with $H_{kv}$ . MQA already fixed the bandwidth problem but lost quality and was unstable to train, so GQA was proposed to recover almost all of MQA’s decode speed at close to MHA quality by keeping a small number of K/V heads (typically 8). Nearly every major recent open-weight family ships with it, including Llama 3, where all sizes use 8 KV heads; the notable exception is DeepSeek-V2 and V3, which replace head sharing with multi-head latent attention.

(1) The Bottleneck It Targets: the KV cache, not the weights, is what grows with batch size and context, and decode re-reads all of it per token, so shrinking it by $H / H_{kv}$ directly raises arithmetic intensity and tokens per second.
(2) A Dial Between Two Extremes: the lineage MHA → MQA → GQA is a single knob $H_{kv}$ ; 8 KV heads recovers most of the memory win of 1 while avoiding the quality drop that MQA shows on summarization and long-input tasks.
(3) Cheap To Retrofit: an existing MHA checkpoint is converted by mean-pooling the K and V projections within each group and then uptraining on roughly 5% of the original pretraining tokens, so no full retrain is needed.
(4) Set $H_{kv}$ To The Tensor-Parallel Degree: with $H_{kv}$ equal to the number of shards, each GPU owns exactly one KV head and nothing is replicated; a smaller $H_{kv}$ forces the same K/V to be duplicated across shards and gives back part of the saving.
(5) Almost No FLOP Change: the query projections and the $QK^{\top}$ product are untouched, so prefill latency barely moves; GQA is a memory and bandwidth optimization, not a compute one.

Three panels each showing four query head boxes above their key/value heads: MHA has four separate KV heads, GQA with group size two has two KV heads each fed by two query heads, and MQA has one KV head shared by all four query heads

Figure 1: The only structural change is how many distinct K/V heads exist; query heads are never merged, so each still has its own projection and its own attention pattern over the shared keys.

Mathematical Formulation:
$H_{kv} = H / G$
$g(i) = \lfloor i / G \rfloor$
$M_{tok} = 2 L H_{kv} d_h b$
$W_K^{(g)} = \frac{1}{G} \sum_{i \in \mathcal{G}_g} W_K^{(i)}$

Where:

$H$ is the number of query heads, $G$ the group size, and $H_{kv}$ the resulting number of key/value heads; $G = 1$ recovers MHA and $G = H$ recovers MQA.
$i \in \{0,\ldots,H-1\}$ indexes query heads and $g(i)$ is the KV head that head $i$ reads, so heads are assigned to groups by contiguous blocks.
$M_{tok}$ is the KV cache bytes per token, with $L$ layers, head dimension $d_h$ , bytes per element $b$ , and the factor 2 covering K and V; total cache is $M_{tok}$ times batch times sequence length.
$\mathcal{G}_g$ is the set of query-head indices in group $g$ , and the last line is the mean-pooling initialization used when converting an MHA checkpoint; the same averaging is applied to $W_V$ .

The payoff is easiest to see with a concrete configuration. A Llama-3-70B-shaped model has $L = 80$ , $H = 64$ , and $d_h = 128$ ; in FP16 that is 2.5 MiB per token of KV cache under MHA, 320 KiB under GQA with 8 KV heads, and 40 KiB under MQA. At a 32K context and batch size 1 those become 85.9 GB, 10.7 GB, and 1.34 GB. On a single H100 with 3.35 TB/s of HBM bandwidth, reading a 10.7 GB cache once costs roughly 3.2 ms per decoded token against roughly 26 ms for the 85.9 GB version; tensor parallelism divides both numbers, but the ratio is fixed by $H_{kv}$ and does not improve with more hardware. That freed memory is also what lets a server hold many more concurrent sequences, which is usually a larger throughput win than the per-token latency itself.

Log-scale bar chart of FP16 KV cache size at 32K context for a 70B-shaped model: MHA with 64 KV heads needs 85.9 GB, GQA with 8 KV heads needs 10.7 GB, and MQA with 1 KV head needs 1.34 GB, with a dashed line marking 80 GB of H100 HBM

Figure 2: FP16 KV cache for one 32K-token sequence on an 80-layer, 64-head, 128-dim model. Under MHA the cache alone would consume an entire 80 GB H100 before any weights are loaded.

Property	MHA	GQA (8 KV heads)	MQA
KV heads	64 (one per query head)	8 (group size 8)	1 (shared by all)
Cache per token (FP16)	2.5 MiB	320 KiB	40 KiB
Decode speed	Slowest, bandwidth bound on the cache	Close to MQA at long context	Fastest
Quality	Reference	Within noise of MHA on most benchmarks	Measurable drop, notably on long-input summarization
Tensor parallelism	Shards cleanly	Shards cleanly when $H_{kv}$ equals the shard count	KV must be replicated on every shard
Training stability	Stable baseline	Stable; uptrainable from an MHA checkpoint	Reported instability at scale

August 15, 2026