LLM Pretraining

BPE (Byte Pair Encoding) iteratively merges the most frequent adjacent byte or character pair. WordPiece chooses merges that maximize language model likelihood rather than raw frequency. Unigram starts with a large vocabulary and prunes subwords that least reduce likelihood. BPE is used by GPT/Llama; WordPiece by BERT; Unigram by SentencePiece (T5, Gemma). BPE and WordPiece are deterministic; SentencePiece Unigram supports stochastic sampling for data augmentation.

RoPE (Su et al., 2021) encodes position by rotating query and key vectors in the complex plane by an angle proportional to token position. The dot product between Q and K naturally encodes relative position, not absolute — Qm·Kn depends only on (m-n). RoPE does not consume embedding dimensions (unlike learned absolute positions) and generalizes to longer sequences than seen during training via extrapolation or scaling tricks (YaRN, LongRoPE). Llama, Mistral, and Qwen use RoPE.

ALiBi adds a static linear penalty to attention scores based on query-key distance: score(i,j) -= m × |i-j|, where m is a head-specific slope. This biases the model to prefer nearby tokens without any positional embedding added to token representations. ALiBi extrapolates beyond training length better than learned absolute positions. Unlike RoPE, ALiBi has no trainable parameters for positions. RoPE is generally preferred in modern LLMs for its stronger empirical performance; ALiBi is used in BLOOM and MPT.

Kaplan et al. showed that LLM loss follows a power law with respect to three factors: model parameters N, dataset tokens D, and compute budget C. L(N) ∝ N^{-0.076}, L(D) ∝ D^{-0.095}, and optimal scaling allocates compute roughly as N ∝ C^{0.73}. This implies model size scales more steeply than data. Chinchilla (2022) later revised this to N ≈ D (20 tokens/param), finding Kaplan models were undertrained. Scaling laws guide how to allocate training compute across model size and data volume.

During the first few thousand steps, model parameters are far from a good optimum and gradients are large and noisy. Starting with a high learning rate causes instability and loss spikes. Linear warmup gradually increases lr from near-zero to the peak value over warmup_steps (typically 1–4% of total steps). This allows the optimizer's second-moment estimates (in Adam) to stabilize before large updates are applied. After warmup, a cosine decay or inverse-square-root schedule reduces lr to near-zero by training end.

from transformers import get_cosine_schedule_with_warmup
scheduler = get_cosine_schedule_with_warmup(
    optimizer, num_warmup_steps=2000, num_training_steps=total_steps)

CLM (next-token prediction) trains the model to predict each token given all preceding tokens — the loss is the mean cross-entropy over all positions. The attention mask is causal (lower-triangular), preventing tokens from seeing future context. MLM (BERT-style) masks 15% of tokens randomly and trains the model to reconstruct them using bidirectional context. CLM models (GPT) are autoregressive decoders used for generation; MLM models (BERT) are encoders used for discriminative tasks. Modern LLMs use CLM.

Deduplication removes near-duplicate documents from training corpora. Without it, the model memorizes repeated content (verbatim memorization), reducing diversity and wasting compute on redundant data. Studies show even single document duplication within a dataset causes memorization proportional to duplication count. MinHash LSH deduplication (used in RedPajama, Dolma) efficiently finds near-duplicates by comparing shingles. Deduplication improves perplexity and reduces memorization risk — a key safety and copyright concern.

Flash Attention 2 (Dao, 2023) improves over FA1 with: (1) better work partitioning — avoids split-K reduction, assigns Q blocks to warps (not K/V blocks), halving warp communication; (2) fewer non-matmul FLOPs via smarter online softmax rescaling; (3) support for multi-query and grouped-query attention natively; (4) head_dim up to 256. FA2 achieves ~2x speedup over FA1 and approaches 70% of A100 theoretical MFU. It is the standard attention kernel in vLLM, Megatron, and HuggingFace Transformers.

GQA is a hybrid between Multi-Head Attention (MHA) and Multi-Query Attention (MQA). MHA has one K/V head per Q head (H heads total); MQA has a single K/V head shared by all Q heads (smallest KV cache). GQA uses G groups where each group shares K/V heads: G heads of K and V, H heads of Q. Llama 3 uses 8 KV heads for 32 Q heads (GQA-8). GQA reduces KV cache size by H/G while maintaining quality close to MHA, significantly improving inference throughput.

SwiGLU (Shazeer, 2020) is a gated linear unit using Swish: SwiGLU(x, W, V) = Swish(xW) ⊙ (xV). It introduces a learned gating mechanism that controls information flow multiplicatively, providing richer gradient signals. PaLM, Llama, and most modern LLMs use SwiGLU in their FFN layers, typically with an intermediate dimension of (8/3)d instead of 4d to keep parameter count constant. Empirically, SwiGLU/GELU consistently outperforms ReLU on language modeling perplexity.

def swiglu(x, W, V, W2):
    return (F.silu(x @ W) * (x @ V)) @ W2

RMSNorm (Root Mean Square Norm) normalizes by the RMS of the input instead of computing mean and variance: RMSNorm(x) = x / RMS(x) × γ. It drops the mean-centering step of LayerNorm, saving ~10% of LayerNorm compute. Despite the simplification, it matches or exceeds LayerNorm quality in LLM pretraining. Llama, Mistral, and most modern decoder LLMs use RMSNorm over LayerNorm. Pre-norm placement (before the sublayer) is used in Llama; post-norm was used in original Transformer.

class RMSNorm(nn.Module):
    def forward(self, x):
        return x / x.pow(2).mean(-1, keepdim=True).add(1e-8).sqrt() * self.weight

Post-norm (original Transformer): LayerNorm is applied after the residual addition: output = LN(x + sublayer(x)). Pre-norm (GPT, Llama): LN is applied before the sublayer: output = x + sublayer(LN(x)). Pre-norm is more stable for deep networks and large learning rates because the residual stream is unaffected by normalization, enabling gradient flow through the identity path. Post-norm sometimes achieves slightly better final performance but is harder to train. Most modern LLMs use pre-norm.

As training progresses, data quality and distribution matter more than quantity. Data mixing assigns weights to different data sources (web, books, code, papers) to control their token contribution rate. Dolma, RedPajama-v2, and FineWeb use quality filtering + domain weighting. Too much web text makes models weak on reasoning; too much code shifts language modeling. The mixing recipe is empirically tuned. Domain-weighted sampling (sampling_weights per source per step) allows changing the mix dynamically during training.

Sequence packing concatenates multiple short documents into a single context window (up to max_seq_len tokens), separated by a delimiter token. Without packing, short documents are padded to max_seq_len, wasting compute on padding tokens. With packing, GPU utilization for real tokens approaches 100%. A cross-document attention mask (block-diagonal) prevents attention from crossing document boundaries. Packing can improve effective throughput by 20–40% for datasets with many short documents.

MoE replaces the dense FFN in each transformer layer with N expert FFNs and a router. The router takes the token representation as input and outputs a softmax over experts; the top-K experts (typically K=2) are selected per token. Only K/N fraction of parameters are activated per forward pass, making MoE models computationally cheaper than dense models of the same size. Load balancing loss encourages uniform expert utilization. Mixtral 8x7B, GPT-4, and DeepSeek use MoE.

The cooldown or annealing phase drastically reduces the learning rate toward the end of training (e.g., from 3e-4 to 3e-5 over the final 5–10% of steps). This allows the optimizer to settle into a sharp minimum and reduces loss artifacts from noisy late-stage gradient steps. Some recipes also change the data mixture during cooldown — increasing high-quality data (e.g., STEM papers, code) proportions. Dolma/OLMo showed cooldown with quality-focused data significantly improves final benchmarks.

Sequence length (max_seq_len or context window) is the maximum number of tokens the model processes in one forward pass — determined by the positional encoding and attention mask. Context length refers to how many tokens of history the model can attend to at inference. During pretraining, training on longer sequences increases compute quadratically (O(N^2) for attention), so many LLMs pretrain at 4096 tokens and extend context (via RoPE scaling, long-context fine-tuning) post-training.

A loss spike is a sudden large increase in training loss during pretraining, often caused by bad micro-batches (outlier data), numerical instability from large gradient norms, or suboptimal data ordering. Mitigations include: gradient clipping (clip_grad_norm ≤ 1.0), skipping batches with extremely high loss (loss thresholding), using bfloat16 instead of float16 (wider dynamic range), monitoring gradient norms and loss curves, and restarting from a checkpoint before the spike. Llama 2 and OLMo encountered and documented loss spikes.

Weight tying shares the embedding matrix (input token embeddings, shape vocab_size × d_model) with the output projection (also vocab_size × d_model). This reduces model parameters by ~10–15% for large vocabularies while often improving performance because the embedding and unembedding spaces are forced to align. GPT-2 and many BERT variants use weight tying. For very large vocabulary models (100k+), untied weights are sometimes preferred to allow independent optimization.

model.lm_head.weight = model.embed_tokens.weight  # tie weights

Chinchilla trained 400 models of varying sizes and token counts and found optimal compute allocation requires equal scaling of model parameters and tokens: optimal_tokens ≈ 20 × parameters. GPT-3 (175B params, 300B tokens) was undertrained by this metric. Chinchilla 70B trained on 1.4T tokens outperformed Gopher 280B on most benchmarks. This shifted the field from maximizing model size to maximizing data quality and training duration. Llama 3.1 8B trained on 15T tokens follows this principle.

Z-loss is an auxiliary regularization term added to the pretraining loss to prevent logit explosion: z_loss = z × log(Z)^2 where Z is the partition function (logsumexp of logits). Without regularization, softmax output distributions can become extremely peaked, making the model numerically unstable (logits approaching ±inf) and sensitive to small changes. Z-loss was introduced in PaLM and ST-MoE to stabilize training without requiring aggressive gradient clipping that might harm learning.

z_loss = router_z_loss_coef * logits.logsumexp(-1).pow(2).mean()

Continual pretraining initializes from an existing pretrained checkpoint and continues training on new data (e.g., domain-specific text, updated web data, new programming languages). It is far more compute-efficient than pretraining from scratch when adapting to a new domain. Key challenges: catastrophic forgetting (model loses general capabilities) and learning rate selection (often restart from a small LR). Llama 3.1 was continually pretrained on code and multilingual data beyond its base training.

FineWeb (HuggingFace, 2024) applies: (1) language ID filtering (English only); (2) URL blocklist (adult content, spam); (3) MinHash deduplication at document and paragraph level; (4) heuristic quality filters: short line ratio, punctuation ratio, repeated n-gram detection, stop-word density, gibberish detection; (5) perplexity filtering using a small reference LM to remove low-quality text. These filters reduced 90PB of raw Common Crawl to ~15T high-quality tokens.

Sink tokens accumulate disproportionately large attention weights — particularly the first token (often a BOS token or punctuation). This occurs because softmax attention must sum to 1 even for irrelevant tokens; sinking excess attention mass on early tokens prevents attention over-dispersion. LM-Infinite and StreamingLLM exploit this by always retaining sink tokens in the KV cache, enabling efficient long-context inference without retraining. Removing sink tokens causes perplexity degradation on tasks requiring long-range attention.

The KV cache stores past keys and values so autoregressive inference avoids recomputing them for previously generated tokens. Memory: KV_cache_bytes = 2 × num_layers × num_kv_heads × head_dim × seq_len × bytes_per_element. For Llama 3.1 70B (80 layers, 8 KV heads, head_dim=128, float16): 2×80×8×128×4096×2 ≈ 10.7 GB for a single 4096-token sequence. GQA reduces this by H/G. KV cache is the main memory constraint at inference time for long sequences.

# KV cache size formula:
# 2 * L * H_kv * d_head * seq_len * dtype_bytes

Decoder-only models (GPT, Llama, Mistral) use causal attention and are trained purely on next-token prediction. They excel at generation and in-context learning. Encoder-decoder models (T5, BART, mT5) have a bidirectional encoder processing the input and a causal decoder generating the output, trained with teacher forcing. Encoder-decoders are stronger on structured output tasks (translation, summarization) but decoder-only models have won out for general-purpose LLMs due to simpler pretraining and better scaling behavior.

In Megatron attention TP, the Q, K, V projection weight matrices are column-split across TP ranks: each rank computes Q/K/V for its own heads independently (no communication needed for QK^T per head). The output projection is row-split, requiring one all-reduce per transformer block. For MHA with H heads and TP degree t, each rank computes H/t heads. GQA complicates this when KV heads < TP degree — requiring replication of KV heads across ranks or custom sharding strategies.

FIM transforms next-token prediction into an infilling task by rearranging document segments: the prefix and suffix surround a masked span, and the model predicts the middle. The standard format is PREFIX→SUFFIX→MIDDLE, inserted at the beginning of context. This enables code completion models (Codex, StarCoder, DeepSeek-Coder) to fill in code given both preceding and following context. FIM training uses ~50% standard CLM and 50% FIM-rearranged sequences, maintaining autoregressive generation capability.

Tokenizer fertility is the average number of tokens per word in a given language. Languages with low vocabulary coverage (e.g., Arabic, Hindi, code-switched text) fragment into more tokens, increasing sequence length, memory usage, and inference cost. Models trained primarily on English data have high fertility for non-English languages, degrading multilingual performance. Multilingual tokenizers (e.g., Llama 3's 128k vocabulary vs Llama 2's 32k) reduce fertility for underrepresented languages by expanding vocabulary coverage.

Pretraining compute C ≈ 6ND (6 × parameters × tokens) in FLOPs, where the factor 6 accounts for 2 FLOPs/multiply-add in forward (2N) and 4 in backward (4N, including gradient w.r.t. input and weights). Additional costs include embedding lookup, LayerNorm, and softmax. For Llama 3 8B trained on 15T tokens: 6 × 8×10^9 × 15×10^12 ≈ 7.2×10^23 FLOPs. H100 peak BF16 throughput is ~2000 TFLOPS; at 40% MFU: 7.2×10^23 / (0.4 × 2×10^15 FLOPs/s) ≈ 900,000 GPU-hours.

BOS (beginning of sequence) is prepended to every document to signal to the model the start of a new context. EOS (end of sequence) is appended at the end of each document so the model learns to generate a termination signal. In packed sequences (multiple documents per context), EOS/BOS pairs serve as document delimiters. Without EOS training, models trained on packed data learn to generate across document boundaries, causing hallucinations. The attention mask in packing should be block-diagonal to prevent cross-document attention.

Perplexity = exp(cross_entropy_loss) measures how well the model predicts a held-out token sequence — lower is better. It is the primary training metric for CLM. Limitations: (1) it depends on tokenization; models with different tokenizers cannot be compared directly; (2) it measures average token prediction, not task performance; (3) it does not capture reasoning, factuality, or instruction following; (4) models optimized for perplexity may fail benchmark tasks. Modern evaluation uses downstream tasks (MMLU, HumanEval) alongside perplexity.

perplexity = torch.exp(cross_entropy_loss)

F.cross_entropy internally computes log-softmax + NLL in a numerically stable way using the logsumexp identity, avoiding explicit softmax exponentiation which can overflow for large logits. Using softmax() followed by log() and NLLLoss separately introduces numerical error and potential -inf from log(0). PyTorch's F.cross_entropy is always preferred over manual softmax → log → NLL decomposition. Temperature scaling and label smoothing are built into F.cross_entropy parameters.

loss = F.cross_entropy(logits, targets,
                       label_smoothing=0.1)

Curriculum learning trains on easier or shorter examples first, gradually introducing harder or longer examples. For LLMs this can mean starting with shorter sequences (cheaper, simpler context) then increasing max_seq_len during training (length curriculum). Data curriculum orders documents by quality score. The hypothesis is that easier examples first establish a useful prior; harder examples then refine it. Results are mixed: some runs show improved early-stage convergence; others see no significant final-loss difference vs random ordering.

bfloat16 has the same dynamic range as float32 (8 exponent bits), preventing overflow and underflow during gradient all-reduces and loss computation without requiring GradScaler. float16's 5 exponent bits cause overflow for large gradient norms common in LLM training. H100 BF16 Tensor Cores deliver the same 2000 TFLOPS peak throughput as FP16 but without numerical instability. Nearly all modern LLM pretraining (Llama 3, Mistral, Falcon) uses bfloat16 as the standard precision.

model = model.to(torch.bfloat16)  # no GradScaler needed

Embedding dimension d_model (hidden size) determines the vector space in which token representations live. In LLM scaling, both depth (number of layers L) and width (d_model) scale together, but the ratio L/d_model affects the model's aspect ratio. Kaplan et al. found relatively more layers than width is better. GPT-3 uses d_model=12288, L=96. However, very wide models (large d_model) are better at in-context learning; very deep models better at multi-step reasoning. Both dimensions are tuned together empirically.

Synthetic data generated by stronger models is increasingly used to supplement natural pretraining corpora, especially for reasoning, math, and code. Phi-1 (Microsoft) showed that training on synthetic 'textbook-quality' code generated by GPT-4 produced a 1.3B model matching much larger code models. Llama 3.1 used synthetically generated reasoning chains during pretraining. Key risks: model collapse (training on synthetic data from the same model family degrades diversity) and distributional shift.

MTP (Gloeckle et al., Meta 2024) trains the model to predict the next N tokens simultaneously using N parallel output heads sharing the main model backbone. Compared to next-single-token prediction, MTP forces the model to plan ahead and represent richer long-range structure in the hidden state. At inference, only the first head is used for standard generation, but all N heads can be used for speculative decoding (draft N tokens in one forward pass). MTP improved code and reasoning benchmarks in Meta's experiments.

Poor initialization causes vanishing/exploding activations and gradients in deep transformers. Common strategies: weights drawn from Normal(0, 0.02/sqrt(2L)) for residual layers (GPT-2 style) where L is depth — this scales down residual projections to prevent variance accumulation. μP (Maximal Update Parametrization) uses principled width-dependent initialization so hyperparameters (LR, std) transfer from small proxy models to large models without re-tuning. Mistral and DeepSeek use μP-inspired initialization.

# GPT-2 style residual projection init:
std = 0.02 / (2 * num_layers) ** 0.5
nn.init.normal_(proj.weight, std=std)

Common Crawl (CC) is a petabyte-scale public dataset of web pages crawled monthly since 2008, available as WARC (raw HTTP) and WET (text-extracted) files. For LLM pretraining, processing includes: (1) text extraction from HTML; (2) language identification; (3) URL and content quality filtering; (4) MinHash deduplication; (5) safety filtering (adult content, PII removal). FineWeb, Dolma, RedPajama, C4, and The Pile all use processed subsets of CC. CC provides ~70–80% of tokens in most open pretraining datasets.

Positional encodings (especially RoPE) are trained on a maximum sequence length during pretraining. Directly extending context beyond this causes poor performance because RoPE angles for unseen positions are out-of-distribution. Extension techniques: RoPE scaling (Position Interpolation, YaRN, LongRoPE) rescale or adjust RoPE frequencies to cover longer positions. Continued pretraining on long-context data is typically needed after RoPE scaling to adapt attention patterns. Llama 3.1 extended to 128K context via long-context fine-tuning.

μP (Yang et al., 2022) is a theoretical framework for initializing and scaling hyperparameters (learning rate, weight init std) with model width so that the optimal hyperparameters found at small scale transfer to large models without re-tuning. Standard parametrization (SP) causes activations to grow or shrink with width, making hyperparameter transfer unreliable. Under μP, optimal LR scales as 1/d_model. This enables transfer of the learning rate found on a cheap proxy model to the full-scale run, saving significant compute.

RLVR trains LLMs using reinforcement learning where rewards come from verifiable correctness signals (e.g., math solution checking, code execution results) rather than human preference labels. Unlike RLHF which requires human annotators for reward modeling, RLVR can scale cheaply by running unit tests or symbolic solvers. DeepSeek-R1 and Qwen-Math use RLVR for math reasoning improvement. It is applied post-pretraining (like SFT and RLHF) but is increasingly integrated into late-stage pretraining.

Temperature T scales logits before softmax: scaled_logits = logits / T. T=1.0 is standard sampling; T<1.0 sharpens the distribution (more greedy, less diverse); T>1.0 flattens it (more random). For perplexity evaluation, greedy decoding (T→0) or exact likelihood is used. For creative generation, T=0.7–1.0. For multi-step reasoning (chain-of-thought), low T (0.1–0.3) reduces variance. Top-p (nucleus) and top-k sampling further filter the distribution after temperature scaling.

probs = F.softmax(logits / temperature, dim=-1)

Chinchilla-optimal allocates compute evenly between model size N and training tokens D: N ≈ D/20. But inference cost depends only on N (parameters), not D. If a model will run inference at scale (millions of users), it is often better to overtrain a smaller model (more tokens per parameter) — the model is cheaper to serve even if compute during training was suboptimal. Llama 3 8B (15T tokens = 1875 tokens/param) is intentionally overtrained vs Chinchilla-optimal, producing a small but high-quality model.

Instruction following is the ability to respond appropriately to natural language commands. Base LLMs from pretraining are completion models — they continue text, not follow instructions. Instruction following is primarily introduced via supervised fine-tuning (SFT) on instruction-response pairs post-pretraining. Some labs incorporate instruction data during late pretraining phases (cooldown with instruction data) or use synthetic instruction data to seed the capability. RLHF/DPO then refines quality and safety.

Model collapse occurs when models are trained iteratively on synthetic data generated by previous model versions. Each generation of synthetic data loses diversity and introduces artifacts, causing the model to forget the tails of the real data distribution. Over many rounds, the model converges to a narrow distribution. This is a concern for web-scale training where AI-generated content increasingly populates the internet. Mitigations include: retaining access to real data, controlling synthetic data fraction, and monitoring distribution diversity.

Dropout regularizes training by randomly zeroing activations, reducing overfitting. In early NLP transformers (BERT, GPT-2), dropout was applied at rate 0.1. Modern LLMs trained on massive diverse datasets (100B+ tokens) do not overfit — the bottleneck is underfitting, not overfitting. Dropout increases gradient noise, reduces MFU (more computation for dropped activations), and slows convergence. Llama, Mistral, and most 7B+ modern LLMs train without dropout, relying on data volume for regularization.

Teacher forcing feeds the ground-truth previous token as input at each training step, regardless of the model's own prediction. This accelerates and stabilizes training because the model always sees correct context. Free-running (auto-regressive) generation uses the model's own predicted tokens as future inputs — this is how inference works. The discrepancy between training (teacher forcing) and inference (auto-regressive) is called exposure bias. Scheduled sampling gradually mixes teacher-forced and model-generated tokens during training to reduce this gap.

In sparse MoE, a router assigns each token to the top-K experts. Without constraint, the router collapses: it learns to route all tokens to a few experts (winner-take-all), leaving most experts unused (dead experts). The auxiliary load balancing loss penalizes unequal expert utilization by encouraging the fraction of tokens each expert receives to be uniform. A common form is: loss = alpha × sum_i(f_i × p_i) where f_i is the fraction of tokens routed to expert i and p_i is the mean router probability for expert i. This is added to the main CLM loss.

load_balance_loss = alpha * sum(fraction_i * router_prob_i for i in range(n_experts))